Nice!

I guess we don't have much of an option beyond using a config parameter for now.

Given that we run forward pass on a submodule, it would be very nice to have some validation so that there are no unexpected failures etc with different distributed training settings. But in this case, I feel there might be compatibility issues with FSDP given the limitation from the docs: https://pytorch.org/docs/stable/fsdp.html

"FSDP does not support running the forward pass of a submodule that is contained in an FSDP instance. This is because the submodule’s parameters will be sharded, but the submodule itself is not an FSDP instance, so its forward pass will not all-gather the full parameters appropriately."

(might be fixed by just making the base model attribute an FSDP instance as well, coz why not)

Beyond that this looks fine! I have a couple nits (don't matter that much):

• Does model.get_decoder() work all the time btw for AutoModelForCausalLM instances? Was wondering if that's a cleaner solution for getting the base model attribute. But I think some base model classes have some further wrapping over the actual decoder (to format outputs, etc) https://github.com/huggingface/transformers/blob/a7f5479b45a8040392af80bf1107a2bdd796931c/src/transformers/models/opt/modeling_opt.py#L1044
• Maybe the config is base_model_attribute_name since its the attribute name of the base model in the CausalLM object?

thanks @SumanthRH yes you are right get_decoder() will work

Yes next is to verify the distributed training cases