[pstjohn]

What does this PR do ?

Allow memory-constrained AutoModel training to use BF16 FSDP outputs instead of the currently hardcoded FP32 outputs.

NeMo-RL forces FSDP module outputs to FP32 even when compute uses BF16. Lower output precision can reduce residual-stream activation storage; total peak savings depend on the model and loss path. Unchunked logprobs still materialize full FP32 logits.

AutoModel already supports independent output-precision configuration through FSDP2Config.mp_policy, including BF16 outputs in its existing configuration. This PR exposes that capability through policy.automodel_cfg.fsdp_output_dtype.

The option accepts float32, bfloat16, or float16. Omitting it preserves the current FP32 output behavior. The setting controls output precision independently of compute precision, FP32 master-weight loading, and FP32 gradient reductions.

Issues

No linked issue. Exposes an existing AutoModel capability through NeMo-RL configuration.

Usage

To opt into BF16 FSDP outputs:

policy:
  automodel_cfg:
    fsdp_output_dtype: bfloat16

Lower output precision changes activation numerics and should be validated for the target model and training recipe.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

Extends tests/unit/models/automodel/test_automodel_setup.py to check omitted, FP32, BF16, and FP16 output settings, preservation of parameter/reduction precision, and rejection of invalid output dtypes. All 108 tests in that module pass; Ruff lint and formatting checks pass for the changed Python files.

Local 8×H100 memory validation

Rebased onto upstream main b2de3de15. The setting now lives under policy.automodel_cfg, matching main's Automodel migration; the removed DTensor V1 guard is no longer needed.

Ran all four requested conditions using the full Policy / AutomodelPolicyWorker training path and ClippedPGLossFn, on 8×H100 80GB. Used the grpo_math_1B.yaml policy defaults: Qwen/Qwen2.5-1.5B, BF16 compute, global batch 512, microbatch 4, sequence length 512, sequence packing enabled, activation checkpointing disabled. Inputs were fixed synthetic token batches (seed 42); no generation engine or reference model was instantiated. This measures policy-worker training memory, not end-to-end GRPO memory or convergence.

For each condition, reset peak statistics before each update, discard one warm-up update, and measure two further updates. Values below are the maximum peak_memory_allocated_mb across all eight ranks (MiB). Both measured updates produced the same peak in each condition.

TP	logprob_chunk_size	FP32 outputs	BF16 outputs	Reduction
1	null	13,216.17	12,016.07	1,200.10 MiB (9.08%)
2	1024	8,957.59	8,069.39	888.20 MiB (9.92%)

Loss and gradient norms were finite in every update. BF16 reduced the measured peak in both cases, including TP=1 without logprob chunking; this does not establish convergence equivalence or a universal memory saving.

Environment: PyTorch 2.13.0+cu130, Transformers 5.12.1, AutoModel b916107a5ce8e77a3f18682f8b0eaa46a4b54389; Qwen model revision 8faed761d45a263340a0528343f099c05c9a4323. Set NCCL_NVLS_ENABLE=0 consistently because this host's NVLS multicast binding failed during initial setup. Training implementation measured at 1219f60347f7170de11f8d9538650c346bf61a3e; subsequent edits only clarify documentation.

Local reproducibility artifacts (benchmark script, launcher, resolved policy configs, per-rank/per-update JSON, logs, and summary): /data/pstjohn/pr4115-output-precision-20261023/.

Local paired GRPO convergence comparison (2026-10-27)

Left: sampled accuracy on 512 fixed held-out questions. Right: training reward with a 10-update moving average; faint lines show individual updates. One training seed per condition.

Ran a real-rollout comparison on the current PR code (29aeb99fc), using the shipped Qwen2.5-Math-1.5B-Instruct Automodel GRPO recipe adapted to GSM8K. Both runs completed 100 optimizer updates with rollout generation, reference-policy KL, reward verification, weight refitting, and final checkpoint saving. Each used four local H100 80GB GPUs in separate Docker/Ray clusters (eight GPUs total), taking about 40 minutes per run.

The two resolved configs differed only in policy.automodel_cfg.fsdp_output_dtype (float32 versus bfloat16) and output paths. Both used seed 42, the same model revision and training-prompt order, 32 prompts × 8 responses per update (25,600 training responses per run), global batch 256, microbatch 4, BF16 compute, TP=1, no logprob chunking, dynamic batching, no sequence packing or activation checkpointing, total sequence length 1536, and a 1024-token generation cap. AdamW learning rate was 5e-6 with the recipe's 50-update linear warm-up. vLLM used eager execution in both conditions.

