feat(automodel): expose FSDP output precision

[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. This increases activation storage and materializes full FP32 logits before downstream loss chunking can help.

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.dtensor_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:
  dtensor_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

Adds test_automodel_output_precision.py, checking the mixed-precision policy produced by distributed setup for omitted, FP32, and BF16 output settings, including preservation of parameter and reduction precision.

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

tests/unit/models/test_automodel_output_precision.py:23

1 action item.

TL;DR — this file errors the four base L0_Unit_Tests_Models_* shards at collection and is deselected from the Automodel shard, so the new code has zero CI coverage. PR-introduced.

It has no nemo_automodel import guard and no @pytest.mark.automodel:

1. L0_Unit_Tests_Models_1.sh:23 runs unit/models/ in the base venv — docker/Dockerfile:263 ends with uv sync --frozen --all-groups --no-install-project (no extras), and nemo-automodel is only in the automodel extra (pyproject.toml:132).
2. Line 23 imports nemo_rl.models.automodel.setup, which imports nemo_automodel unconditionally at setup.py:26. Collection runs before marker filtering (conftest.py:151), so only a module-level guard can prevent this.
3. L0_Unit_Tests_Automodel.sh:21 passes --automodel-only, and conftest.py:149 keeps only marked items — so the unmarked test is dropped there.

Both halves reproduced locally:

$ pytest unit/models/ --collect-only --shard-id=0 --num-shards=4   # base env
ERROR collecting tests/unit/models/test_automodel_output_precision.py
E   ModuleNotFoundError: No module named 'nemo_automodel'
530 tests collected, 1 error in 6.26s

$ pytest unit/models/test_automodel_output_precision.py --automodel-only
20 warnings in 5.24s          # 0 tests collected

Action: fold both cases into the existing TestSetupDistributed class in tests/unit/models/automodel/test_automodel_setup.py and delete this file. That class already carries the marker, the module-level guard, and mock_config / mock_runtime_config / mock_device_mesh fixtures next to five existing setup_distributed tests — so the CI fix comes for free and the ad-hoc SimpleNamespace / bare-dict stand-ins go away.

Use a (configured, expected) table rather than computing the expectation in the assertion — the current form (torch.bfloat16 if output_dtype == "bfloat16" else torch.float32) would assert float32 if someone added "float16" to the list. Verified passing (5 passed; whole file 103 vs 98 before) and mutation-checked — reverting output_dtype=STRING_TO_DTYPE[...] to torch.float32 fails the bfloat16 case:

    @pytest.mark.parametrize(
        "configured_dtype,expected_output_dtype",
        [
            (None, torch.float32),
            ("float32", torch.float32),
            ("bfloat16", torch.bfloat16),
            ("float16", torch.float16),
        ],
    )
    @patch("nemo_rl.models.automodel.setup.MoEParallelizerConfig")
    @patch("nemo_rl.models.automodel.setup.MeshContext")
    @patch("nemo_rl.models.automodel.setup.torch.distributed")
    def test_setup_distributed_fsdp_output_dtype(
        self, mock_torch_dist, mock_mesh_context, mock_moe_config,
        mock_config, mock_runtime_config, mock_device_mesh,
        configured_dtype, expected_output_dtype,
    ):
        """fsdp_output_dtype only changes output_dtype; params/reductions are untouched."""
        mock_torch_dist.get_world_size.return_value = 8
        mock_moe_config.return_value = MagicMock()
        mock_mesh_context.build.return_value = SimpleNamespace(
            device_mesh=mock_device_mesh, moe_mesh=None
        )
        if configured_dtype is not None:
            mock_config["dtensor_cfg"]["fsdp_output_dtype"] = configured_dtype

        result = setup_distributed(mock_config, mock_runtime_config)

        mp_policy = result.fsdp2_config.mp_policy
        assert mp_policy.param_dtype == torch.bfloat16
        assert mp_policy.reduce_dtype == torch.float32
        assert mp_policy.output_dtype == expected_output_dtype

It deliberately does not patch FSDP2Config, so the real MixedPrecisionPolicy is built — worth noting that mp_policy contents have never been asserted anywhere in the repo before, so this is genuinely new coverage, just misplaced. Also worth adding a test_fsdp_output_dtype_validation_invalid beside test_precision_validation_invalid — the new ValueError has no test today.

[pstjohn]

Done in 5a9638c. test_automodel_output_precision.py is deleted; both the parametrized (configured, expected) table (test_setup_distributed_fsdp_output_dtype, 4 cases incl. float16) and test_fsdp_output_dtype_validation_invalid now live in test_automodel_setup.py — the former in TestSetupDistributed (inheriting the guard, marker, and fixtures), the latter beside test_precision_validation_invalid. The new tests deliberately do not patch FSDP2Config, so the real MixedPrecisionPolicy is asserted.

Verified locally: 98 → 103 collected under --automodel-only (matching the count from the review), unit/models/ collects cleanly in the base env (2175 tests, 0 errors), and the mutation check from the review still holds — reverting output_dtype=STRING_TO_DTYPE[...] to torch.float32 fails the bfloat16 case.

One note: your snippet patched only MoEParallelizerConfig, MeshContext, and torch.distributed; the shipped test matches that exactly.

[pstjohn]

nemo_rl/models/policy/__init__.py:206

1 action item.

TL;DR — DTensorConfig is shared by both DTensor backends, but only v2 reads this key, so on v1 fsdp_output_dtype: bfloat16 is a silent no-op. PR-introduced.

The v1 worker parallelizes through nemo_rl/models/dtensor/parallelize.py, which hardcodes output_dtype=torch.float32 at L510 and L748 and never reads this key. _v2 defaults to False — lm_policy.py:233: use_v2 = config.get("dtensor_cfg", {}).get("_v2", False).

This isn't hypothetical: seven shipped exemplars have a dtensor_cfg block with no _v2 key — dpo.yaml, rm.yaml, sft_openmathinstruct2.yaml, grpo_math_8B.yaml, grpo_sliding_puzzle.yaml, sft_avlm.yaml, vlm_mpo.yaml. A user derives from one of those, sets fsdp_output_dtype: bfloat16, gets no warning, and outputs stay fp32.

Action: add a guard in the v1 else: branch of nemo_rl/models/policy/lm_policy.py. The exact precedent is five lines away at lm_policy.py:254-259, which does this for dp_replicate_size, another v2-only key:

if config["dtensor_cfg"].get("fsdp_output_dtype", "float32") != "float32":
    raise ValueError(
        "fsdp_output_dtype requires policy.dtensor_cfg._v2: true "
        "(Automodel DTensor v2 backend). The V1 DTensor worker always "
        "produces float32 module outputs."
    )

This blocks nothing legitimate — on v1 the setting has no effect, so the only configurations it rejects are ones that are already silently doing nothing.

[pstjohn]

Done in 5a9638c — the guard is in the v1 else: branch of lm_policy.py, five lines below the dp_replicate_size precedent, with the suggested wording:

if (
    config["dtensor_cfg"].get("fsdp_output_dtype", "float32")
    != "float32"
):
    raise ValueError(
        "fsdp_output_dtype requires policy.dtensor_cfg._v2: true "
        "(Automodel DTensor v2 backend). The V1 DTensor worker always "
        "produces float32 module outputs."
    )

Covered by test_dtensor_fsdp_output_dtype_requires_v2 in tests/unit/models/policy/test_policy_validation.py, mirroring test_dtensor_dp_replicate_size_requires_v2 (asserts the raise and that RayWorkerGroup is never constructed). Passing locally; the one failure in that file (test_world_size_validation_megatron[...]) is pre-existing on this box — it fails identically on the unmodified branch because megatron.bridge isn't installed in the local env.

[pstjohn]

nemo_rl/models/automodel/setup.py:498

1 action item.

TL;DR — at the exemplar defaults (TP=1, logprob_chunk_size: null) the logit buffer measurably gets worse, so the claimed memory win needs one number to back it.

Two independent mechanisms are in play and they point opposite ways:

• Logits (works against you at the defaults). use_chunking is gated on vocab_parallel_group is not None and chunk_size is not None (model_utils.py:1838); otherwise the full .to(torch.float32) happens anyway. vocab_parallel_group is non-None only at TP>1. grpo_math_1B.yaml defaults logprob_chunk_size: null and tensor_parallel_size: 1. At those defaults today's fp32 output makes .to(torch.float32) a no-op returning self; with bf16 you get a bf16 tensor plus a fresh fp32 copy.
• Residual stream (works for you, TP- and chunking-independent). Each transformer block is its own FSDP unit sharded with the caller's mp_policy (Automodel moe/parallelizer.py:750-758), and torch casts every unit's forward output. Upstream names this mechanism explicitly (config.py:262-266): "Use with output_dtype=float32 in mp_policy to keep the residual stream in fp32 while running matmuls in lower precision."

So the win is real in principle, but whether it dominates depends on whether those per-layer fp32 boundary tensors are actually retained for backward — which isn't determinable from the code (the cast's backward is itself a cast and need not save its input, and activation checkpointing changes what's saved).

Action: post one before/after peak_memory_allocated_mb (already tracked, policy/utils.py:217) for fsdp_output_dtype: float32 vs bfloat16, ideally at both tensor_parallel_size=1, logprob_chunk_size=null (where the logit buffer gets worse) and tensor_parallel_size=2, logprob_chunk_size=1024. If the TP=1 case doesn't improve, narrowing the claim in the PR body and the field comment would help users pick the right setting.

[pstjohn]

Fair ask — deferring the number rather than guessing it. The measurement needs a full policy-worker run (and ideally the TP=2/chunked variant on multi-GPU), which I don't want to bolt onto this branch blind; the local box here is a single heterogeneous-GPU host, so a representative TP=2 datapoint isn't available in this environment.

Everything else from this comment is addressed in 5a9638c: the field comment now names the three opt-in precision edges so users can see the trade-off (see reply on the __init__.py comment). I'll follow up with peak_memory_allocated_mb before/after numbers at the exemplar defaults and at TP=2/chunked — and if the TP=1 case doesn't improve, I'll narrow the claim in the PR body and the field comment as suggested.

[pstjohn]

nemo_rl/models/policy/__init__.py:204

1 action item.

TL;DR — three places lose precision when you opt in, and none of them are mentioned; the comment should name them.

To be clear up front: all three are opt-in only. .get(..., "float32") preserves the previous hardcoded value exactly, so nothing changes for existing users. These are sharp edges to document, not regressions.

1. Value/RM training losses run in bf16. LossInputType.LOGIT is the one branch that hands model output straight to a loss with no upcast — loss/utils.py:416 is loss_input = {"logits": logits}, in contrast to the LOGPROB branch three lines below which does .to(torch.float32). So PreferenceLoss (loss_functions.py:1293) and ValueLoss (:1843) run logsigmoid / (values - returns)**2 / mse_loss in bf16. The knob reaches these workers via dtensor_value_worker_v2.py:161. RM scoring is fine — ScorePostProcessor upcasts at train.py:1105.
2. Temperature scaling is in-place, pre-upcast. train.py:245 does logits.div_(T) on the raw output (called from L362), before any fp32 cast. Measured at V=128256: ~2.7% mean multiplicative probability error at T=0.7, zero at T=1.0 (the function early-returns). No shipped v2 config uses T≠1.0 today, so this is latent rather than live.
3. moe_parallelizer.lm_head_precision: float32 is silently undone. Upstream gives lm_head its own all-fp32 FSDP unit (moe/parallelizer.py:794-806), but the root is then sharded with the caller's policy (:858) and torch's _post_forward casts the whole output tree — round-tripping the fp32 logits back to bf16. Requires a hand-written moe_parallelizer block, so nobody hits it today.

Action: extend the field comment so users can see the trade-off without tracing the code:

Suggested change

	# AutoModel FSDP2 output precision. Master weights/reductions remain FP32;
	# AutoModel FSDP2 output precision. Master weights/reductions remain FP32;
	# BF16 outputs avoid retaining full FP32 residual/logit buffers.
	# DTensor v2 only. Opting out of float32 lowers precision in three places:
	# value/RM losses (LossInputType.LOGIT is not upcast), in-place temperature
	# scaling when temperature != 1.0, and moe_parallelizer.lm_head_precision.

[pstjohn]

Done in 5a9638c — the field comment now reads:

# AutoModel FSDP2 output precision. Master weights/reductions remain FP32;
# BF16 outputs avoid retaining full FP32 residual/logit buffers.
# DTensor v2 only. Opting out of float32 lowers precision in three places:
# value/RM losses (LossInputType.LOGIT is not upcast), in-place temperature
# scaling when temperature != 1.0, and moe_parallelizer.lm_head_precision.

Good catch on all three, especially the LOGIT branch — none of those were visible from the config surface alone.

[pstjohn]

nemo_rl/models/automodel/setup.py:298

1 action item. PR-introduced, trivial.

This validates against a hand-written tuple but then reads from STRING_TO_DTYPE at L498 — two sources of truth for one set. The precision check three lines above already does it the other way (L294): if precision not in STRING_TO_DTYPE. No drift today (STRING_TO_DTYPE holds exactly these three keys), but the tuple is the copy that has to be updated by hand.

Action: match the sibling idiom.

Suggested change

	if fsdp_output_dtype not in ("float32", "bfloat16", "float16"):
	if fsdp_output_dtype not in STRING_TO_DTYPE:

[pstjohn]

Done in 5a9638c — now if fsdp_output_dtype not in STRING_TO_DTYPE:, matching the sibling precision check. No test change needed (the invalid-value test passes through the same ValueError).

[pstjohn]

nemo_rl/models/automodel/setup.py:297

No action needed — consistency note, your call.

Per config-conventions, for a v1 TypedDict the default belongs in the exemplar YAML and NotRequired fields shouldn't take a non-None default at the call site; fsdp_output_dtype currently appears in zero YAMLs and takes "float32" at two call sites (here and L469).

Flagging as a note rather than an ask because the three nearest analogs in this same TypedDict do exactly the same thing — defer_fsdp_grad_sync, moe_parallelizer and lm_head_precision are each read with a non-None .get default and each absent from every YAML. You're consistent with the neighbors; if you'd rather be consistent with the rule, grpo_math_1B.yaml's dtensor_cfg block is the place (and tests/unit/reference_configs/ would need the same key, per test_reference_configs_up_to_date).

[pstjohn]

Acknowledged, and thanks for the census on the neighbors — staying consistent with defer_fsdp_grad_sync/moe_parallelizer/lm_head_precision for now. If we later move to exemplar-YAML defaults per the convention, I'll do it for all four keys in one sweep plus the reference_configs update, rather than making fsdp_output_dtype the odd one out.

[pstjohn]

/ok to test 5a9638c