Optimize end-to-end Pipeline RL throughput

[FurtherAI]

Summary

This PR improves ART Pipeline RL throughput across rollout scheduling, policy freshness, trajectory packing, MoE communication, Megatron execution, and observability.

It adds online runtime tuning, in-flight LoRA updates with token-level policy provenance, generalized prefix-tree packing, binary MoE route transport, and a substantially faster Megatron critical path.

During development, we optimized around a compact throughput score:

score = accepted_train_tok_per_s * freshness_discount * batch_factor

ART reports the score and its components to W&B. It is intended as an interpretable systems benchmark that rewards useful throughput without ignoring policy staleness or scenario efficiency. It is not a training objective or a substitute for measuring real learning progress.

Benchmark Score

The score was inspired by PipelineRL's score framing and is recorded in wandb. It combines three terms:

• Accepted train throughput measures trainable assistant tokens accepted by the pipeline per wall-clock second. Work discarded as stale or zero-variance does not count.
• Freshness discount reduces the value of tokens generated by older policies. It uses token-level policy provenance and penalizes old tails more strongly than simply using mean policy age.
• Batch factor accounts for the diminishing statistical benefit of increasing optimizer batch size relative to the additional scenarios consumed. Increasing rollouts for the same scenario can improve the factor because it strengthens the grouped learning signal without consuming more scenarios.

The batch factor follows critical-batch-size models in which increasing batch size reduces steps to target with diminishing returns. Its current critical rollout scale of 300 is grounded in early GRPO/RLVR measurements, but remains a heuristic rather than a universal constant.

The freshness term is motivated by importance-sampling effective sample size: the second moment of behavior-to-learner importance ratios determines effective sample size and is related to order-2 Rényi divergence. ART uses policy age with a decay scale of eight training steps as a practical divergence proxy. This is likewise a grounded benchmark approximation, not an empirical learning law.

Adaptive Pipeline Scheduling

• Add an optional online autotuner for rollout workers, batch bounds, and queue capacity.

This is really cool, this removes the need to configure batch size, queue size, number of rollout workers, etc. from the user. These are things you'd have to tune per run and understanding the impact on efficiency is difficult. The autotuner performs really well and also can adapt to changing workloads over time.

• Infer packed-sequence capacity, data-parallel requirements, and inference scale from the active runtime.
• Estimate counterfactual packing outcomes from real trajectory shapes and target a configurable spill probability.
• Bound queued work using the configured policy-age limit.
• Use strict metric contracts: missing or excessively delayed serving metrics fail instead of silently degrading.
• Persist tuner profiles at each decision window and support resuming them.
• Record pipeline settings and tuner decisions as canonical metrics.

In-Flight Policy Updates

• Add an ART vLLM endpoint for atomically updating a stable LoRA slot in place.
• Allow active requests to use updated weights on subsequent decode tokens.
• Return compact token spans describing which policy generated each completion region.
• Use these spans for token-weighted policy age, filtering, metrics, and scoring.
• Salt new-request prefix-cache entries by policy version while preserving the KV state of already-running requests.
• Preserve versioned public model aliases and exact adapter leases for evaluation and checkpoint use.

Prefix-Tree Packing and MoE Routes

• Generalize the prefix-tree packing algorithm past greedy to a FFD-bin-packing inspired form with a cleaner design, though both perform similarly in packing efficiency. Prefix-tree packing happens across arbitrary common regions of trajectory groups in the batch, so a close to optimal packed representation of a batch is expected. Pretty cool.
• Replace JSON/base64 route metadata with a compact binary route bundle.
• Remove redundant route-conflict computation and validation from the production critical path.
• Expose the same packing model to the autotuner when estimating how many groups fit.

