Enhance the robustness of module_state_dict

[LZHgrla]

If user re-implements the state_dict() method of self.module and manually excludes the frozen parameters, this line will lead to KeyError.

[tjruwase]

@LZHgrla, thanks for this PR. Could you please add a unit test like the follow?

DeepSpeed/tests/unit/checkpoint/test_zero_optimizer.py

Line 500 in 4199dc2

def test_save_exclude_frozen_weights(self, tmpdir, zero_stage):

[LZHgrla]

@LZHgrla please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

• (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.

@microsoft-github-policy-service agree

• (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.

@microsoft-github-policy-service agree company="Microsoft"

Contributor License Agreement

@microsoft-github-policy-service agree

[LZHgrla]

@tjruwase Thanks for your reviews! The unit tests have been added~

[LZHgrla]

Hi @tjruwase
It seems that there are some problems with CI.. What can I do for it?

[tjruwase]

Hi @tjruwase It seems that there are some problems with CI.. What can I do for it?

@LZHgrla, apologies for the inconvenience. I don't there is any action needed on your side, right it looks like we need to resolve the CI issues. Please stay tuned.

[LZHgrla]

@tjruwase
Thanks! It seems CI works now!