Save memory using main_param for moe in param_l2_norm (#2249)

Signed-off-by: lit <[email protected]>
Co-authored-by: oliver könig <[email protected]>

Author: BestJuly

megatron/training/utils.py

@@ -80,10 +80,21 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False):
                 continue
             assert is_not_tp_duplicate
             if not getattr(param, 'allreduce', True):
-                # TODO: Implement memory optimization for MoE parameters.
                 assert param_is_not_shared(param)
                 param = to_local_if_dtensor(param)
-                moe_params_data.append(param.data.float() if args.bf16 else param.data)
+                if args.bf16:
+                    if not force_create_fp32_copy and hasattr(param, 'main_param'):
+                        if getattr(param, 'main_param_sharded', False):
+                            if param.main_param is not None:
+                                sharded_params_data.append(param.main_param)
+                        else:
+                            moe_params_data.append(param.main_param)
+                    else:
+                        # Fallback to original logic of making a fp32 copy of the
+                        # parameter if `.main_param` attribute is not available.
+                        moe_params_data.append(param.data.float())
+                else:
+                    moe_params_data.append(param.data)
             else:
                 if param_is_not_shared(param):
                     param = to_local_if_dtensor(param)

tests/unit_tests/test_utils.py

@@ -15,8 +15,10 @@
 import megatron.training.utils as training_util
 from megatron.core import config
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
 from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
 from megatron.core.transformer import TransformerConfig
+from megatron.core.transformer.moe.moe_layer import MoELayer
 from tests.unit_tests.test_utilities import Utils
 
 success_string = "hello,world"
@@ -238,7 +240,7 @@ def test_cross_check_param_hashes_across_dp_replicas():
 @pytest.mark.flaky
 @pytest.mark.flaky_in_dev
 @pytest.mark.internal
-def test_param_norm(use_distributed_optimizer: bool):
+def test_param_norm_linear(use_distributed_optimizer: bool):
     world = int(os.getenv('WORLD_SIZE', '1'))
     rank = int(os.getenv('RANK', '0'))
 
@@ -286,6 +288,73 @@ def test_param_norm(use_distributed_optimizer: bool):
     _deinit_distributed()
 
 
+@pytest.mark.parametrize("use_distributed_optimizer", [False, True])
+@pytest.mark.flaky
+@pytest.mark.flaky_in_dev
+@pytest.mark.internal
+def test_param_norm_moe(use_distributed_optimizer: bool):
+    world = int(os.getenv('WORLD_SIZE', '1'))
+    rank = int(os.getenv('RANK', '0'))
+
+    # Setup: distributed, model, mock_args.
+    _init_distributed(world, rank)
+    Utils.initialize_model_parallel()
+    transformer_config = TransformerConfig(
+        num_layers=1,
+        hidden_size=12,
+        num_attention_heads=4,
+        num_moe_experts=2,
+        use_cpu_initialization=True,
+        moe_token_dispatcher_type="alltoall",
+        moe_router_topk=2,
+        moe_aux_loss_coeff=0.01,
+        moe_grouped_gemm=True,
+        moe_ffn_hidden_size=128,
+        add_bias_linear=False,
+        bf16=True,
+    )
+    transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+        num_experts=2, moe_grouped_gemm=True
+    )
+    model = MoELayer(transformer_config, transformer_layer_spec.submodules.mlp.submodules).to(
+        device='cuda'
+    )
+    model.requires_grad_(True)
+    # Initialize the model with all 1.0 for weights.
+    for param in model.parameters():
+        param.data.fill_(1.0)
+    ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=use_distributed_optimizer)
+    model = DistributedDataParallel(transformer_config, ddp_config, model)
+    for param in model.parameters():
+        assert param.requires_grad
+    mock_args = SimpleNamespace(bf16=True)
+
+    with mock.patch('megatron.training.utils.get_args', new=lambda: mock_args):
+        # Make sure norm is correct when `main_param` attribute is not available.
+        norm_no_fp32_copy = training_util.calc_params_l2_norm(model, force_create_fp32_copy=False)
+        norm_fp32_copy = training_util.calc_params_l2_norm(model, force_create_fp32_copy=True)
+        assert norm_no_fp32_copy == pytest.approx(norm_fp32_copy)
+
+        # Make sure norm is correct when `main_param` attribute is available.
+        optimizer_config = OptimizerConfig(
+            bf16=True, use_distributed_optimizer=use_distributed_optimizer
+        )
+        _ = get_megatron_optimizer(optimizer_config, [model])
+        for param in model.parameters():
+            # Only bf16/fp16 parameters get main_param attribute.
+            # Router weights are always fp32, so they won't have main_param.
+            if param.dtype in [torch.bfloat16, torch.float16]:
+                assert hasattr(param, 'main_param')
+                if use_distributed_optimizer:
+                    assert getattr(param, 'main_param_sharded', False)
+        norm_no_fp32_copy = training_util.calc_params_l2_norm(model, force_create_fp32_copy=False)
+        norm_fp32_copy = training_util.calc_params_l2_norm(model, force_create_fp32_copy=True)
+        assert norm_no_fp32_copy == pytest.approx(norm_fp32_copy)
+
+    # Teardown.
+    _deinit_distributed()
+
+
 @pytest.mark.flaky
 @pytest.mark.flaky_in_dev
 def test_straggler_detector():