Method Signature Updated: The sharding_scheme parameter in allgather_weights_and_update_full_parameter is made optional by setting a default value of None.
# Original
def allgather_weights_and_update_full_parameter(self, sharding_scheme):
# method body
# Change
def allgather_weights_and_update_full_parameter(self, sharding_scheme=None):
if sharding_scheme is None:
sharding_scheme = self._get_sharding_scheme({})
# method body
Issue Addressed: This change fixes a bug where the load_state_dict method calls allgather_weights_and_update_full_parameter without providing a sharding_scheme, which would cause a TypeError.
# In both versions within load_state_dict
self.allgather_weights_and_update_full_parameter()
Impact:
Bug Fix: By making sharding_scheme optional and handling the None case inside the method, the second version prevents the TypeError and ensures that load_state_dict functions correctly without requiring additional arguments.
Change Summary:
Method Signature Updated: The
sharding_schemeparameter inallgather_weights_and_update_full_parameteris made optional by setting a default value ofNone.Issue Addressed: This change fixes a bug where the
load_state_dictmethod callsallgather_weights_and_update_full_parameterwithout providing asharding_scheme, which would cause aTypeError.Impact:
sharding_schemeoptional and handling theNonecase inside the method, the second version prevents theTypeErrorand ensures thatload_state_dictfunctions correctly without requiring additional arguments.