Ultra 3 vllm tuning

[fsiino-nvidia]

(branch originally cut from fsiino/super35-vllm-model-tuning. After #2599 merged, this branch was rebuilt on top of the updated main. The Ultra commits were replayed and the overlapping launcher changes were reconciled against the merged implementation.)

What this PR does

• Adds benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh, an Ultra-specific vLLM configuration using eager prefill, piecewise CUDA-graph decode, MTP3, FP8 KV cache, expert parallelism, and the memory and scheduler limits validated during tuning.
• Extends benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh with a coupled prefill/decode deployment. Each tier forms one distributed engine across four nodes because the Ultra weights cannot fit on one four-GPU node; the existing independent deployment remains available for smaller models.
• Makes evaluation concurrency, walltime, and Slurm segmentation configurable. Adds RESUME_EVAL_ON_REQUEUE primarily for automatic recovery after Slurm preemptions: it keeps the results path stable across restarts and enables Gym's existing resume-from-cache behavior.
• Detects prefill process failures while waiting for either prefill or decode readiness, so failed coupled launches exit with an accurate error instead of continuing to wait.
• Removes experiment-only tuning switches after testing and retains the known-good settings as focused defaults.
• The completed 11,213-row validation averaged 105.01 tokens/sec/request, with nonzero waiting in 2.04% of telemetry windows and 1.78% average decode KV-cache usage. MTP3 was the best validated throughput configuration, although rare repeated-output failures remain documented in the tuning handoff.

Benchmark scores

Completed validation run: Slurm job 7014044, 11,213 evaluation samples. These are measured results; accuracy parity with the expected baseline is not yet established, particularly for Banking and AA-LCR.

Benchmark	Score	vs. tracker
Tau3-Banking	11.13%	18.40%
Tau3-Average	58.86%	64.60%
SciCode — full-problem accuracy	11.54%	11.90%
AA-LCR	52.13%	64.10%
AA-Omniscience — answer accuracy	24.38%	23.80%
GPQA Diamond	85.98%	86.50%

Source artifact: results/nemotron-3-ultra-bf16-p4d4-c16-mtp3-v0271-piecewise-copyinputs-full-suite-validation-20260908/resumable_aggregate_metrics.json.

Run command

Replace the placeholder values before submitting and use a new, unused experiment name for each new suite. Set RESUME_EVAL_ON_REQUEUE=1 on the first submission so Gym can reuse completed samples if Slurm preempts and requeues the job; it starts fresh when no saved results exist.

MODEL=/lustre/fsw/portfolios/llmservice/users/igitman/hf_models/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 \
VLLM_CONFIG=benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh \
EXPORT_TO_CSV=0 \
EXPERIMENT_NAME="<UNIQUE_EXPERIMENT_NAME>" \
RESUME_EVAL_ON_REQUEUE=1 \
NUM_SAMPLES_IN_PARALLEL=16 \
NUM_PREFILL_NODES=4 \
NUM_DECODE_NODES=4 \
VLLM_PD_DEPLOYMENT_MODE=coupled \
VLLM_SLURM_SEGMENT=4 \
ULTRA_ENABLE_MTP=1 \
ULTRA_NUM_SPECULATIVE_TOKENS=3 \
ULTRA_DECODE_ENFORCE_EAGER=0 \
ULTRA_DECODE_CUDAGRAPH_MODE=PIECEWISE \
ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS=1 \
SBATCH_ACCOUNT="<SBATCH_ACCOUNT>" \
SBATCH_PARTITION=batch_long \
SBATCH_QOS=normal \
SBATCH_TIME=20:00:00 \
SBATCH_GRES=gpu:4 \
CONTAINER="<ABSOLUTE_PATH_TO_EVAL_CONTAINER.sqsh>" \
MOUNTS="/lustre:/lustre,$(pwd)/env.yaml:/opt/Gym/env.yaml:x-create=file,$(pwd)/results:/opt/Gym/results" \
bash benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh \
  --config benchmarks/nemotron_3.5_super/eval_container_config.yaml \
  '++critpt_benchmark_agent.responses_api_agents.critpt_agent.datasets=[]' \
  '++hle_benchmark_equivalence_llm_judge_simple_agent.responses_api_agents.simple_agent.datasets=[]'

The two Hydra overrides set the CRITPT and HLE dataset lists to empty, excluding those benchmarks from the suite. Remove them to include those benchmarks in the run.

Resuming after preemption or a manual restart

The main use case for RESUME_EVAL_ON_REQUEUE=1 is unattended recovery when Slurm preempts and automatically requeues a job. Without the switch or an explicit fixed results path, each restart uses a new timestamped path and cannot find the earlier results. The switch uses a stable experiment/results name and injects ++resume_from_cache=true; it does not configure Slurm's preemption or requeue policy.

For a manual resubmission, such as after reaching the walltime limit, wait until the previous job has stopped and repeat the command with the same experiment name, model, benchmark settings, and RESUME_EVAL_ON_REQUEUE=1. Preserve both resumable.jsonl and resumable_materialized_inputs.jsonl (the exact tasks selected for the run); Gym skips saved completions. Without both files, Gym starts fresh instead of resuming. For an older run with a different filename, also set ROLLOUTS_FPATH to its existing rollout JSONL file.

Branch-added command controls

Explicit Hydra arguments such as ++num_samples_in_parallel=32 and ++resume_from_cache=false take precedence over the corresponding environment defaults. Disabling resume explicitly does not disable the switch's stable experiment/results naming.

Variable	Valid values	Purpose and known-good choice
SBATCH_TIME	Any Slurm duration, such as 20:00:00 or 7-00:00:00	Makes the previously fixed four-hour limit configurable. The known-good validation job completed in about 14.6 active hours across three allocations.
NUM_SAMPLES_IN_PARALLEL	Unset or a positive integer	Controls how many benchmark samples Gym processes concurrently. We used 16; 32 increased total throughput but reduced per-request speed and increased waiting.
RESUME_EVAL_ON_REQUEUE	0 (default) or 1	Primarily handles automatic restarts after Slurm preemption. 1 removes job-ID/timestamp components from the experiment name and enables cached-result reuse. The validation used 1 and preserved progress across two preemptions.
VLLM_PD_DEPLOYMENT_MODE	independent or coupled	coupled joins multiple nodes into one distributed prefill/decode engine. Ultra requires it because its weights cannot fit on one four-GPU node.
VLLM_SLURM_SEGMENT	Unset or a positive integer	Sets the Slurm allocation segment size. We used 4 to align each segment with one four-node prefill or decode tier and avoid distributed startup hangs.
ULTRA_ENABLE_MTP	0 or 1	Enables multi-token prediction, where the model drafts several future tokens and verifies them together. We used 1 because it was necessary to exceed 90 tokens/sec/request.
ULTRA_NUM_SPECULATIVE_TOKENS	A positive integer supported by the model	Controls how many extra tokens MTP drafts. We used 3; it cleared the throughput target while producing substantially fewer repetition failures than MTP5.
ULTRA_DECODE_ENFORCE_EAGER	0 or 1	1 disables compilation and CUDA graphs; 0 permits them. We used 0 to retain the faster graph-backed decode path.
ULTRA_DECODE_CUDAGRAPH_MODE	FULL_DECODE_ONLY, PIECEWISE, or NONE	CUDA graphs replay captured GPU work with less CPU launch overhead. PIECEWISE was the stable, favorable decode configuration.
ULTRA_DECODE_CUDAGRAPH_COPY_INPUTS	0 or 1	Copies inputs into stable buffers before graph replay. We used 1 for the validated piecewise CUDA-graph path.

[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.

[fsiino-nvidia]

/ok to test a26a1e2

[fsiino-nvidia]

/ok to test f656822

[bxyu-nvidia]

can we hardcode these once we have a roughly tuned config?

[bxyu-nvidia]

same as above, can we hardcode this one things are tuned?

[bxyu-nvidia]

i think this should be fixed with more recent containers

[bxyu-nvidia]

we should not enforce eager because it's much slower

[bxyu-nvidia]

we recently got async scheduling accuracy working!