feat(dsv4): add MXFP4 QAT FP8 expert path

megatron/core/dist_checkpointing/utils.py

@@ -2,10 +2,13 @@
 
 """ Helpers for manipulating sharded tensors and sharded state dicts. """
 import logging
+import os
 from contextlib import contextmanager
 from time import time
 from typing import Dict, Optional, Tuple
 
+import torch
+
 from .dict_utils import dict_list_map_inplace, extract_matching_values, nested_values
 from .mapping import (
     LocalNonpersistentObject,
@@ -233,6 +236,27 @@ def _replace_prefixes(x):
     dict_list_map_inplace(_replace_prefixes, sharded_state_dict)
 
 
+def _dist_ckpt_fp8_dequant_dtype() -> Optional[torch.dtype]:
+    value = os.getenv("MEGATRON_DIST_CKPT_FP8_DEQUANT_DTYPE", "").strip().lower()
+    if value in ("", "default", "none"):
+        return None
+    dtype_by_name = {
+        "bf16": torch.bfloat16,
+        "bfloat16": torch.bfloat16,
+        "fp16": torch.float16,
+        "float16": torch.float16,
+        "half": torch.float16,
+        "fp32": torch.float32,
+        "float32": torch.float32,
+    }
+    if value not in dtype_by_name:
+        raise ValueError(
+            "MEGATRON_DIST_CKPT_FP8_DEQUANT_DTYPE must be one of "
+            f"{sorted(dtype_by_name)} plus default/none, got {value!r}"
+        )
+    return dtype_by_name[value]
+
+
 def force_all_tensors_to_non_fp8(sharded_state_dict: ShardedStateDict):
     """Force all tensors in state dict to be non-fp8.
 
@@ -241,9 +265,11 @@ def force_all_tensors_to_non_fp8(sharded_state_dict: ShardedStateDict):
     """
     from ..fp8_utils import dequantize_fp8_tensor, is_float8tensor  # Avoid circular import
 
+    dequant_dtype = _dist_ckpt_fp8_dequant_dtype()
+
     for v in nested_values(sharded_state_dict):
         if hasattr(v, "data") and is_float8tensor(v.data):
-            v.data = dequantize_fp8_tensor(v.data)
+            v.data = dequantize_fp8_tensor(v.data, dtype=dequant_dtype)
 
 
 fallback_logger = logging.getLogger(__name__)

megatron/core/distributed/finalize_model_grads.py

@@ -281,7 +281,7 @@ def reset_model_temporary_tensors(config: TransformerConfig, model: List[torch.n
     """
     for model_chunk in model:
         for module in get_attr_wrapped_model(model_chunk, 'modules')():
-            if config.moe_router_enable_expert_bias and hasattr(module, 'expert_bias'):
+            if config.moe_router_enable_expert_bias and getattr(module, 'expert_bias', None) is not None:
                 module.local_tokens_per_expert.zero_()
             if (
                 config.moe_router_load_balancing_type == "global_aux_loss"
@@ -299,7 +299,7 @@ def _update_router_expert_bias(model: List[torch.nn.Module], config: Transformer
     expert_bias_list = []
     for model_chunk in model:
         for module in get_attr_wrapped_model(model_chunk, 'modules')():
-            if hasattr(module, 'expert_bias'):
+            if getattr(module, 'expert_bias', None) is not None:
                 tokens_per_expert_list.append(module.local_tokens_per_expert)
                 expert_bias_list.append(module.expert_bias)
     # For hybrid models with both MoE and Dense layers, this list can be empty.
@@ -472,7 +472,7 @@ def finalize_model_grads(
     if config.timers is not None:
         config.timers('embedding-grads-all-reduce').stop()
 
-    if config.moe_router_enable_expert_bias:
+    if config.moe_router_enable_expert_bias and not config.freeze_e_score_correction_bias:
         _update_router_expert_bias(model, config)
 
     reset_model_temporary_tensors(config, model)