DeepSpeed

01af3e1d - Enhance the robustness of `module_state_dict` (#4587)

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