perf(sft): add a direct path for Megatron-LM prepacked inputs by seonjinn · Pull Request #3380 · NVIDIA-NeMo/RL

seonjinn · 2026-07-28T06:28:35Z

What does this PR do ?

This PR adds an opt-in SFT input path that consumes Megatron-LM-style
offline-packed JSONL rows without repacking their constituent conversations
inside NeMo-RL.

The new path:

registers the megatron_sft_packed response dataset;
reconstructs shifted input/target tensors, loss masks, reset position IDs,
and packed sequence boundaries from each offline-packed row;
bypasses NeMo-RL online sequence packing only when direct-packed fields are
present;
preserves context-parallel and MTP target/mask alignment;
supports validation rows with deterministic padding before data-parallel
sharding; and
defers scalar metric materialization to the reporting boundary and avoids the
PP1 object-broadcast path.

Normal response datasets continue through the existing collator and packing
path. The direct path is SFT-only and opt-in.

Issues

No linked issue.

Usage

sft:
  only_unmask_final: false

policy:
  train_micro_batch_size: 1
  dynamic_batching:
    enabled: false
  sequence_packing:
    enabled: false
  megatron_cfg:
    enabled: true
    use_fused_linear_logprobs: false

data:
  train:
    dataset_name: megatron_sft_packed
    data_path: /path/to/train.jsonl.packed
    chat_key: messages
  validation:
    dataset_name: megatron_sft_packed
    data_path: /path/to/validation.jsonl.packed
    chat_key: messages
  default:
    megatron_sft_prompt_format: identity
    megatron_sft_pad_token: "<unk>"
    megatron_sft_assistant_prefix_len: 0
    megatron_sft_context_parallel_size: ${policy.megatron_cfg.context_parallel_size}

The direct path currently requires the Megatron policy backend, micro-batch
size 1, dynamic batching disabled, sft.only_unmask_final=false, the standard
NLLLossFn, fused linear logprobs disabled, delegate_pack_to_model=false,
and no draft-model, split-training, or router-replay input. Unsupported
combinations fail closed during setup or before model execution.

Performance

A strict 64-node comparison used 512 H100 GPUs, the same static packed row
schedule, and matched performance-sensitive settings: TP8, CP16, EP32, ETP1,
PP1, GBS64, MBS1, sequence length 262144, selective MoE recompute, overlapping
gradient reduction and parameter gathering, two repeated MTP layers, P2P
batching enabled, and NVLS disabled.

System	Mean step time	Window	Relative throughput
Megatron-LM	51.9744s	Steps 3-20	100.00%
NeMo-RL	52.3428s	Steps 3-20	99.30%
NeMo-RL gap	+0.3684s (+0.71%)	18 steps	-0.70%

In this strict comparison, NeMo-RL was within 0.71% of standalone Megatron-LM.

Detailed performance and convergence results are available in the 200-step W&B report. The lm_loss and validation loss closely matched between the two systems.

Correctness

For steps 3-20, training loss closely matched between the two systems, with a
mean absolute difference of 0.001631 and a maximum difference of 0.004211.
All compared losses were finite.

Before your PR is "Ready for review"

Pre checks:

Read and followed the contributor guidelines.
Added focused tests for dataset parsing, collation, CP/MTP handling,
validation DP remainders, pipeline transport, loss handling, metric payloads,
and worker integration.
Run focused Linux unit tests on the latest PR head (68 passed).
Run the full upstream unit and functional test suites; full PR CI is
pending.
Add user-facing documentation after the dataset/config interface is
accepted in review.

Additional Information

The input must contain Megatron-LM-style packed conversation rows in a
supported SFT prompt format.
The adapter retokenizes each packed row with the configured tokenizer; it
does not deserialize prebuilt PyTorch tensors.
The direct model-loss path intentionally rejects custom loss functions and
fused linear logprobs.
This branch contains current upstream main commit
80555d3a0595ce3cf76f2ca1b2bf123339064556; full CI and the matched
latest-main 200-step run are pending.

Reviewer focus:

Megatron-LM token/target masking and packed-boundary equivalence;
CP-local label and loss-mask alignment, including MTP;
fail-fast guards for unsupported configurations;
pipeline-parallel metric payload materialization; and
absence of behavior changes for non-packed datasets.