Skip to content

Commit

Permalink
Build PT aarch64 on arm runner (pytorch#131964)
Browse files Browse the repository at this point in the history
Another fix is needed to address https://github.com/pytorch/pytorch/actions/runs/10118374576/job/27985575620.  The build needs to be done on arm runner to stay compatible with the Docker image.

### Testing

https://github.com/pytorch/pytorch/actions/runs/10118589329/job/27985670691

Pull Request resolved: pytorch#131964
Approved by: https://github.com/malfet
  • Loading branch information
huydhn authored and pull[bot] committed Aug 2, 2024
1 parent 6c47ceb commit 760c928
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 2 deletions.
1 change: 1 addition & 0 deletions .github/workflows/inductor-perf-test-nightly-aarch64.yml
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ jobs:
name: linux-jammy-aarch64-py3.10-inductor
uses: ./.github/workflows/_linux-build.yml
with:
runner: linux.arm64.m7g.4xlarge
build-environment: linux-jammy-aarch64-py3.10
docker-image-name: pytorch-linux-jammy-aarch64-py3.10-gcc11-inductor-benchmarks
test-matrix: |
Expand Down
6 changes: 4 additions & 2 deletions test/dynamo/test_model_output.py
Original file line number Diff line number Diff line change
Expand Up @@ -239,10 +239,12 @@ def fn(inp):

@maybe_skip
def test_HF_bert_model_output(self):
device = "cuda" if torch.cuda.is_available() else "cpu"

class BertPooler(torch.nn.Module):
def __init__(self):
super().__init__()
self.dense = torch.nn.Linear(768, 768).to("cuda")
self.dense = torch.nn.Linear(768, 768).to(device)
self.activation = torch.nn.Tanh()

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
Expand Down Expand Up @@ -300,7 +302,7 @@ def forward(
result["pooler_output"] = pooled_output
return result

sequence_output = torch.rand(1, 12, 768).to("cuda")
sequence_output = torch.rand(1, 12, 768).to(device)
model = BertModel()
orig_result = model(sequence_output)
compiled_model = torch.compile(model, backend="eager")
Expand Down

0 comments on commit 760c928

Please sign in to comment.