Extend FP32 MoE numerics to low-precision recipes and registered activations

.gitignore

@@ -16,6 +16,7 @@ onelogger.err
 runs/
 /test_cases/
 **/dist/
+.idea
 
 # Sphinx documentation
 docs/_build

megatron/core/dist_checkpointing/dict_utils.py

@@ -217,8 +217,25 @@ def dict_list_map_outplace(f: Callable[[U], V], x: Union[Dict, List, U]) -> Unio
         return f(x)
 
 
+def _is_optimizer_param_state_key(key: Tuple) -> bool:
+    """Check if key path matches (..., 'optimizer', ..., 'param_state', ...)."""
+    try:
+        idx = key.index("optimizer")
+    except ValueError:
+        return False
+    try:
+        key.index("param_state", idx + 1)
+    except ValueError:
+        return False
+    return True
+
+
 def merge(x1: Union[dict, list], x2: Union[dict, list], key: Tuple[Union[str, int], ...] = ()):
-    """Merges dicts and lists recursively."""
+    """Merges dicts and lists recursively.
+
+    For optimizer param_state paths, allows x1 to be longer than x2
+    (dp_reshardable padding entries are truncated).
+    """
     if isinstance(x1, dict) and isinstance(x2, dict):
         for k, v2 in x2.items():
             if k not in x1:
