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

[yueming-yuan]

Three model-agnostic fixes to the standalone run_megatron worker, found while
using it for cross-stack forward/backward alignment.

1. --run-backward produced no parameter gradients at all

The worker called dumper.register_non_intrusive_dumper(m) (activations) and
_finalize_dumper() (step() + configure(enable=False)), but never
dump_model. DUMPER_ENABLE_MODEL_GRAD is only consumed inside dump_model
(enable_curr_grad=self._config.enable_model_grad), so setting it produced
activations and zero parameter gradients — no error, no warning, empty dump.

The training path reaches dump_model through DumperMegatronUtil.finalize
(miles/utils/dumper_utils.py); the standalone worker had no equivalent.

get_grad is required alongside, not optional. On sglang-miles,
dumper.py:509 reads:

g = get_grad(value) if get_grad is not None else value.grad

a bare .grad with no main_grad fallback. Under Megatron DDP the gradient
lives in param.main_grad with param.grad left None, so calling
dump_model without a getter still yields a silently empty dump. The heavier
distributed-optimizer shard gather in dumper_utils._build_full_grad_getter is
only needed for the FT trainer, whose reduce-scatter leaves each rank holding a
partial buffer — the standalone worker has no reduce-scatter, so the cheap
.grad or main_grad read is sufficient.

step= is deliberately not passed: dump_model on sglang-miles has no step
parameter, so it lands in **kwargs → extra_kwargs → tags, and
_dump_single then does dict(step=..., rank=..., **tags) → TypeError: dict() got multiple values for keyword argument 'step'.

2. loss_func normalised by the local valid-token count under CP

reduction="mean" divides by each rank's own valid count, but under context
parallelism the labels are sharded. Every activation gradient comes out scaled
by global/local — about 2x at CP=2 — so a CP run compared against CP=1 shows a
spurious ~0.5 relative difference that swamps any real discrepancy.

Summing locally and dividing by the CP-reduced count 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.

3. Forward-only runs still allocated DDP grad buffers

get_model was always called with the default wrap_with_ddp=True, so a
forward-only run paid for per-parameter grad buffers it never used, OOMing on
large models when only logits/logprobs were wanted. Now keyed off
script.run_backward, with m.train(script.run_backward) to match.

Also mirrors the miles backend's variable_seq_lengths, which Megatron's
validate_args unconditionally resets to False.

Adjacent issue, not fixed here

By the same step= mechanism, miles/utils/dumper_utils.py calling
dumper.dump_model(extracted_model, get_grad=get_grad, step=0) would raise
TypeError against sglang-miles. Left alone since it is a different code
path; flagging it because anyone enabling DUMPER_ENABLE_MODEL_GRAD on the
training path will hit it. Happy to fix in a follow-up.

Note also that dumper_utils.finalize only builds a get_grad when
enable_experimental_ft_trainer() is true and passes None otherwise — which,
per the bare-.grad default above, means non-FT Megatron DDP runs dump empty
gradients there too. _log_model_grad_coverage uses the correct
.grad or main_grad read, so the two disagree.

Testing

pre-commit and py_compile pass. I could not run the fast tests locally — no
sglang / megatron on this host — so CI is the first real exercise.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.