Merge branch 'main' of ssh://gitlab-master.nvidia.com:12051/ADLR/mega…

…tron-lm into moe-tests
Rivendile · Oct 5, 2023 · e1e6926 · e1e6926
2 parents be6f63e + 4e79e71
commit e1e6926
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/megatron/utils.py b/megatron/utils.py
@@ -54,14 +54,13 @@ def calc_params_l2_norm(model):
     params_data = []
     for model_ in model:
         for param in model_.parameters():
+            is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param)
             if args.expert_parallel and mpu.get_data_parallel_rank() > 0:
-                if not getattr(param, 'allreduce', True):
+                if not getattr(param, 'allreduce', True) and is_not_tp_duplicate:
                     assert param_is_not_shared(param)
-                    assert param_is_not_tensor_parallel_duplicate(param)
                     params_data.append(param.data.float() if args.bf16 else param.data)
             else:
                 is_not_shared = param_is_not_shared(param)
-                is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param)
                 if is_not_shared and is_not_tp_duplicate:
                     params_data.append(param.data.float() if args.bf16 else param.data)