[3/15] Select SGLang parity paths from transformer config

megatron/core/models/gpt/gpt_layer_specs.py

@@ -82,6 +82,29 @@
     HAVE_APEX = False
 
 
+_DENSE_MLP_SHARDED_STATE_DICT_KEYS_MAP = {
+    "mlp.0.weight": "mlp.linear_fc1.layer_norm_weight",
+    "mlp.0.bias": "mlp.linear_fc1.layer_norm_bias",
+    "mlp.1.basic_ops.0.weight": "mlp.linear_fc1.weight",
+    "mlp.1.basic_ops.1.bias": "mlp.linear_fc1.bias",
+    "mlp.3.basic_ops.0.weight": "mlp.linear_fc2.weight",
+    "mlp.3.basic_ops.1.bias": "mlp.linear_fc2.bias",
+}
+
+
+def _dense_mlp_sharded_state_dict_keys_map(num_experts: Optional[int]) -> dict[str, str]:
+    if num_experts is not None:
+        return {}
+    return dict(_DENSE_MLP_SHARDED_STATE_DICT_KEYS_MAP)
+
+
+def _local_layer_norm_sharded_state_dict_keys_map(num_experts: Optional[int]) -> dict[str, str]:
+    keys_map = {"input_layernorm.": "self_attention.linear_qkv.layer_norm_"}
+    if num_experts is None:
+        keys_map["pre_mlp_layernorm."] = "mlp.linear_fc1.layer_norm_"
+    return keys_map
+
+
 def get_gpt_layer_with_inference_spec(
     qk_layernorm: Optional[bool] = False,
     multi_latent_attention: Optional[bool] = False,
@@ -166,14 +189,7 @@ def get_gpt_layer_with_inference_spec(
                 pre_mlp_layernorm=IdentityOp,
                 mlp=mlp,
                 mlp_bda=get_bias_dropout_add,
-                sharded_state_dict_keys_map={
-                    "mlp.0.weight": "mlp.linear_fc1.layer_norm_weight",
-                    "mlp.0.bias": "mlp.linear_fc1.layer_norm_bias",
-                    "mlp.1.basic_ops.0.weight": "mlp.linear_fc1.weight",
-                    "mlp.1.basic_ops.1.bias": "mlp.linear_fc1.bias",
-                    "mlp.3.basic_ops.0.weight": "mlp.linear_fc2.weight",
-                    "mlp.3.basic_ops.1.bias": "mlp.linear_fc2.bias",
-                },
+                sharded_state_dict_keys_map=dict(_DENSE_MLP_SHARDED_STATE_DICT_KEYS_MAP),
             ),
         )
 
@@ -308,14 +324,7 @@ def get_gpt_layer_with_transformer_engine_spec(
                 mlp=mlp,
                 mlp_bda=get_bias_dropout_add,
                 post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp,
-                sharded_state_dict_keys_map={
-                    "mlp.0.weight": "mlp.linear_fc1.layer_norm_weight",
-                    "mlp.0.bias": "mlp.linear_fc1.layer_norm_bias",
-                    "mlp.1.basic_ops.0.weight": "mlp.linear_fc1.weight",
-                    "mlp.1.basic_ops.1.bias": "mlp.linear_fc1.bias",
-                    "mlp.3.basic_ops.0.weight": "mlp.linear_fc2.weight",
-                    "mlp.3.basic_ops.1.bias": "mlp.linear_fc2.bias",
-                },
+                sharded_state_dict_keys_map=_dense_mlp_sharded_state_dict_keys_map(num_experts),
             ),
         )
 
@@ -455,10 +464,9 @@ def get_gpt_layer_local_spec(
                 pre_mlp_layernorm=layer_norm,
                 mlp=mlp,
                 mlp_bda=bias_dropout_add,
-                sharded_state_dict_keys_map={
-                    "input_layernorm.": "self_attention.linear_qkv.layer_norm_",
-                    "pre_mlp_layernorm.": "mlp.linear_fc1.layer_norm_",
-                },
+                sharded_state_dict_keys_map=_local_layer_norm_sharded_state_dict_keys_map(
+                    num_experts
+                ),
             ),
         )
 

megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py

