HF FSDP wrap BLOOM and OPT as well

[samhavens]

This makes the FSDP wrap function a bit more generic so it will work on OPT and BLOOM models.

However, it removes an assert, so the user could now call it on any model — we may want some sanity check that they aren't calling it on e.g. T5, which this doesn't support yet (would need to look at model.encoder for some properties, I think). For now, I raise a ValueError if any of the expected parts of the model aren't found.

Also, I wasn't sure on how to correctly attribute the helper functions, see the file header.

[abhi-mosaic]

Looking good so far! I think we should also consider replacing ComposerHFCausalLM(ComposerModel): with Composer's HuggingFaceModel class and see if it works/breaks. It gives us a lot of good functionality for free.

But also totally understand if it makes to separate into a followup PR.

[samhavens]

Following up: this code works on my LegalGPT fork to FSDP-wrap OPT and BLOOM. I tested on 4x A100 40G, and OPT 6.7B and BLOOM 3B works. However, I got errors when using using it in this repo which I need to investigate.

I will update it to use HuggingFaceModel

[alextrott16]

Love this. Simple and easy to follow/extend.

It would be good to confirm that it works on T5 (an Encoder-Decoder) and BERT (an Encoder-only). Along those lines, I might recommend explicitly defining the model types that this supports (with unsupported models triggering a warning or error) since there are a lot of HF model types out there and the 10 most popular ones will probably 99% of use cases.

[samhavens]

@alextrott16 I agree that we should give information about confidence. There are 2 levels that I see: "we haven't checked this model yet, but it has all the properties that we expect in the helper functions" (likely will work?) vs "we tried to wrap this model, and were unable to find needed components" (almost certainly won't work)

[samhavens]

@alextrott16 I agree that we should give information about confidence. There are 2 levels that I see: "we haven't checked this model yet, but it has all the properties that we expect in the helper functions" (likely will work?) vs "we tried to wrap this model, and were unable to find needed components" (almost certainly won't work)

I've implemented the later as a ValueError. Now need to test BERT-style & T5-style

[samhavens]

@alextrott16 RE: model incompatibility concerns:

• this seems to wrap Codegen (a decoder-only model I hadn't looked at).
• the codebase errors if you try to pass T5, so didn't check that one (we call AutoModelForCausalLM, which substantially narrows down the potential models)
• it raises a ValueError if you try to use Bert

So I think we are good on that front!

[abhi-mosaic]

Thanks for untangling FSDP and expanding our supported models!