[sync] Fix bug in param_norm computation where some ranks might call collective and some might not (#852)

* sync with mr 3918

Signed-off-by: Ananth Subramaniam <[email protected]>

* add unit tests

Signed-off-by: Ananth Subramaniam <[email protected]>

* rebase

Signed-off-by: Ananth Subramaniam <[email protected]>

---------

Signed-off-by: Ananth Subramaniam <[email protected]>
Signed-off-by: mollys <[email protected]>

Author: ananthsub

src/megatron/bridge/training/utils/train_utils.py

@@ -186,14 +186,16 @@ def calc_params_l2_norm(
             False,  # no per-parameter norm.
         )
         sharded_norm_2 = sharded_norm * sharded_norm
-        # Sum over all DP groups, including CP since distributed optimizer state is
-        # sharded jointly over DP+CP.
-        torch.distributed.all_reduce(
-            sharded_norm_2,
-            op=torch.distributed.ReduceOp.SUM,
-            group=parallel_state.get_data_parallel_group(with_context_parallel=True),
-        )
-        norm_2 += sharded_norm_2
+    else:
+        sharded_norm_2 = torch.zeros((1,), dtype=torch.float32, device="cuda")
+    # Sum over all DP groups, including CP since distributed optimizer state is
+    # sharded jointly over DP+CP.
+    torch.distributed.all_reduce(
+        sharded_norm_2,
+        op=torch.distributed.ReduceOp.SUM,
+        group=parallel_state.get_data_parallel_group(with_context_parallel=True),
+    )
+    norm_2 += sharded_norm_2
 
     # Add norm contribution from expert layers in MoEs.
     if len(moe_params_data) > 0: