Skip to content

Fix flaky SwitchTransformersModelTest::test_training_gradient #35587

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ydshieh merged 2 commits into from
Jan 9, 2025
Merged

Fix flaky SwitchTransformersModelTest::test_training_gradient #35587

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2bb9d89
fix
ydshieh Jan 9, 2025
92869f6
Update tests/models/switch_transformers/test_modeling_switch_transfor…
ydshieh Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions tests/models/switch_transformers/test_modeling_switch_transformers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -576,6 +576,8 @@ class SwitchTransformersModelTest(ModelTesterMixin, GenerationTesterMixin, Pipel
test_torchscript = False
# The small SWITCH_TRANSFORMERS model needs higher percentages for CPU/MP tests
model_split_percents = [0.5, 0.8, 0.9]
# `SwitchTransformers` is a MOE in which not all experts will get gradients
Comment thread
ydshieh marked this conversation as resolved.
Outdated
test_all_params_have_gradient = False

def setUp(self):
self.model_tester = SwitchTransformersModelTester(self)
Expand Down
9 changes: 6 additions & 3 deletions tests/test_modeling_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,8 @@ class ModelTesterMixin:
test_mismatched_shapes = True
test_missing_keys = True
test_model_parallel = False
# Used in `check_training_gradient_checkpointing` to NOT check all params having gradient (e.g. for some MOE models)
test_all_params_have_gradient = True
is_encoder_decoder = False
has_attentions = True
_is_composite = False
Expand Down Expand Up @@ -895,9 +897,10 @@ def check_training_gradient_checkpointing(self, gradient_checkpointing_kwargs=No
loss.backward()
optimizer.step()

for k, v in model.named_parameters():
if v.requires_grad:
self.assertTrue(v.grad is not None, f"{k} in {model_class.__name__} has no gradient!")
if self.test_all_params_have_gradient:
for k, v in model.named_parameters():
if v.requires_grad:
self.assertTrue(v.grad is not None, f"{k} in {model_class.__name__} has no gradient!")

def test_training(self):
if not self.model_tester.is_training:
Expand Down