set supports_gradient_checkpointing = False in SwitchTransformer
Author: nhamanasu
Pull request: yes

# What does this PR do?

Recently, CircleCI's `tests_torch` sometimes fails (Reference: [tests_torch](https://app.circleci.com/pipelines/github/huggingface/transformers/113436/workflows/802101b9-ed24-49ac-bc1d-7f2ae9dce240/jobs/1515817?utm_campaign=vcs-integration-link&utm_medium=referral&utm_source=github-checks-link&utm_content=summary) - Failed), because some models which don't support gradient_checkpointing or have MoE module like SwitchTransformer will occasionally have weights without gradients.

I found maybe-related PR: https://github.com/huggingface/transformers/pull/34806, and the situation described in the PR was already fixed (https://github.com/pytest-dev/pytest-subtests/pull/169 was already merged).

In https://github.com/huggingface/transformers/pull/34806, @ydshieh said: `TODO (ydshieh): use skipTest once pytest-dev/pytest-subtests/pull/169 is merged`, so I reverted the stopgap fix.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

cc @ydshieh @muellerzr @SunMarc

Comments:
Hi @nhamanasu 

Thank you for the PR. Although https://github.com/pytest-dev/pytest-subtests/pull/169 is merged, we will still need to review the impact of installing and using `pytest-subtests` as it will change the report format.

In the meantime, #34806 should not be the cause of failing `check_training_gradient_checkpointing` as it skips some cases (using `continue`). It means there are some `model_class` that are still failing this test.

Change it to `skipTest` along with https://github.com/pytest-dev/pytest-subtests/pull/169 won't help here, as that is somehow equivalent `continue`.

Could you check if `SwitchTransformers` is MOE type models? Maybe this model just doesn't support gradient checkpoint and we need to set `supports_gradient_checkpointing = False` for this model


Hi @ydshieh 

Thank you for your response! I got the situation (sorry for my misunderstanding).

Related to SwitchTransformer, it currently has [`supports_gradient_checkpointing = True`](https://github.com/huggingface/transformers/blob/31f9a289a6207be6cae746e009d8e0db523be203/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L764C1-L775C52).

And also, SwitchTransformer has experts in MLP modules:
- SwitchTransformersSparseMLP has [router](https://github.com/huggingface/transformers/blob/31f9a289a6207be6cae746e009d8e0db523be203/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L273C9-L274C59) and [experts](https://github.com/huggingface/transformers/blob/31f9a289a6207be6cae746e009d8e0db523be203/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L276C9-L279C65)
- [The router only returns the most plausible expert indices mask](https://github.com/huggingface/transformers/blob/31f9a289a6207be6cae746e009d8e0db523be203/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L182C5-L211C57), and [SwitchTransformersSparseMLP seems to use it to select experts in a discrete way](https://github.com/huggingface/transformers/blob/31f9a289a6207be6cae746e009d8e0db523be203/src/transformers/models/switch_transformers/modeling_switch_transformers.py#L305C9-L311C14).

I think the last part should be the cause of non-gradient weights. So as you suggested, I think we need to set supports_gradient_checkpointing = False for this model.

- CircleCI's `tests_generate` has failed, but it'll be resolved in this PR: https://github.com/huggingface/transformers/pull/35246
- `tests_torch` (The target of this PR) has passed after setting `supports_gradient_checkpointing = False` in SwitchTransformersPreTrainedModel.

Thanks. Just 2 nits and we are ready to go

Thank you! I've reflected your comments.

Thanks! I think recently @ydshieh had a new look and this hsould be fixed (it does support GC just that some expert are sometimes not used) 

For the record: #35587

Sorry for not including you as a contributor in that PR (my limited short memory ... 😢 ), and thank you !

I am closing this PR.