@@ -375,10 +375,21 @@ def _update_fp32_params_by_new_state(self):
 
     def update_fp32_param_by_new_param(self):
         """
-        Update the fp32 parameters by the new parameters.
+        Update optimizer-owned parameter copies from the live model parameters.
+
+        ``DistributedOptimizer.reload_model_params`` calls this after checkpoint
+        loading. For offloaded fp32 model parameters there is no separate
+        ``param_to_fp32_param`` entry: the CPU copy tracked by
+        ``gpu_params_map_cpu_copy`` is the optimizer-owned parameter. Refresh it
+        as well, otherwise the first optimizer step can copy constructor-time
+        values back over checkpoint-loaded weights.
         """
         for param, fp32_param in self.param_to_fp32_param.items():
             fp32_param.data.copy_(param)
+        for param, cpu_copy in self.gpu_params_map_cpu_copy.items():
+            if param in self.param_to_fp32_param:
+                continue
+            cpu_copy.data.copy_(param)
 
     def _register_load_state_dict_hooks(self):
         def pre_load_state_dict_hook(self, state_dict):

megatron/core/transformer/transformer_block.py

@@ -30,14 +30,14 @@
     get_transformer_layer_offset,
 )
 from megatron.core.transformer.utils import sharded_state_dict_default
-from miles_megatron_plugins.true_on_policy.contracts import resolve_true_on_policy_runtime_policy
-from miles_megatron_plugins.true_on_policy.sglang_backend import SGLangFinalRMSNorm, SGLangNorm
 from megatron.core.utils import (
     WrappedTensor,
     deprecate_inference_params,
     get_pg_rank,
     make_viewless_tensor,
 )
+from miles_megatron_plugins.true_on_policy.contracts import resolve_true_on_policy_runtime_policy
+from miles_megatron_plugins.true_on_policy.sglang_backend import SGLangFinalRMSNorm, SGLangNorm
 
 try:
     import transformer_engine.pytorch as te  # pylint: disable=unused-import
@@ -117,7 +117,7 @@ def hook(grad: Tensor) -> Tensor:
         path = Path(dump_dir) / f"rank_{rank}_{counter:05d}_{name}.pt"
         path.parent.mkdir(parents=True, exist_ok=True)
         torch.save({"stats": stats}, path)
-        print(f"[MILES_FINAL_LAYERNORM_BACKWARD_DEBUG] {stats} wrote {path}", flush=True)
+        logger.info("[MILES_FINAL_LAYERNORM_BACKWARD_DEBUG] %s wrote %s", stats, path)
         return grad
 
     tensor.register_hook(hook)
@@ -519,6 +519,9 @@ def _checkpointed_forward(
     ):
         """Forward method with activation checkpointing."""
 
