add compatibility toggle moe_router_padding_for_fp8 back

Signed-off-by: Zhongbo Zhu <[email protected]>

Author: zhongbozhu

megatron/core/transformer/transformer_config.py

@@ -468,6 +468,10 @@ class TransformerConfig(ModelParallelConfig):
     is a multiple of 16/32 for quantized precision (e.g., FP8, FP4). This can remove the explicit
     padding in the GroupedMLP layer."""
 
+    moe_router_padding_for_fp8: Optional[bool] = False
+    """[Compatibility alias for moe_router_padding_for_quantization]
+    Enabling this will also enable moe_router_padding_for_quantization."""
+
     moe_router_num_groups: Optional[int] = None
     """Number of groups to divide experts into for group-limited routing.
     When using group-limited routing:
@@ -1306,6 +1310,13 @@ def __post_init__(self):
                     f"but your version is {get_te_version()}."
                 )
 
+        if self.moe_router_padding_for_fp8:
+            # enable moe_router_padding_for_quantization
+            warnings.warn(
+                "--moe-router-padding-for-fp8 is going to be deprecated. Use --moe-router-padding-for-quantization instead."
+            )
+            self.moe_router_padding_for_quantization = True
+
         if self.moe_router_padding_for_quantization:
             if self.fp8 is None and self.fp4 is None:
                 raise ValueError("fp8/fp4 must be specified when moe_router_padding_for_quantization is True.")

megatron/training/arguments.py

@@ -3066,6 +3066,9 @@ def _add_moe_args(parser):
                        'dropless training with FP8/FP4 precision when num_local_experts > 1. This is a more '
                        'efficient way to pad for FP8/FP4 which eliminates the explicit padding in the '
                        'GroupedMLP layer.')
+    group.add_argument('--moe-router-padding-for-fp8', action='store_true',
+                       help='[Compatibility alias for --moe-router-padding-for-quantization] '
+                       'Enabling this will also enable --moe-router-padding-for-quantization.')
     group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0,
                        help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.')
     group.add_argument('--moe-z-loss-coeff', type=float, default=None,