[omri-manolo]

Implements RFC #4108: a new opt-in --log-probs-backend fused_lm_head. The actor's model returns the output layer's input instead of its logits, and the loss runs the output layer on the rows it scores, a tile at a time under activation checkpointing, reducing each tile with #3741's fused_log_probs_and_entropy. The micro-batch's [tokens × vocab/TP] logits and their gradient are never materialized. torch and fused are unchanged.

How it works

• Hook: lm_head.py passes Megatron's GPTModel.forward(output_processor=...). The processor returns the decoder output as [b, s, h] (gathered over TP when sequence parallel is on) and records the output layer for the loss (batch["lm_head"]). Only the actor's forwards use it; the critic keeps its value head.
• Op: fused_lm_head_log_probs_and_entropy (fused_log_probs.py). Per tile of scored rows, under torch.utils.checkpoint: the output layer with sequence parallel off, then fused_log_probs_and_entropy with inplace_backward. Backward replays the tile, so the output layer computes its own gradients and the hidden gradient is all-reduced over TP by the layer's backward. Tile size is --log-probs-chunk-size rows (default 4096). Extra memory is about one [tile, vocab/TP] tile, independent of the token budget.
• Selection: get_log_probs_and_entropy selects response rows and targets exactly as fused does, so prompt rows are never multiplied by the output weight.

Scope

• Covers --loss-type policy_loss and sft_loss, with entropy (on-policy distillation goes through the policy loss).
• Rejected at startup with a clear message: true-on-policy mode, sampling-support replay, a missing --vocab-size, a train backend other than Megatron, other loss types, multi-LoRA. Models whose forward has no output_processor (e.g. HybridModel) are rejected at model build.
• With --enable-mtp-training a startup warning says the MTP heads are not covered: Megatron computes the MTP loss in GPTModel._postprocess before the output_processor runs, so each MTP layer still materializes its fp32 logits. A fused MTP loss needs a Megatron hook; follow-up.
• Score centering's K-candidate scoring through this backend is ready on top of Fused log-prob op: K targets per row; score centering trains on it #3998 (fused op with K targets per row); it will follow as a small PR once Fused log-prob op: K targets per row; score centering trains on it #3998 lands.

Testing

• Unit/GPU tests (tests/fast-gpu/test_fused_lm_head_log_probs.py, tests/fast/backends/megatron_utils/test_lm_head.py, argument validation): equality against fused on hidden @ W.T for log-probs, entropy and the hidden/weight gradients; uneven tiles, no scored rows, frozen weight, two TP shards over NCCL; policy and SFT losses in thd and bshd. With the existing fused log-prob tests: 624 passed on 1× H100.
• End to end on this branch: Qwen3.6-27B (TP4, 2×8 H100), the same recorded rollouts replayed into both backends (--load-debug-rollout-data --debug-train-only):

tokens/GPU	backend	step-0 trainer-vs-rollout |Δlogp|	peak allocated (GiB)	result
114,688	fused	—	—	OOM ("Tried to allocate 13.03 GiB")
114,688	fused_lm_head	0.005929194	64.75	3 steps
49,152	fused	0.005926252	52.97	3 steps
49,152	fused_lm_head	0.005926252 (identical)	49.44	3 steps

Steps 1–2 differ within the run-to-run variation we see within one backend.

• Also run with the same two commits on top of Fused log-prob op: K targets per row; score centering trains on it #3998: Qwen3.6-35B-A3B (TP2 × EP8) with and without MTP training (step-0 log-probs and MTP loss identical across backends; 9.8 GiB less at 36,864 tokens/GPU without MTP), and GLM-5.3 744B LoRA with PP8 × CP8 up to 128k context.

Not in this PR

• Score centering through fused_lm_head (after Fused log-prob op: K targets per row; score centering trains on it #3998).
• A fused MTP loss (needs a hook in Megatron's process_mtp_loss).
• Tile size derived from free memory (fixed --log-probs-chunk-size for now).

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[HJSang]

Design question: everything LMHeadContext records (output_layer, output_weight, scale_logits) is constant for the model's lifetime, but lm_head_forward_kwargs builds a new context per micro-batch in both forward steps and the processor fills it at call time. Could the context be built once at model build (next to check_model_takes_output_processor) and bound to the forward, with the loss receiving that run-level LMHead once instead of batch["lm_head"] plus the forward_only callback kwarg? That would make hidden_states_output_processor stateless (gather + transpose only) and leave the forward steps and the batch dict as they were. Nice that this builds on the upstream output_processor hook (NVIDIA/Megatron-LM#4686) rather than a patch.

[HJSang]

Could the testing table also carry latency? Per-step train time (and the log-probs forward time) for fused vs fused_lm_head at the same token budget, so the memory saving is read next to the extra GEMM per tile (the RFC bench shows 277 ms vs 330 ms at 49K tokens; the end-to-end number at 49,152 and 114,688 tokens/GPU would complete the picture).

[HJSang]

With a LoRA target that resolves to the output layer (unembed, the default for Inkling models in hf_lora_targets.py), Bridge wraps output_layer in LoRALinear, which has no sequence_parallel or tp_group attribute, so this raises AttributeError on the first micro-batch after build and load (project would also pass weight= and runtime_gather_output= into the adapter). Could _validate_fused_lm_head_args reject that case at startup, next to the multi-LoRA check, until the wrapped layer is supported?

[omri-manolo]

great catch! adding this check to _validate_fused_lm_head_args

[HJSang]

This fires per micro-batch, after the model is built and the checkpoint loaded. Megatron selects the schedule plan from overlap_moe_expert_parallel_comm, so could the rejection move to _validate_fused_lm_head_args with the other startup checks, keeping this assert as a guard?

[omri-manolo]

agreed, will add this as a startup check as well.

[HJSang]

The 4096 default lives here, in fused_lm_head_log_probs_and_entropy (tile_rows: int = 4096) and in the --log-probs-backend help text. Could one of them own it, e.g. make tile_rows a required keyword and let this constant be the single source?

[omri-manolo]

Design question: everything LMHeadContext records (output_layer, output_weight, scale_logits) is constant for the model's lifetime, but lm_head_forward_kwargs builds a new context per micro-batch in both forward steps and the processor fills it at call time. Could the context be built once at model build (next to check_model_takes_output_processor) and bound to the forward, with the loss receiving that run-level LMHead once instead of batch["lm_head"] plus the forward_only callback kwarg? That would make hidden_states_output_processor stateless (gather + transpose only) and leave the forward steps and the batch dict as they were. Nice that this builds on the upstream output_processor hook (NVIDIA/Megatron-LM#4686) rather than a patch.

Agreed, and thanks for the pointer. Two points before I rework it:

1. Building at model build. output_layer and shared_embedding_or_output_weight() are public, but MuP's logit scaling is only exposed through the hook's scale_logits argument (otherwise it's the private _scale_logits). Options: reject use_mup at startup with this backend, or keep taking the scaling from the hook, which leaves the processor slightly stateful. WDYT?

2. Handing the run-level LMHead to the loss. The losses only see (args, batch, logits), never the model, so whatever owns the LMHead has to pass it down somehow. Do you have an idea or preference here?