FIX Make special LoRA inits DeepSpeed compatible

[BenjaminBossan]

Resolves huggingface/accelerate#2886

Possibly resolves
#896 (comment)

Some LoRA init methods need to access the base layer weight. Getting this access can fail or stall in distributed settings. For DeepSpeed, the weight is now gathered before trying to access it.

Note: Without DeepSpeed, this is a no-op and should thus not have any disadvantage. We don't have DS in our CI, so this is not tested.

I also made some small changes to OLoRA init to use self.get_base_layer() instead of self.base_layer.

[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.

[BenjaminBossan]

ping @tokenizer-decode as there are some slight changes to OLoRA.

[SunMarc]

LGTM ! Maybe we can leave a comment on why we added that for these inits since we don't do it for other inits

[tokenizer-decode]

Could you explain to me why self.get_base_layer() rather than self.base_layer? No mutation to base weights occured at that part so I'd assume both should return the same? How did this solve the problem? I am surprised.

[BenjaminBossan]

Could you explain to me why self.get_base_layer() rather than self.base_layer? No mutation to base weights occured at that part so I'd assume both should return the same?

The reason why we have get_base_layer() is that we support nesting LoRA adapters and other adapter types. Calling the method ensures we really get back the "original" base layer. When there is no nesting, both approaches are indeed the same.

How did this solve the problem? I am surprised.

This change is not related to the DeepSpeed issue, that's what I wanted to convey when I wrote "I also made some small changes to OLoRA".

[tokenizer-decode]

Oh now I see. Thanks for notifying. Looks good.