params partition for skip_init #4722
Conversation
inkcherry
requested review from
jeffra,
tjruwase and
mrwyattii
as code owners
tjruwase
requested review from
tohtana
and removed request for
mrwyattii and
loadams
|
Hi @inkcherry, |
Thanks for the review. @tohtana |
|
Thank you for the clarification, @inkcherry. I am wondering what happens if almost all modules are declared with I understand that it is difficult to set a hook in |
|
@tohtana ， Thanks for your suggestion! currently, I have implemented a hook, without encountering such memory issues[after the child module initialization is completed by I tested full-parameters finetuning with chatglm2 6B zero3, without using skip_init and this patch, and compared it with using the default skip_init of the model with this patch. The loss is exactly the same for the first 50 steps. |
|
@tohtana I just make some changes ,could you please take a look, thanks! |
|
Thank you @inkcherry, the hook you implemented should work. This is an intricate and refined work! |
|
@inkcherry This PR was merged. Thank you for your great contribution! |
Some models use ```skip_init``` to initialize weights. ```skip_init``` first initializes on a meta device in ```__init__``` of a module and then uses ```to_empty()```. This conflicts with the deepspeed hook ```module.__init__``` mechanism. it's necessary to wait for ```skip_init``` to finish before executing ```_post_init_method```. However, the ```from ... import skip_init``` behavior typically occurs outside the context, there seems to be no good way to directly hook into ```skip_init```. Hence, the approach here is to delay the execution of ```_post_init_method``` to resolve this issue. Known affected models include HuggingFace models like chatglm2 and chatglm3." --------- Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
Some models use
skip_initto initialize weights.skip_initfirst initializes on a meta device in__init__of a module and then usesto_empty(). This conflicts with the deepspeed hookmodule.__init__mechanism. it's necessary to wait forskip_initto finish before executing_post_init_method. However, thefrom ... import skip_initbehavior typically occurs outside the context, there seems to be no good way to directly hook intoskip_init. Hence, the approach here is to delay the execution of_post_init_methodto resolve this issue.Known affected models include HuggingFace models like chatglm2 and chatglm3."