[arnavkomaragiri]

Superseded

Retired in favor of NRL-owned Hybrid MTP control in NVIDIA-NeMo/RL#4147: freeze retained MTP weights before DDP/optimizer setup and pass the existing upstream compute_mtp_loss=False argument through policy forwards. Unsupported MTP-bearing models fail explicitly; no new MCore config field or checkpoint topology change is needed.

The replacement passed FP32/BF16 Hybrid logits/base-gradient parity, checkpoint-key and bypass checks, plus the async 1x16 TQ/R3/CP2 gate. No active PR in this stack depends on this proposal. #7387 is the sole remaining MCore change.

Historical Proposal

Add opt-in TransformerConfig.disable_mtp_loss for fine-tuning an MTP checkpoint without the auxiliary MTP training forward or vocabulary loss.

• Retain MTP modules and checkpoint keys; this is not a change to mtp_num_layers.
• Freeze retained MTP parameters so they do not acquire optimizer state.
• Skip the auxiliary training forward and loss in GPT and Hybrid models, while retaining inference paths.
• Preserve Hybrid's existing per-forward compute_mtp_loss control.
• Reject the contradictory combination with freeze_base_model_for_mtp.

Setting an MTP loss coefficient to zero alone still computes the auxiliary vocabulary projection. This option avoids that compute and peak allocation without removing model structure.

Tests

Added native model tests for strict checkpoint loading with unchanged keys, frozen MTP/trainable base parameters, skipped MTP forward/loss, and a functioning base backward pass for GPT and Hybrid. These new model tests require CUDA/Transformer Engine and have not been executed on the login host; CI execution is required before readiness.

The extracted production behavior was exercised in the successful integrated async 1x16 Harbor training gate, with finite nonzero gradient norm. Tested MCore base: 54c62dfb6852e9218bf38590be9b90930bf7d589. New upstream changes have not yet been reconciled and revalidated.

Stack

No longer a prerequisite. The historical patch is retained here for reference; its implementation is not carried by the active shared-prefix PR.

Review Status

Superseded; closed without merging.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[arnavkomaragiri]

Superseded by the NRL-owned Hybrid MTP control in NVIDIA-NeMo/RL#4147, using the existing compute_mtp_loss forward argument and pre-wrap parameter freezing. The replacement passed FP32/BF16 logits/base-gradient parity, checkpoint-key/bypass checks and the async 1x16 gate. Closing without merging; #7387 no longer depends on this proposal.