Add Apple M1 tests (#796) · addianto/diffusers@cca59ce · GitHub
Skip to content

Commit cca59ce

Browse files
anton-lpcuenca
anton-l
and
pcuenca
authored
Add Apple M1 tests (huggingface#796)
* [CI] Add Apple M1 tests * setup-python * python build * conda install * remove branch * only 3.8 is built for osx-arm * try fetching prebuilt tokenizers * use user cache * update shells * Reports and cleanup * -> MPS * Disable parallel tests * Better naming * investigate worker crash * return xdist * restart * num_workers=2 * still crashing? * faulthandler for segfaults * faulthandler for segfaults * remove restarts, stop on segfault * torch version * change installation order * Use pre-RC version of PyTorch. To be updated when it is released. * Skip crashing test on MPS, add new one that works. * Skip cuda tests in mps device. * Actually use generator in test. I think this was a typo. * make style Co-authored-by: Pedro Cuenca <pedro@huggingface.co>
1 parent 627ad6e commit cca59ce

7 files changed

Lines changed: 235 additions & 12 deletions

File tree

‎.github/actions/setup-miniconda/action.yml‎

Lines changed: 146 additions & 0 deletions

‎.github/workflows/pr_tests.yml‎

Lines changed: 53 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
name: Run non-slow tests
1+
name: Run fast tests
22

33
on:
44
pull_request:
@@ -10,14 +10,14 @@ concurrency:
1010
cancel-in-progress: true
1111

1212
env:
13-
HF_HOME: /mnt/cache
1413
OMP_NUM_THREADS: 8
1514
MKL_NUM_THREADS: 8
1615
PYTEST_TIMEOUT: 60
16+
MPS_TORCH_VERSION: 1.13.0
1717

1818
jobs:
1919
run_tests_cpu:
20-
name: Diffusers tests
20+
name: CPU tests on Ubuntu
2121
runs-on: [ self-hosted, docker-gpu ]
2222
container:
2323
image: python:3.7
@@ -39,7 +39,7 @@ jobs:
3939
run: |
4040
python utils/print_env.py
4141
42-
- name: Run all non-slow selected tests on CPU
42+
- name: Run all fast tests on CPU
4343
run: |
4444
python -m pytest -n 2 --max-worker-restart=0 --dist=loadfile -s -v --make-reports=tests_torch_cpu tests/
4545
@@ -51,5 +51,53 @@ jobs:
5151
if: ${{ always() }}
5252
uses: actions/upload-artifact@v2
5353
with:
54-
name: pr_torch_test_reports
54+
name: pr_torch_cpu_test_reports
55+
path: reports
56+
57+
run_tests_apple_m1:
58+
name: MPS tests on Apple M1
59+
runs-on: [ self-hosted, apple-m1 ]
60+
61+
steps:
62+
- name: Checkout diffusers
63+
uses: actions/checkout@v3
64+
with:
65+
fetch-depth: 2
66+
67+
- name: Clean checkout
68+
shell: arch -arch arm64 bash {0}
69+
run: |
70+
git clean -fxd
71+
72+
- name: Setup miniconda
73+
uses: ./.github/actions/setup-miniconda
74+
with:
75+
python-version: 3.9
76+
77+
- name: Install dependencies
78+
shell: arch -arch arm64 bash {0}
79+
run: |
80+
${CONDA_RUN} python -m pip install --upgrade pip
81+
${CONDA_RUN} python -m pip install -e .[quality,test]
82+
${CONDA_RUN} python -m pip install --pre torch==${MPS_TORCH_VERSION} --extra-index-url https://download.pytorch.org/whl/test/cpu
83+
84+
- name: Environment
85+
shell: arch -arch arm64 bash {0}
86+
run: |
87+
${CONDA_RUN} python utils/print_env.py
88+
89+
- name: Run all fast tests on MPS
90+
shell: arch -arch arm64 bash {0}
91+
run: |
92+
${CONDA_RUN} python -m pytest -n 4 --max-worker-restart=0 --dist=loadfile -s -v --make-reports=tests_torch_mps tests/
93+
94+
- name: Failure short reports
95+
if: ${{ failure() }}
96+
run: cat reports/tests_torch_mps_failures_short.txt
97+
98+
- name: Test suite reports artifacts
99+
if: ${{ always() }}
100+
uses: actions/upload-artifact@v2
101+
with:
102+
name: pr_torch_mps_test_reports
55103
path: reports

‎tests/test_layers_utils.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,9 @@ def test_downsample_with_conv_out_dim(self):
221221

222222

223223
class AttentionBlockTests(unittest.TestCase):
224+
@unittest.skipIf(
225+
torch_device == "mps", "Matmul crashes on MPS, see https://github.com/pytorch/pytorch/issues/84039"
226+
)
224227
def test_attention_block_default(self):
225228
torch.manual_seed(0)
226229
if torch.cuda.is_available():
@@ -245,6 +248,30 @@ def test_attention_block_default(self):
245248
)
246249
assert torch.allclose(output_slice.flatten(), expected_slice, atol=1e-3)
247250

251+
def test_attention_block_sd(self):
252+
# This version uses SD params and is compatible with mps
253+
torch.manual_seed(0)
254+
if torch.cuda.is_available():
255+
torch.cuda.manual_seed_all(0)
256+
257+
sample = torch.randn(1, 512, 64, 64).to(torch_device)
258+
attentionBlock = AttentionBlock(
259+
channels=512,
260+
rescale_output_factor=1.0,
261+
eps=1e-6,
262+
num_groups=32,
263+
).to(torch_device)
264+
with torch.no_grad():
265+
attention_scores = attentionBlock(sample)
266+
267+
assert attention_scores.shape == (1, 512, 64, 64)
268+
output_slice = attention_scores[0, -1, -3:, -3:]
269+
270+
expected_slice = torch.tensor(
271+
[-0.6621, -0.0156, -3.2766, 0.8025, -0.8609, 0.2820, 0.0905, -1.1179, -3.2126], device=torch_device
272+
)
273+
assert torch.allclose(output_slice.flatten(), expected_slice, atol=1e-3)
274+
248275

249276
class SpatialTransformerTests(unittest.TestCase):
250277
def test_spatial_transformer_default(self):

‎tests/test_modeling_common.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,7 @@ def recursive_check(tuple_object, dict_object):
247247

248248
recursive_check(outputs_tuple, outputs_dict)
249249

250+
@unittest.skipIf(torch_device == "mps", "Gradient checkpointing skipped on MPS")
250251
def test_enable_disable_gradient_checkpointing(self):
251252
if not self.model_class._supports_gradient_checkpointing:
252253
return # Skip test if model does not support gradient checkpointing

‎tests/test_models_unet.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ def test_from_pretrained_hub(self):
135135

136136
assert image is not None, "Make sure output is not None"
137137

138-
@unittest.skipIf(torch_device == "cpu", "This test is supposed to run on GPU")
138+
@unittest.skipIf(torch_device != "cuda", "This test is supposed to run on GPU")
139139
def test_from_pretrained_accelerate(self):
140140
model, _ = UNet2DModel.from_pretrained(
141141
"fusing/unet-ldm-dummy-update", output_loading_info=True, device_map="auto"
@@ -145,7 +145,7 @@ def test_from_pretrained_accelerate(self):
145145

146146
assert image is not None, "Make sure output is not None"
147147

148-
@unittest.skipIf(torch_device == "cpu", "This test is supposed to run on GPU")
148+
@unittest.skipIf(torch_device != "cuda", "This test is supposed to run on GPU")
149149
def test_from_pretrained_accelerate_wont_change_results(self):
150150
model_accelerate, _ = UNet2DModel.from_pretrained(
151151
"fusing/unet-ldm-dummy-update", output_loading_info=True, device_map="auto"
@@ -177,7 +177,7 @@ def test_from_pretrained_accelerate_wont_change_results(self):
177177

178178
assert torch.allclose(arr_accelerate, arr_normal_load, rtol=1e-3)
179179

180-
@unittest.skipIf(torch_device == "cpu", "This test is supposed to run on GPU")
180+
@unittest.skipIf(torch_device != "cuda", "This test is supposed to run on GPU")
181181
def test_memory_footprint_gets_reduced(self):
182182
torch.cuda.empty_cache()
183183
gc.collect()
@@ -267,6 +267,7 @@ def prepare_init_args_and_inputs_for_common(self):
267267
inputs_dict = self.dummy_input
268268
return init_dict, inputs_dict
269269

270+
@unittest.skipIf(torch_device == "mps", "Gradient checkpointing skipped on MPS")
270271
def test_gradient_checkpointing(self):
271272
# enable deterministic behavior for gradient checkpointing
272273
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()

‎tests/test_models_vae.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ def test_output_pretrained(self):
9696
model.config.in_channels,
9797
model.config.sample_size,
9898
model.config.sample_size,
99-
generator=torch.manual_seed(0),
99+
generator=generator,
100100
)
101101
image = image.to(torch_device)
102102
with torch.no_grad():

‎tests/test_pipelines.py‎

Lines changed: 3 additions & 3 deletions

0 commit comments

Comments
 (0)