Evaluated the same first 512 held-out GSM8K test questions before training and every 20 updates. Validation uses the native recipe's sampled decoding (temperature 1.0, top-p 1.0, one response per question), not greedy decoding. The two initial evaluation records were byte-identical; user prompts were verified identical across every validation checkpoint and across all 100 paired training batches.

GRPO updates	FP32 outputs: held-out accuracy	BF16 outputs: held-out accuracy
0	80.86%	80.86%
20	80.86%	80.27%
40	83.79%	83.01%
60	83.01%	84.38%
80	86.13%	84.77%
100	85.35%	86.52%

At update 100, FP32 scored 437/512 (85.35%) and BF16 443/512 (86.52%), from the same 414/512 (80.86%) starting point. Mean training reward over the last 20 updates was 95.10% versus 95.21%. Every recorded training loss and gradient norm was finite, and both runs exited successfully with saved step-100 checkpoints.

Interpretation: both conditions showed similar early learning and stability in this small example, with no obvious BF16 regression. The endpoint difference was +1.17 percentage points for BF16; a paired bootstrap over the 512 evaluated questions gave a 95% interval of [-1.37, +3.91] points. That interval is conditional on these two trained models and sampled answers; it does not include training-seed variability. This is one seed and 100 updates, so it supports this representative early-convergence check rather than a claim of general convergence equivalence or BF16 superiority.

Model revision: aafeb0fc6f22cbf0eaeed126eff8be45b0360a35. Training/generation backend environment: PyTorch 2.13.0+cu130, Transformers 5.12.1, AutoModel b916107a5ce8e77a3f18682f8b0eaa46a4b54389; Ray 2.58.0. The driver used Transformers 5.16.1. NCCL_NVLS_ENABLE=0 was set consistently for this host.

Local reproducibility records, committed experiment configs, complete TensorBoard curves, per-question outputs, final checkpoints, summary JSON, and convergence plots: /data/pstjohn/pr4115-convergence-20261027/. Experiment commits: FP32 b96714375, BF16 43580413a (local research branches; no training-code changes to the PR).

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

[pstjohn]

Reviewed by a team of six agents (RL/codebase, Automodel upstream, tests, bug scan, design, devil's advocate), with every finding put through an adversarial pass.

The mechanism is sound. Verified against upstream that output_dtype is a real torch.distributed.fsdp.MixedPrecisionPolicy field, honored at every FSDP unit boundary, and that bf16 outputs are Automodel's own default — this moves NeMo-RL toward upstream, not away from it. Default behavior is byte-identical to before, and there is no logprob precision regression on the default path: every numerically sensitive reduction already upcasts to fp32 before log_softmax.

Two blocking items: the new test file breaks the base L0_Unit_Tests_Models_* shards at collection and runs in no CI job at all; and fsdp_output_dtype is a silent no-op on the DTensor v1 path, which seven shipped exemplars still use.

The remaining four are suggestions. The one I'd most encourage is a single peak-memory number — the benefit is currently unmeasured, and at the exemplar defaults one of the two mechanisms works against you.

Five candidate findings were dropped during the adversarial pass, including a recommendation to remove float16 (the premise was false — Automodel ships and tests output_dtype=torch.float16) and a request to plumb the value through RuntimeConfig (tp_size/cp_size do the identical double-read, so the PR follows the local shape).

pre-commit run --all-files passes.

Generated by Claude Code

[pstjohn]

/ok to test 5a9638c

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[sharonyu-115]

This PR makes Automodel’s FSDP output dtype configurable as FP32, BF16, or FP16, preserving FP32 as NeMo-RL’s default. Parameter precision and FP32 gradient reduction remain unchanged.

Tests cover the default, all supported values, and invalid input. The reported GPU benchmarks demonstrate memory savings with finite losses and gradients. Upstream Automodel already defaults to BF16 output.

LGTM. Longer-run RL convergence remains unvalidated by the short benchmark.

[sharonyu-115]

/ok to test 29aeb99

[sharonyu-115]

/ok to test 9c7d58a