Support PP-free Megatron forward/backward warmup

[nvcastet]

Add PP-free Megatron forward/backward warmup

Summary

Add an opt-in --enable-pp-free-warmup option for Megatron actors.

Before the first actor step, every pipeline stage performs one local forward and backward pass concurrently. This initializes JIT kernels without pipeline dependencies, then initializes the bidirectional NCCL pipeline channels before the first real pipeline execution.

The warmup pass shaves ~20min on 1st step of DSV4+GB300 RL training (~4x speedup).

The warmup:

• Performs no optimizer step.
• Disables gradient synchronization.
• Supports virtual pipeline chunks.
• Restores RNG, FP8, registered buffers, gradients, iterator position, model state, and temporary trackers.
• Runs once and is skipped when pipeline parallelism is disabled.

Motivation

During the first normal pipeline execution, a stage cannot compile its kernels until it receives activations from the preceding stage. Forward compilation therefore propagates from the first stage to the last. Backward compilation then propagates in reverse.

The first use of each pipeline communication path also triggers lazy NCCL setup.

Before: serialized pipeline warmup

flowchart LR
    R0F["Rank 0<br/>JIT compile FWD kernels"]
    C0F["NCCL PP<br/>comm setup"]
    MF["..."]
    RNF["Rank N<br/>JIT compile FWD kernels"]
    RNB["Rank N<br/>JIT compile BWD kernels"]
    CNB["NCCL PP<br/>comm setup"]
    MB["..."]
    R0B["Rank 0<br/>JIT compile BWD kernels"]

    R0F --> C0F --> MF --> RNF --> RNB
    RNB --> CNB --> MB --> R0B

Loading

After: PP-free warmup on all ranks concurrently

flowchart TB
    subgraph S0["Rank 0"]
        direction LR
        R0F["JIT compile<br/>FWD kernels"] --> R0B["JIT compile<br/>BWD kernels"] --> R0C["PP comm setup"]
    end

    DOTS["..."]

    subgraph SN["Rank N"]
        direction LR
        RNF["JIT compile<br/>FWD kernels"] --> RNB["JIT compile<br/>BWD kernels"] --> RNC["PP comm setup"]
    end

    S0 ~~~ DOTS
    DOTS ~~~ SN

Loading

Performance

Measured on 16 GB300 nodes:

• 8 actor nodes and 8 rollout nodes, 4 GPUs per node
• TP=2, PP=8, CP=2, EP=4
• The no-warmup control used --distributed-timeout-minutes 60 so the serialized first pass could complete

First-step phase	Without warmup	With warmup	Time saved
Local F/B warmup	—	3m 14.6s	−3m 14.6s
Log-probs	13m 28.9s	48.0s	12m 40.9s
Actor train	14m 06.9s	2m 55.5s	11m 11.4s
Warmup + first training	27m 35.9s	6m 58.2s	20m 37.6s

Validation

Activations and gradients were dumped with and without warmup flag to validate that warmup did not impact training. The raw reward and KL were also compared.

[gemini-code-assist]

Caution

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

[nvcastet]

CC @kaixih @yueming-yuan