megatron-fsdp 0.2.0.dev123727__tar.gz → 0.2.0.dev124909__tar.gz

{megatron_fsdp-0.2.0.dev123727 → megatron_fsdp-0.2.0.dev124909}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-fsdp
-Version: 0.2.0.dev123727
+Version: 0.2.0.dev124909
 Summary: **Megatron-FSDP** is an NVIDIA-developed PyTorch extension that provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP)
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>

{megatron_fsdp-0.2.0.dev123727 → megatron_fsdp-0.2.0.dev124909}/megatron_fsdp/megatron_fsdp.py

@@ -898,9 +898,10 @@ class MegatronFSDP(torch.nn.Module):
 
         # Register pre state_dict hook to ensure that the module parameters are
         # distributed before saving the state_dict.
-        self._state_dict_pre_hook = self.module.register_state_dict_pre_hook(
-            lambda *args, **kwargs: self._replace_param_with_distributed_if_needed()
-        )
+        for name, module in self.named_modules():
+            module.register_state_dict_pre_hook(
+                lambda *args, **kwargs: self._replace_param_with_distributed_if_needed()
+            )
 
     @contextmanager
     def no_sync(self):

{megatron_fsdp-0.2.0.dev123727 → megatron_fsdp-0.2.0.dev124909}/megatron_fsdp/package_info.py

@@ -4,7 +4,7 @@
 MAJOR = 0
 MINOR = 2
 PATCH = 0
-PRE_RELEASE = '0.dev123727'
+PRE_RELEASE = '0.dev124909'
 
 # Use the following formatting: (major, minor, patch, pre-release)
 VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)

{megatron_fsdp-0.2.0.dev123727 → megatron_fsdp-0.2.0.dev124909}/megatron_fsdp/param_and_grad_buffer.py

@@ -2782,6 +2782,9 @@ class GradReducePipeline:
             outer_fsdp_group_grad_reduce (bool, optional): Whether to reduce gradients
                 across outer-DP groups. Defaults to False.
         """
+        # Sort parameters by their bucket IDs to ensure a deterministic processing order.
+        # Performing reduce-scatter operations out of order can lead to hangs.
+        params = sorted(list(params), key=lambda x: self.buffer.param_to_param_group[x])
         for param in params:
             bucket_id = self.buffer.param_to_param_group[param]
             param_group = self.buffer.parameter_groups[bucket_id]

{megatron_fsdp-0.2.0.dev123727 → megatron_fsdp-0.2.0.dev124909}/megatron_fsdp.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-fsdp
-Version: 0.2.0.dev123727
+Version: 0.2.0.dev124909
 Summary: **Megatron-FSDP** is an NVIDIA-developed PyTorch extension that provides a high-performance implementation of Fully Sharded Data Parallelism (FSDP)
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>