DeepSpeed Zero3 is Incompatible with Freeze Range Code

[josharian]

Please check that this issue hasn't been reported before.

• I searched previous Bug Reports didn't find any similar reports.

Expected Behavior

Set up a config like:

unfrozen_parameters:
 - ^model.embed_tokens.weight$[128256:] # only train the new tokens

deepspeed: deepspeed_configs/zero3.json

Train.

Expect something like:

Unfrozen model.embed_tokens.weight with ranges [(128256, 130304)]

Got:

Unfrozen model.embed_tokens.weight with ranges [(128256, 0)]

This leads to things...not working as intended.

#1686 will make diagnosis/recognition of this easier. But it doesn't fix the root problem.

AFAICT, the root problem is that deepspeed/zero3.json changes model loading such that the parameters no longer have their original shapes, like this:

>>> print(model.state_dict()["model.embed_tokens.weight"].shape)
torch.Size([0])

As a result, when range end is None, it gets set to 0.

(It also appears that this may mess with model saving as well. My saved models with deepspeed/zero3.json are way too small, possible because they have shape torch.Size([0]) for almost all layers.)

Current behaviour

see above

Steps to reproduce

see above

Config yaml

No response

Possible solution

No response

Which Operating Systems are you using?

• Linux
• macOS
• Windows

Python Version

3.11

axolotl branch-commit

whatever docker image has (how do i get this from the docker image?)

Acknowledgements

• My issue title is concise, descriptive, and in title casing.
• I have searched the existing issues to make sure this bug has not been reported yet.
• I am using the latest version of axolotl.
• I have provided enough information for the maintainers to reproduce and diagnose the issue.