The module's original class is added after HSDPModule in the runtime
MRO, so static analysis cannot resolve the delegated ``_apply``.
"""
if not hasattr(self, "hsdp_scheduler") or self.hsdp_scheduler is None:
return super()._apply(*args, **kwargs) # pylint: disable=no-member
self.reshard()
result = super()._apply(*args, **kwargs) # pylint: disable=no-member
hsdp_state = self.hsdp_scheduler.hsdp_state
if hsdp_state is None:
return result
with torch.no_grad():
for hsdp_param in hsdp_state.hsdp_params:
hsdp_param.reset_sharded_param()
return result
def hsdp_init(self, module, mesh, reshard_after_forward,
shard_placement_fn, mp_policy, offload_policy, ignored_params, replicate_params, device,
comm_fusion, comm_fusion_zero_copy: Optional[bool] = None,