run_megatron: dump parameter grads, fix CP loss normalisation, skip DDP on forward-only runs

miles/utils/debug_utils/run_megatron/worker/batch.py

@@ -4,6 +4,9 @@
 from typing import Any
 
 import torch
+import torch.distributed as dist
+
+from miles.backends.training_utils.parallel import get_parallel_state
 
 
 def prepare_batch(
@@ -95,17 +98,32 @@ def loss_func(
     labels: torch.Tensor,
     output_tensor: torch.Tensor,
 ) -> tuple[torch.Tensor, dict[str, Any]]:
-    """Cross-entropy loss for forward-backward pipeline schedule.
+    """Cross-entropy loss, normalised by the *global* valid-token count.
+
+    A local ``reduction="mean"`` is wrong under context parallelism: the labels are
+    sharded, so each rank divides by its own valid count rather than the global
+    one, and every activation gradient comes out scaled by global/local -- about 2x
+    at CP=2, which shows up as a ~0.5 relative difference against the CP=1 run and
+    swamps any real discrepancy.
 
-    Uses ignore_index=-100 to handle CP-aware label masking.
+    Summing locally and dividing by the CP-reduced count instead makes each rank
+    contribute its true share of the global mean: the per-rank losses add up to the
+    CP=1 loss, and the gradients match position for position.
     """
     logits: torch.Tensor = output_tensor.float()
     vocab_size: int = logits.size(-1)
+    flat_labels: torch.Tensor = labels.view(-1)
 
-    loss: torch.Tensor = torch.nn.functional.cross_entropy(
+    loss_sum: torch.Tensor = torch.nn.functional.cross_entropy(
         logits.view(-1, vocab_size),
-        labels.view(-1),
+        flat_labels,
         ignore_index=-100,
-        reduction="mean",
+        reduction="sum",
     )
+    num_valid: torch.Tensor = (flat_labels != -100).sum()
+    cp = get_parallel_state().cp
+    if cp.size > 1:
+        dist.all_reduce(num_valid, group=cp.group)
+
+    loss: torch.Tensor = loss_sum / num_valid.clamp(min=1)
     return loss, {"loss": loss.detach()}

miles/utils/debug_utils/run_megatron/worker/main.py

@@ -97,6 +97,7 @@ def main() -> None:
         )
 
     save_replay_data(script, rank=rank)
+    _dump_model_params_and_grads(model)
     _finalize_dumper()
 
     if rank == 0:
@@ -127,13 +128,20 @@ def _initialize_megatron(args: argparse.Namespace) -> None:
     args.__dict__.setdefault("debug_deterministic_collective", False)
     set_default_megatron_args(args)
     validate_args(args)
+    # Megatron's validate_args unconditionally resets this to False; the miles
+    # backend always runs with variable_seq_lengths (miles/utils/arguments.py),
+    # so mirror it here to keep the worker consistent with the training path.
+    args.variable_seq_lengths = True
 
     init(args)
 
 
 def _build_and_load_model(args: argparse.Namespace, script: WorkerScriptArgs) -> list[Any]:
     model_provider: Callable[..., Any] = get_model_provider_func(args, role=script.role)
-    model: list[Any] = get_model(model_provider, ModelType.encoder_or_decoder)
+    # Forward-only runs never call backward, so skip the DDP wrapper and its
+    # per-parameter grad buffers (which otherwise OOM for large models even
+    # when only logits/logprobs are needed).
+    model: list[Any] = get_model(model_provider, ModelType.encoder_or_decoder, wrap_with_ddp=script.run_backward)
 
     if args.load is not None:
         load_checkpoint(
@@ -145,7 +153,7 @@ def _build_and_load_model(args: argparse.Namespace, script: WorkerScriptArgs) ->
         )
 
     for m in model:
-        m.train()
+        m.train(script.run_backward)
     return model
 
 
@@ -210,6 +218,33 @@ def _print_config(args: argparse.Namespace, script: WorkerScriptArgs) -> None:
     print(f"[worker] run_backward={script.run_backward}, role={script.role}", flush=True)
 
 
+def _dump_model_params_and_grads(model: list[Any]) -> None:
+    """Dump parameter gradients so runs can be diffed per parameter.
+
+    ``DUMPER_ENABLE_MODEL_GRAD`` only takes effect inside ``dump_model``, which
+    the training path reaches through ``DumperMegatronUtil.finalize`` but this
+    standalone worker never called -- so ``--run-backward`` produced no
+    parameter gradients at all.
+
+    ``get_grad`` is not optional here. The dumper's default is a bare
+    ``param.grad`` read, and under Megatron DDP the gradient lives in
+    ``param.main_grad`` with ``param.grad`` left as None, so omitting it yields a
+    silently empty dump. The heavier distributed-optimizer shard gather in
+    ``dumper_utils`` is only needed for the FT trainer, whose reduce-scatter
+    leaves each rank holding a partial buffer.
+    """
+    if os.environ.get("DUMPER_ENABLE", "0") != "1":
+        return
+    if os.environ.get("DUMPER_ENABLE_MODEL_GRAD", "0") != "1":
+        return
+    assert len(model) == 1, f"model dump does not support virtual pipeline parallelism ({len(model)} chunks)"
+
+    def get_grad(param: torch.nn.Parameter) -> torch.Tensor | None:
+        return param.grad if param.grad is not None else getattr(param, "main_grad", None)
+
+    dumper.dump_model(model[0], get_grad=get_grad)
+
+
 def _finalize_dumper() -> None:
     """Step + disable dumper after forward/backward."""
     if os.environ.get("DUMPER_ENABLE", "0") == "1":