+        true_on_policy_policy = resolve_true_on_policy_runtime_policy(self.config)
+        use_sglang_residual_pair = true_on_policy_policy.use_sglang_residual_pair
+
         def custom(start: int, end: int):
             def custom_forward(
                 hidden_states,
@@ -527,7 +530,11 @@ def custom_forward(
                 context_mask,
                 rotary_pos_emb,
                 padding_mask=None,
+                sglang_residual=None,
             ):
+                if use_sglang_residual_pair and sglang_residual is not None:
+                    hidden_states._sglang_residual = sglang_residual
+
                 for index in range(start, end):
                     layer = self._get_layer(index)
 
@@ -559,13 +566,25 @@ def custom_forward(
                             packed_seq_params=packed_seq_params,
                             padding_mask=padding_mask,
                         )
+                if use_sglang_residual_pair:
+                    return hidden_states, context, getattr(hidden_states, "_sglang_residual", None)
                 return hidden_states, context
 
             return custom_forward
 
-        def checkpoint_handler(forward_func):
+        def checkpoint_handler(forward_func, sglang_residual=None):
             """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""
-            true_on_policy_policy = resolve_true_on_policy_runtime_policy(self.config)
+            checkpoint_args = (
+                hidden_states,
+                attention_mask,
+                context,
+                context_mask,
+                rotary_pos_emb,
+                padding_mask,
+            )
+            if sglang_residual is not None:
+                checkpoint_args = (*checkpoint_args, sglang_residual)
+
             if true_on_policy_policy.use_ulysses_cp_recompute_fallback:
                 mode = _get_sglang_cp_recompute_mode()
                 global _WARNED_SGLANG_CP_RECOMPUTE_FALLBACK
@@ -577,14 +596,7 @@ def checkpoint_handler(forward_func):
                             "'non_reentrant' or 'reentrant' to override."
                         )
                         _WARNED_SGLANG_CP_RECOMPUTE_FALLBACK = True
-                    return forward_func(
-                        hidden_states,
-                        attention_mask,
-                        context,
-                        context_mask,
-                        rotary_pos_emb,
-                        padding_mask,
-                    )
+                    return forward_func(*checkpoint_args)
                 if mode == "non_reentrant":
                     if not _WARNED_SGLANG_CP_RECOMPUTE_FALLBACK:
                         logger.info(
@@ -593,15 +605,7 @@ def checkpoint_handler(forward_func):
                         )
                         _WARNED_SGLANG_CP_RECOMPUTE_FALLBACK = True
                     return torch_checkpoint(
-                        forward_func,
-                        hidden_states,
-                        attention_mask,
-                        context,
-                        context_mask,
-                        rotary_pos_emb,
-                        padding_mask,
-                        use_reentrant=False,
-                        preserve_rng_state=True,
+                        forward_func, *checkpoint_args, use_reentrant=False, preserve_rng_state=True
                     )
                 if mode != "reentrant":
                     raise ValueError(
@@ -616,34 +620,29 @@ def checkpoint_handler(forward_func):
                     self.config.distribute_saved_activations,
                     tensor_parallel.random.get_cuda_rng_tracker,
                     self.pg_collection.tp,
-                    hidden_states,
-                    attention_mask,
-                    context,
-                    context_mask,
-                    rotary_pos_emb,
-                    padding_mask,
+                    *checkpoint_args,
                 )
             else:
                 return tensor_parallel.checkpoint(
-                    forward_func,
-                    self.config.distribute_saved_activations,
-                    hidden_states,
-                    attention_mask,
-                    context,
-                    context_mask,
-                    rotary_pos_emb,
-                    padding_mask,
+                    forward_func, self.config.distribute_saved_activations, *checkpoint_args
                 )
 
+        sglang_residual = (
+            getattr(hidden_states, "_sglang_residual", None) if use_sglang_residual_pair else None
+        )
         if self.config.recompute_method == 'uniform':
             # Uniformly divide the total number of Transformer layers and checkpoint
             # the input activation of each divided chunk.
             # A method to further reduce memory usage reducing checkpoints.
             layer_idx = 0
             while layer_idx < self.num_layers_per_pipeline_rank:
-                hidden_states, context = checkpoint_handler(
-                    custom(layer_idx, layer_idx + self.config.recompute_num_layers)
+                checkpoint_output = checkpoint_handler(
+                    custom(layer_idx, layer_idx + self.config.recompute_num_layers), sglang_residual
                 )
+                if use_sglang_residual_pair:
+                    hidden_states, context, sglang_residual = checkpoint_output
+                else:
+                    hidden_states, context = checkpoint_output
 
                 layer_idx += self.config.recompute_num_layers
 
@@ -663,14 +662,32 @@ def checkpoint_handler(forward_func):
                     layer_idx >= recompute_skip_num_layers
                     and layer_idx < self.config.recompute_num_layers + recompute_skip_num_layers
                 ):
-                    hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1))
+                    checkpoint_output = checkpoint_handler(
+                        custom(layer_idx, layer_idx + 1), sglang_residual
+                    )
+                    if use_sglang_residual_pair:
+                        hidden_states, context, sglang_residual = checkpoint_output
+                    else:
+                        hidden_states, context = checkpoint_output
                 else:
-                    hidden_states, context = custom(layer_idx, layer_idx + 1)(
-                        hidden_states, attention_mask, context, context_mask, rotary_pos_emb
+                    forward_output = custom(layer_idx, layer_idx + 1)(
+                        hidden_states,
+                        attention_mask,
+                        context,
+                        context_mask,
+                        rotary_pos_emb,
+                        padding_mask,
+                        sglang_residual,
                     )
+                    if use_sglang_residual_pair:
+                        hidden_states, context, sglang_residual = forward_output
+                    else:
+                        hidden_states, context = forward_output
         else:
             raise ValueError("Invalid activation recompute method.")
 
+        if use_sglang_residual_pair and sglang_residual is not None:
+            hidden_states._sglang_residual = sglang_residual
         return hidden_states
 
     def set_input_tensor(self, input_tensor: Tensor):