The complete path from trajectory groups to packed tensors, including MoE routing replay data, typically takes approximately 0.5-1 second on single-turn-shaped batches. This removes a major preprocessing bottleneck present when routing replay is enabled (current version is very slow, wasn't measured properly and is poorly optimized).

Megatron Throughput

• Keep compatible optimizer and LoRA state resident across training jobs.
• Save optimizer state every five steps by default and unconditionally on the final step.
• Signal adapter readiness as soon as LoRA saving completes so serving reload can overlap worker cleanup.
• Require the CUTLASS grouped-GEMM path, avoiding the routed-shape cache churn encountered in the cuBLAS implementation.
• Replace DeepEP with vendored HybridEP for MoE dispatch and combine.
• Cache HybridEP builds by source, environment, and GPU architecture, with hard architecture validation.
• Add configurable streaming weight offload to the Megatron runtime configuration.

Compile specialization and kernel JIT behavior were significant barriers to reaching isolated Megatron throughput before this PR. Changing sequence and routing shapes could repeatedly trigger TorchDynamo, Triton, Transformer Engine, or model-specific specialization work rather than settling into steady execution.

This PR removes the recurring specialization paths found across every currently supported handler while retaining compilation around the operations that benefit from it. It includes fixes for GDN mutable state, dynamic normalization, sparse attention dimensions, MoE dispatch, grouped linear operations, and model-specific postprocessing.

HybridEP was finite and correct across 90 measured model, EP, and token-count combinations. Forward-plus-backward geometric speedups over DeepEP ranged from 1.21x to 1.38x by model, with individual points between 1.10x and 1.56x.

Metrics and Model Support

• Add a canonical metric taxonomy and builder shared by PipelineTrainer components.
• Add a low-overhead ART vLLM metrics endpoint.
• Report Megatron packed throughput, padding, accepted throughput, policy age, discard rates, pipeline settings, and the score components.
• Add distributed importance-ratio and clipping diagnostics (these are likely better to watch than number of steps off policy to determine how much you want to accept and how it affects your run).
• Improve support for Qwen3 MoE, Qwen3.5 MoE, Gemma4 MoE, GPT-OSS, and DSV4.
• Add internal expert-dimension padding where required by CUTLASS while preserving logical weights, gradients, and exported adapters.
• Standardize DSV4 LoRA tensors on the fused 3D representation, including down projection.

Results

Online autotuning improved the last-10-step score over the same fixed reasonable defaults on all six synthetic workloads:

Workload	Fixed	Autotuned	Improvement
Single-turn tool/history	271.2	326.7	1.20x
Long generation	1285.5	2126.8	1.65x
Shared-prefix packing	581.9	594.2	1.02x
Multi-turn agent	789.5	901.7	1.14x
Short single-turn	506.7	813.5	1.61x
Weight-update-bound	51.0	118.5	2.32x

Final 20-step single-turn-shaped profiles produced these last-10-step means:

Model	Packed train tok/s	Accepted tok/s	Score
Qwen3 30B-A3B	18,610	380.8	281.0
Qwen3.5 35B-A3B	19,472	435.1	319.4
Gemma4 26B-A4B	8,210	306.8	222.7
GPT-OSS 20B	18,629	527.7	390.8
DSV4 Flash	3,212	92.6	72.7

The Qwen, Gemma, and GPT-OSS profiles used two training and two inference H200s with 122,880-token packed sequences. DSV4 used an eight-GPU TP2/EP8 trainer and a separate four-GPU EP4 inference deployment.

Evaluation was disabled in these profiles to isolate Pipeline RL throughput.

Validation

• Full formatting, lint, type, and lockfile checks.
• ART wheel and source-distribution builds, including bundled runtime and HybridEP source checks.
• HybridEP BF16 and FP8 dispatch/combine validation on H200s.
• Packing invariance and length-trainability coverage across the supported MoE handlers.
• Targeted oracle and train/inference parity checks for changed model paths.
• Twenty-step end-to-end profiles for all five model families.
• Production profiles accounted for effectively all PipelineTrainer wall time.

Scope

The autotuner intentionally requires ART vLLM capabilities, including fast metrics and in-flight LoRA updates. Required capabilities do not silently fall back to weaker behavior.

This PR bundles HybridEP source and supports deterministic local Megatron setup. Broader release-environment automation and prebuilt distribution are deferred to a follow-up.

[bradhilton]

Requesting changes before merge. Performance gains look real, and the 2xH200 TrainerRank gate passed, but correctness, recovery, and compatibility blockers remain.

1. Cover dynamic LoRA slots in internal-padding cleanup. Gemma4 and GPT-OSS cleanup currently touches only legacy A_T/B_T parameters. Dynamic TrainerRank checkpoint slots retain padded gradients and parameters. Add one stable LoRA-site iterator in art.megatron.lora and use the iterator from both handlers.

2. Make optimizer persistence upgrade-safe, crash-atomic, and objective-safe. Existing fixed shards without a numeric marker must remain resumable. Replace in-place shard overwrites with versioned temporary shards plus an atomically committed manifest. Key resident optimizer reuse by training session, objective/mode, and optimizer path so RL Adam moments cannot leak into SFT. Save a final committed generation during clean stop/finalization even when final_training_step is unknown. A root-level versioned shard format remains compatible with both ../../coreweave/serverless-training prod and staging: both branches scan root .pt files containing -of-, validate list payloads, preserve root training_mode, and recursively copy the optimizer directory.

3. Map external vLLM in-flight adapter paths. Normal reload and exact-eval paths translate trainer-local paths to server-visible paths; the in-flight endpoint currently bypasses the mapping.

4. Preserve PipelineTrainer constructor and default behavior. Keep deprecated direct pipeline kwargs as aliases into PipelineRuntimeConfig. Preserve the previous unset max_batch_size = 10 * min_batch_size; the new fallback changes the default from 40 to 4 while autotuning remains off.

5. Preserve metric API compatibility and settle one migration rule. The migration is not a uniform {split}/{metric} to {metric}/{split} inversion: built-in reward becomes reward/val, trajectory metrics become task/val/<metric>, group metrics become task/val/group/<metric>, direct custom metrics retain split prefixes locally, and the W&B filter silently drops non-canonical custom metrics. Infrastructure keys and cumulative-state keys also change. Define one naming contract, preserve arbitrary user W&B metrics, dual-read or dual-write compatibility aliases during migration, and accept historical val/reward for checkpoint retention so prior best checkpoints remain eligible.

6. Restore base-wheel importability. A clean base install now fails first on NumPy and then Torch because top-level art imports reach optional pipeline/preprocessing modules. Keep optional exports lazy or add required base dependencies; the backend-extra smoke test does not cover the base wheel contract.

7. Propagate fatal rollout-worker exceptions before reconciliation. Completed tasks are removed before exception inspection, allowing fatal local-serving failures to disappear and workers to be replaced indefinitely.

8. Drain finite scenario sources without losing in-flight tail work. Source exhaustion should stop new fetches, allow active rollouts to enqueue, and only then end the stage. Current global stop can discard up to num_workers - 1 scenarios after persisted offsets have advanced.

9. Keep cross-backend training kwargs executable. Backend-specific arguments are expected, but unconditional optimizer_save_interval and final_training_step currently make PipelineTrainer + ServerlessBackend fail at runtime. Add optional accepted kwargs or conditionally emit supported kwargs.

10. Preserve existing completion-token counts when exact reconstruction is unavailable. Clearing trajectory.metrics["completion_tokens"] before reconstruction can remove a valid caller-provided count; policy-age accounting later requires a positive value.

11. Optional follow-up: define streaming policy provenance behavior. Streaming responses bypass the full-response metadata wrapper. Either attach complete token policy spans to streaming output or reject streaming while in-flight policy tracking is required.

Add targeted regression tests for every corrected behavior, including legacy optimizer migration, interrupted shard writes, RL-to-SFT resident transitions, external path mapping, base-only wheel import, worker exception propagation, finite-source draining, ServerlessBackend execution, metric aliases/W&B custom metrics, and dynamic-slot padding.

[bradhilton]

Dynamic checkpoint slots live under _slot_modules, so direct module.A_T/B_T access clears only legacy parameters. A targeted H200 run left a padded GPT-OSS slot grad at 0.0206298828125 and moved the padded parameter by 0.00099945068359375. Add a stable iterator in art.megatron.lora over legacy and dynamic slot matrices, then use the iterator in both Gemma4 and GPT-OSS padding cleanup.

[bradhilton]

Persistent optimizer presence does not prove optimizer state belongs to the incoming SFT adapter. An RL-resident optimizer can reach this return path and carry Adam moments across objectives. Key resident state by training session + objective/mode + optimizer path, and rebuild/load whenever the complete key changes.

[bradhilton]

In-place overwrite can leave mixed or corrupt shards behind an old marker after crash or ENOSPC. Write root-level versioned *.pt.tmp shards, os.replace each shard, barrier, then atomically replace CURRENT.json listing step/world/files; retain the previous committed generation until manifest advance. Keep root training_mode and legacy 01-of-04.pt reads. The format remains compatible with serverless-training prod/staging root shard validation and recursive checkpoint copying.

[bradhilton]

External vLLM deployments can use a different server-visible checkpoint root. Match the reload and exact-eval paths here: "lora_path": self._vllm_checkpoint_path(checkpoint_path), plus an external-runtime mapping test.

[bradhilton]

Preserve constructor compatibility by accepting the five former direct pipeline kwargs as deprecated aliases into PipelineRuntimeConfig. Also preserve the unset maximum as 10 * min_batch_size; the current fallback reduces the default from 40 to 4 while autotuning is off.

[bradhilton]

_reconcile() removes every completed task before _raise_finished_errors() can inspect exceptions. Reverse the calls (await _raise_finished_errors() before _reconcile()) and add a fatal-worker regression test.

[bradhilton]

Global stop on first exhaustion cancels other active workers after scenario offsets have already advanced. Record source exhaustion separately, stop new fetches, drain active rollouts into the output queue, then publish the terminal sentinel after all workers finish.

[bradhilton]

These kwargs are emitted unconditionally, while ServerlessBackend.train() accepts neither. Add optional ignored kwargs to ServerlessBackend or emit backend-supported kwargs conditionally; the permissive Backend.train callable type hides rather than resolves the runtime incompatibility.

[bradhilton]

Do not delete a valid caller-provided completion_tokens value before proving exact reconstruction is available. Overwrite only when every choice has an exact count; otherwise retain the existing positive value. Pipeline policy-age accounting later raises when the metric is missing.

[bradhilton]

Canonical filtering silently drops arbitrary user metrics from W&B even though the same metrics remain in local history. Preserve user-provided metrics in the W&B payload; canonical definitions can control dashboards and step axes without becoming an allowlist.

[bradhilton]

Four additional findings from a follow-up review:

1. Support the declared minimum Weave version. _suppress_weave_trace() imports override_settings, but the permitted weave==0.52.24 lacks that symbol. Every proxied completion fails with ImportError before HTTP dispatch. Raise the dependency minimum or provide a compatible fallback, with a lower-bound dependency test.

2. Make binary MoE metadata serializable. attach_moe_routing_metadata_to_choice() stores a NumPy array in Choice.model_extra. Routed trajectory groups then fail during serverless logging/training with PydanticSerializationError from model_dump(mode="json"). Add a JSON-safe representation or explicitly exclude/reconstruct the metadata at serialization boundaries.

3. Prevent stale checkpoint aliases from serving newer weights. register_lora_alias() and _resolve_lora_alias() map every model@step name to the same mutable slot. After a later update, an older explicit name silently serves current weights despite get_inference_name(step) promising a specific checkpoint. Track the expected version per alias and reject stale aliases, or preserve exact checkpoint semantics.

4. Limit Weave suppression to ART-managed inference. create() suppresses tracing for every proxied completion, including ordinary OpenAI and OpenRouter calls because the proxy is installed universally. Apply suppression only to local or binary ART vLLM responses and retain normal tracing for remote providers.

The four items above are requested changes. Everything below is explicitly optional, low priority, and not a merge blocker:

• New sensitive /art/* endpoints inherit the existing vLLM authentication-prefix gap. Binary inference, cache reset, and in-flight LoRA updates sit outside vLLM's guarded prefixes. Route-level authentication would provide useful defense in depth.
• Online autotuning has no absolute rollout-worker cap. Repeated underfed windows can continually increase the worker target. A configurable upper bound would prevent long-running resource growth.
• Cancellation while begin_update() waits for active admissions can leave the slot blocked. Cancellation cleanup would avoid requiring a runtime restart.
• Autotuner warmup uses absolute training steps, so a resumed run can skip restart warmup. Applying warmup relative to the resumed starting step would avoid reacting to restart transients.
• NumPy assignments in _materialize_prefix_tree_row() can broadcast malformed one-element token metadata instead of rejecting inconsistent lengths. Explicit length validation would improve defensive error handling.

[vivekkalyan]

when running vLLM with n=8 rollouts, policy provenance is lost for all but the final vLLM child under FINAL_ONLY parallel sampling. RequestState.make_request_output() returns None while early children are buffered, so the current wrapper never attaches their spans; ART then rejects those choices because their completion tokens have zero provenance coverage. This is a small fix to records each child’s cumulative spans on the shared ParentRequest before aggregation, then attaches the complete per-choice map when vLLM returns the final parent response.