Only disallow DeepSpeed Zero-3 for auto bs finder

[muellerzr]

What does this PR do?

Looking into re-enabling the auto bs finder for FSDP and DeepSpeed. Did a small scale test with FSDP and found that training works fine. Will move onto DeepSpeed soon but here is an interim.

At the worst users can raise an issue here (a new one) with what specifically breaks with FSDP. But my attempts to break it based on FSDP have yet to fail (as FSDP doesn't care about batch sizes actually!)

I don't enjoy not having a test here, in my tests I was using tinyllama on 48GB (2x4090). What's the general guidelines for using non-custom models in this case? Do we make test models of size N? tinyllama is a 1.1B Llama model, so I can probably write a chunk test with a model that size if we need to/should. Just curious if we have any methods of dynamically creating models of n size for testing purposes (I know we do so in accelerate)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• 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.

@amyeroberts @SunMarc

[muellerzr]

I don't think the failing tests are on my end 😅

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[muellerzr]

Update: Zero-2 works, so it's Zero-3 that we have issues with. I've pinged the DeepSpeed team for help, but this gets us most of the way there :)

[muellerzr]

Essentially zero-3 does some pretty intense monkey patching into their Init class that persists despite the model references being deleted, because they literally add attributes to the original model itself. Once we can figure out how to wipe everything + the state of Init, we can fully enable everything again :)

[amyeroberts]

LGTM - thanks!

Re small models @ydshieh is the best person to ask :)

[ydshieh]

Hi @muellerzr !

Just curious if we have any methods of dynamically creating models of n size for testing purposes

We have a script to create tiny models for testing purpose, however:

• it is not designed to be able to create a specific wanted size
• it is not guaranteed to always work (for some model classes, it may fail)

[muellerzr]

Thanks @ydshieh, I'll think today on how I can maybe write a test for this then on a small enough model/large enough BS that this can work no matter the GPU and still is efficient and add a test in. Worst case we may not be able to