@@ -227,10 +244,13 @@ def merge(x1: Union[dict, list], x2: Union[dict, list], key: Tuple[Union[str, in
                 x1[k] = merge(x1[k], v2, key=key + (k,))
     elif isinstance(x1, list) and isinstance(x2, list):
         if len(x1) != len(x2):
-            raise ValueError(
-                f"Cannot merge two lists with different lengths ({len(x1)} and {len(x2)}, "
-                f"encountered at level {key})"
-            )
+            if _is_optimizer_param_state_key(key) and len(x1) > len(x2):
+                del x1[len(x2):]
+            else:
+                raise ValueError(
+                    f"Cannot merge two lists with different lengths ({len(x1)} and {len(x2)}, "
+                    f"encountered at level {key})"
+                )
         for i, v2 in enumerate(x2):
             x1[i] = merge(x1[i], v2, key=key + (i,))
     else:

megatron/core/dist_checkpointing/mapping.py

@@ -131,7 +131,12 @@ class with `from_rank_offsets` or `from_rank_offsets_flat` constructors.
                     )
 
         if self.flattened_range is not None:
-            raise CheckpointingException("ShardedTensor.flattened_range is not supported.")
+            if not _logged_deprecations.get("flattened_range", False):
+                logger.warning(
+                    "ShardedTensor.flattened_range is deprecated."
+                    " Use latest DistributedOptimizer formats."
+                )
+                _logged_deprecations["flattened_range"] = True
 
     @property
     def has_regular_grid(self):

megatron/core/dist_checkpointing/strategies/common.py

@@ -86,7 +86,7 @@ def load_common(self, checkpoint_dir: Union[str, Path]):
                 msc = MultiStorageClientFeature.import_package()
                 return msc.torch.load(load_path, map_location='cpu')
             else:
-                return torch.load(load_path, map_location='cpu')
+                return torch.load(load_path, map_location='cpu', weights_only=False)
         except FileNotFoundError as e:
             err_msg = f'Common file {load_path} does not exist'
             if MultiStorageClientFeature.is_enabled():

megatron/core/dist_checkpointing/strategies/torch.py

@@ -385,6 +385,10 @@ def _replace_sharded_keys_with_state_dict_keys(
     """Inverse of _replace_state_dict_keys_with_sharded_keys."""
     recovered_sd = {}
     for k, tensors in state_dict.items():
+        if isinstance(tensors, io.BytesIO):
+            # TE FP8 _extra_state arrives as a bare io.BytesIO, so len(tensors) raises.
+            # An empty uint8 tensor makes TE.set_extra_state skip restoring FP8 state.
+            tensors = [torch.empty(0, dtype=torch.uint8)]
         assert len(tensors) == len(rename_mapping[k])
         for ten, recovered_k in zip(tensors, rename_mapping[k]):
             recovered_sd[recovered_k] = ten
@@ -503,10 +507,12 @@ def __init__(
     def _validate_global_shapes(self, metadata, sharded_tensors):
         for sh_ten in sharded_tensors:
             if sh_ten.key not in metadata.state_dict_metadata:
-                raise KeyError(
-                    f"{sh_ten.key} from model not in state dict:"
-                    f" {sorted(metadata.state_dict_metadata.keys())}"
-                )
+                # raise KeyError(
+                #     f"{sh_ten.key} from model not in state dict:"
+                #     f" {sorted(metadata.state_dict_metadata.keys())}"
+                # )
+                print(f"{sh_ten.key} from model not in state dict, will skip")
+                continue
             loaded_shape = metadata.state_dict_metadata[sh_ten.key].size
             expected_shape = sh_ten.global_shape
             if loaded_shape != expected_shape:
@@ -530,7 +536,7 @@ def _temporarily_bypass_shape_validation(self):
         tensor_metadata = self.metadata.state_dict_metadata
         metadata_with_sizes = [
             (tensor_metadata[key], tensor_metadata[key].size, sharded_tensor)
-            for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items()
+            for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items() if key in tensor_metadata
         ]
         try:
             # Temporarily set sizes to expected shapes
@@ -802,6 +808,7 @@ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> St
             planner=MCoreLoadPlanner(
                 shapes_validation_sharded_tensors=flexible_shape_sharded_tensors,
                 allow_shape_mismatch_sharded_tensors=allow_shape_mismatch_sharded_tensors,
+                allow_partial_load=True,
             ),
         )
 

megatron/core/distributed/distributed_data_parallel.py

@@ -11,6 +11,7 @@
 from ..process_groups_config import ProcessGroupCollection
 from ..transformer.cuda_graphs import is_graph_capturing
 from ..transformer.transformer_config import TransformerConfig
+from miles_megatron_plugins.true_on_policy.contracts import resolve_true_on_policy_runtime_policy
 from ..utils import log_single_rank
 from .data_parallel_base import _BaseDataParallel
 from .distributed_data_parallel_config import DistributedDataParallelConfig
@@ -19,6 +20,29 @@
 logger = logging.getLogger(__name__)
 
 
+def _first_cp_comm_type(cp_comm_type):
+    if isinstance(cp_comm_type, list):
+        return cp_comm_type[0] if cp_comm_type else None
+    return cp_comm_type
+
+
+def _use_true_on_policy_ulysses_cp_gradient_scaling(config: TransformerConfig) -> bool:
+    """Return true when CP loss scaling is deferred to the CP gradient sum."""
+
+    return (
+        resolve_true_on_policy_runtime_policy(config).defer_ulysses_cp_loss_scaling_to_grad_sum
+        and getattr(config, "context_parallel_size", 1) > 1
+        and _first_cp_comm_type(getattr(config, "cp_comm_type", None)) == "a2a"
+        and not getattr(config, "calculate_per_token_loss", False)
+    )
+
+
+def _dense_gradient_scaling_factor(config: TransformerConfig, dp_cp_world_size: int) -> float:
+    if _use_true_on_policy_ulysses_cp_gradient_scaling(config):
+        return float(getattr(config, "context_parallel_size", 1)) / float(dp_cp_world_size)
+    return 1.0 / float(dp_cp_world_size)
+
+
 class DistributedDataParallel(_BaseDataParallel):
     """
     DDP wrapper which stores grads in contiguous buffers. Also has option of overlapping
@@ -45,6 +69,8 @@ def __init__(
         module: torch.nn.Module,
         disable_bucketing: bool = False,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        disable_grad_buffers_cpu_backup: bool = False,
+        disable_param_buffers_cpu_backup: bool = False,
     ):
         super().__init__(config=config, module=module)
         if has_config_logger_enabled(config):
@@ -125,7 +151,10 @@ def __init__(
                 expert_parallel_params.append(param)
 
         def _allocate_buffers_for_parameters(
-            input_params, data_parallel_group, gradient_scaling_factor
+            input_params,
+            data_parallel_group,
+            gradient_scaling_factor,
+            target_gradient_scaling_factor,
         ):
             param_and_grad_dtype_to_params = {}
             param_and_grad_dtype_to_offsets = {}
@@ -172,7 +201,6 @@ def _allocate_buffers_for_parameters(
                 param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)] = indices
 
             if not config.calculate_per_token_loss:
-                target_gradient_scaling_factor = 1.0 / self.dp_cp_group.size()
                 if self.ddp_config.average_in_collective:
                     if self.ddp_config.num_distributed_optimizer_instances == 1:
                         # Collective is averaging gradients in collective with data_parallel_group.
@@ -209,6 +237,8 @@ def _allocate_buffers_for_parameters(
                         param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)],
                         self.ddp_config.nccl_ub,
                         pg_collection,
+                        disable_grad_buffers_cpu_backup=disable_grad_buffers_cpu_backup,
+                        disable_param_buffers_cpu_backup=disable_param_buffers_cpu_backup,
                     )
                 )
 
@@ -284,12 +314,19 @@ def _allocate_buffers_for_parameters(
             else:
                 data_parallel_world_size = self.dp_cp_group.size()
 
-                gradient_scaling_factor = 1.0 / data_parallel_world_size
+                gradient_scaling_factor = _dense_gradient_scaling_factor(
+                    config, data_parallel_world_size
+                )
                 expert_gradient_scaling_factor = 1.0 / data_parallel_world_size
 
         # Allocate the param+grad buffers for dense params' grads.
         self.buffers, self.bucket_groups = _allocate_buffers_for_parameters(
-            dense_params, self.intra_dp_cp_group, gradient_scaling_factor=gradient_scaling_factor
+            dense_params,
+            self.intra_dp_cp_group,
+            gradient_scaling_factor=gradient_scaling_factor,
+            target_gradient_scaling_factor=_dense_gradient_scaling_factor(
+                config, self.dp_cp_group.size()
+            ),
         )
 
         # Allocate separate param+grad buffers for expert parallel params' grads.
@@ -298,6 +335,7 @@ def _allocate_buffers_for_parameters(
                 expert_parallel_params,
                 self.intra_expt_dp_group,
                 gradient_scaling_factor=expert_gradient_scaling_factor,
+                target_gradient_scaling_factor=1.0 / self.dp_cp_group.size(),
             )
         )
 

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,8 @@ 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:
+                assert module.local_tokens_per_expert 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.
@@ -473,7 +474,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)

megatron/core/distributed/param_and_grad_buffer.py

@@ -3,6 +3,7 @@
 import functools
 import logging
 import math
+import os
 import warnings
 from contextlib import nullcontext
 from enum import Enum
@@ -17,6 +18,7 @@
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.rerun_state_machine import get_rerun_state_machine
 from megatron.core.utils import log_single_rank
+from miles_megatron_plugins.true_on_policy.debug import dump_param_and_grad_buffer_debug
 
 from ..fp8_utils import (
     is_float8tensor,
@@ -93,6 +95,7 @@ def __init__(
         gradient_scaling_factor: float,
         bucket_id: int,
         param_index_map: Dict[torch.nn.Parameter, tuple],
+        param_to_name: Dict[torch.nn.Parameter, str],
     ):
         self.params_list = params
         self.params = set(params)
@@ -106,6 +109,7 @@ def __init__(
         self.numel_unpadded = numel_unpadded
         self.gradient_scaling_factor = gradient_scaling_factor
         self.bucket_id = bucket_id
+        self.param_to_name = param_to_name
         # Derive bucket-local param offsets from the global param_index_map.
         self.param_to_index = {}
         for param in params:
@@ -199,6 +203,7 @@ def __init__(
         # or bucket.grad_data.
         self.cached_param_buffer_shard_list = [None] * len(self.buckets)
         self.cached_grad_buffer_shard_list = [None] * len(self.buckets)
+        self._grad_debug_dumped_buckets = set()
 
     def reset(self):
         """
@@ -218,8 +223,16 @@ def check_grads(self, check_for_nan_or_inf, check_for_large):
         all-reduce / reduce-scatter.
         """
         rerun_state_machine = get_rerun_state_machine()
+        grad_debug_dir = os.environ.get("MILES_GRAD_DEBUG_DIR")
         for i in range(len(self.buckets)):
             grad_norm = self.buckets[i].grad_data.norm(p=2)
+            if grad_debug_dir and not torch.isfinite(grad_norm).item():
+                dump_param_and_grad_buffer_debug(
+                    self,
+                    bucket_index=i,
+                    grad_norm=grad_norm,
+                    grad_debug_dir=grad_debug_dir,
+                )
             # check for NaN, Inf and unexpectedly large grads
             if check_for_nan_or_inf:
                 rerun_state_machine.validate_result(
@@ -599,6 +612,8 @@ def __init__(
         param_indices: List[int],
         nccl_ub: bool,
         pg_collection: Optional[ProcessGroupCollection] = None,
+        disable_grad_buffers_cpu_backup: bool = False,
+        disable_param_buffers_cpu_backup: bool = False,
     ):
 
         if pg_collection is None:
@@ -614,6 +629,7 @@ def __init__(
         self.ddp_config = ddp_config
         self.params = params
         self.param_indices = param_indices
+        self.param_to_name = param_to_name
 
         # Check that params are unique.
         unique_params = set()
@@ -754,6 +770,9 @@ def _does_param_require_new_bucket(param):
         self.param_data = None
 
         if self.nccl_ub:
+            assert not disable_grad_buffers_cpu_backup and not disable_param_buffers_cpu_backup, (
+                "disable_grad/param_buffers_cpu_backup is not supported with nccl_ub=True"
+            )
             # If nccl_ub is True, use nccl_allocator to allocate memory for param_data/grad_data.
             nccl_allocator.init()
             pool = nccl_allocator.create_nccl_mem_pool(
@@ -773,8 +792,23 @@ def _does_param_require_new_bucket(param):
             torch.distributed.barrier()
         else:
             # If nccl_ub is False, mem_alloc_context is nullcontext.
+            # Individual param/grad contexts below handle TMS regions separately.
             mem_alloc_context = nullcontext
 
+        def _make_no_backup_context(tag, disable):
+            if disable:
+                from torch_memory_saver import torch_memory_saver
+
+                return partial(
+                    torch_memory_saver.region,
+                    tag=tag,
+                    enable_cpu_backup=False,
+                )
+            return nullcontext
+
+        grad_mem_alloc_context = _make_no_backup_context("grad_buffer", disable_grad_buffers_cpu_backup)
+        param_mem_alloc_context = _make_no_backup_context("param_buffer", disable_param_buffers_cpu_backup)
+
         with mem_alloc_context():
             # For MXFP8 param: Create a shared buffer for param AG and grad RS for memory efficiency
             # The buffer is mapped to weight gradients whose dtype is either bf16 or FP32.
@@ -797,18 +831,20 @@ def _does_param_require_new_bucket(param):
             else:
                 # Only re-map param tensors if using distributed optimizer.
                 if self.ddp_config.use_distributed_optimizer:
-                    self.param_data = torch.zeros(
+                    with param_mem_alloc_context():
+                        self.param_data = torch.zeros(
+                            self.numel,
+                            dtype=self.param_dtype,
+                            device=torch.cuda.current_device(),
+                            requires_grad=False,
+                        )
+                with grad_mem_alloc_context():
+                    self.grad_data = torch.zeros(
                         self.numel,
-                        dtype=self.param_dtype,
+                        dtype=self.grad_dtype,
                         device=torch.cuda.current_device(),
                         requires_grad=False,
                     )
-                self.grad_data = torch.zeros(
-                    self.numel,
-                    dtype=self.grad_dtype,
-                    device=torch.cuda.current_device(),
-                    requires_grad=False,
-                )
 
         self.grad_data_size = 0
         self.param_data_size = 0
@@ -952,6 +988,7 @@ def _new_bucket(
             gradient_scaling_factor=self.gradient_scaling_factor,
             bucket_id=bucket_id,
             param_index_map=self.param_index_map,
+            param_to_name=self.param_to_name,
         )
         for bucket_param in bucket_params:
             assert bucket_param not in self.param_to_bucket