fix(distributed): control frozen multimodal FSDP sharding

[yuhezhang-ai]

Summary

This PR adds the FSDP2 knob distributed.multimodal.frozen_sharding for fully frozen multimodal modules. It accepts:

• root (new default): keep frozen vision/audio towers and projectors in the always-run outer FSDP root. This keeps collective ordering aligned when modality execution differs across data-parallel ranks.
• per_layer: shard the frozen modules individually. This preserves the previous dense-VLM behavior, but requires every rank in the FSDP group to execute or skip those units identically on every microbatch; selecting it logs a warning.
• replicate: exclude the frozen parameters from FSDP so every rank keeps a full copy without tower collectives.

Before this PR, dense VLMs implicitly used per_layer, while common MoE VLMs implicitly used root. This PR makes root the default for both code paths. The measured Gemma4 MoE example explicitly selects replicate; five rank-uniform text-only Gemma4 examples explicitly retain per_layer, including two joint-drafter examples selected from paired measurements.

Modules containing any trainable parameters keep their normal layer/container sharding. Nested dense and MoE VLM modules use the same policy resolution, including Gemma4's nested vision tower and projector.

For nested MoE/VLM models, root is rejected only when wrap_outer_model=False leaves a fully frozen multimodal module outside the inner text-model FSDP root. Those models can enable the outer root or choose per_layer with rank-uniform execution or replicate. The only current VLM-labeled wrap_outer_model=False examples are LoRA benchmarks, whose adapters make their towers trainable at FSDP application time, so they are unaffected.

The configuration guide and breaking-changes page document the knob, default change, tradeoffs, outer-root constraint, and per_layer execution contract.

Config audit and measurements

The final audit found 85 tracked frozen-multimodal FSDP2 configs:

• 52 dense-vision configs were initially affected by the default change. Five rank-uniform text-only configs now explicitly retain per_layer, leaving 47 that inherit root; excluding four Qwen2.5 configs slated for deprecation leaves 43 current configs inheriting the changed default.
• 29 common-MoE configs keep root ownership.
• 3 dense-audio configs already used root ownership.
• 6 configs explicitly select a policy: 5 text-only Gemma4 configs use per_layer (2 based on paired measurements), and 1 Gemma4 MoE config uses replicate based on paired measurements.

Paired 30-step H100 measurements used warmup-excluded mean throughput as the preselected decision metric. Each policy cell reports mean throughput delta; peak memory/GPU delta relative to root; a negative memory delta means lower peak memory:

Tracked config	GPUs	per_layer vs root (throughput; memory/GPU)	replicate vs root (throughput; memory/GPU)	Code path
Qwen3-VL 4B	2	-5.82%; -0.64 GiB	+3.89%; -0.38 GiB	per_layer (before) → root (after)
Qwen3-VL 4B	8	+1.11%; -0.61 GiB	+4.70% first run; +0.68% reversed confirmation (+2.65% across run means); -0.07 GiB	per_layer (before) → root (after)
Gemma3 4B	8	+1.24%; -0.74 GiB	+0.92%; -0.08 GiB	per_layer (before) → root (after)
Ministral3 3B	8	-0.99%; -0.72 GiB	+2.91%; -0.10 GiB	per_layer (before) → root (after)
Gemma4 26B-A4B MoE	8	+5.70%; -0.01 GiB	+4.72%; -0.28 GiB	root (before) → replicate (after)
Gemma4 4B joint drafter, text-only	8	+25.45%; -0.84 GiB	—	per_layer (before) → per_layer (after)
Gemma4 31B joint drafter, text-only	8	+5.45%; -1.01 GiB	—	per_layer (before) → per_layer (after)
Gemma4 31B Tulu, text-only CP8	—	Not measured; the dataset is structurally text-only and rank-uniform	—	per_layer (before) → per_layer (after)
Gemma4 E4B Tulu, text-only CP16	16	Not measurable: current dense-CP/shared-KV path stalls before step 0 under both policies	—	per_layer (before) → per_layer (after)
Gemma4 31B CoderForge, text-only CP8	128	Not measured: the exact-topology job failed before training and its backfill was cancelled	—	per_layer (before) → per_layer (after)

The full Gemma4 and Ministral3 train-loss sequences and validation losses matched exactly across all three policies. For Gemma4, replicate is only 0.93% slower than per_layer by measured mean throughput, so it is the safer example selection with essentially the same performance.

The two joint-drafter pairs also had identical 30-step train-loss sequences and validation loss. The 4B result was repeated in reversed order and reproduced the throughput delta (+25.24%).

All five explicit per_layer examples provide only text batches to the VLM recipe, so every rank uniformly skips the frozen multimodal modules and satisfies the collective-ordering contract. The E4B CP16 recipe was originally validated when added, but current main has since moved from Transformers 5.8.1 to 5.12.1 and refactored CP dispatch. Its untouched path first fails the newer explicit shared-KV contract; a benchmark-only compatibility bridge gets past that error, after which both policies reproduce the same CP-ring stall before step 0. This is a separate current-main dense-CP regression and does not change the rank-uniform safety argument for retaining its previous per_layer topology.

The earlier retained Gemma4 E4B matrix measured only +0.57% for replication. The deprecated Qwen2.5-VL result is supplemental only. Qwen3, Gemma3, and Ministral3 keep the root default because their gains were small or non-reproducible.

All Gemma3, Gemma4, and Ministral3 policy matrices had identical train-loss sequences; the dedicated distributed regression provides root/per-layer optimizer-update parity and asymmetric-execution coverage.

Validation

• Latest real two-rank FSDP2 regression: Slurm job 14341819, 1 passed.
• Latest merged-head focused unit aggregate: 635 passed, 1 skipped.
• Text-only example policy regression: 5 passed.
• Joint-drafter policy pairs: jobs 14324681 (4B) and 14324720 (31B), 30-step losses and validation loss identical.
• Gemma4 MoE replicate rerun: job 14337370, 30/30 steps completed.
• Policy matrices completed successfully: jobs 14297661, 14297677, 14297753, 14297767, 14299169, and 14323820.
• ruff format, ruff check, and git diff --check on the latest merged head.

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

[yuhezhang-ai]

/ok to test 217af2b

[yuhezhang-ai]

Triggered scoped real GPU CI for correctness on the current PR head:

• Parent pipeline: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/pipelines/55757619
• AUTOMODEL_REF: 217af2b0a62ab6a76fe918ba8cf010038a9d8ce8
• AUTOMODEL_TEST_FOLDER: vlm
• TESTS_TO_RUN_ON_THIS_COMMIT: release
• TEST_CLUSTER_REGEX: /eos/
• ARCH: x86
• Scoped configs: gemma3_vl_4b_medpix, gemma4_26b_a4b_moe, gemma4_4b, qwen3_5_4b

This is correctness-only. I still consider the performance question WIP until we run more controlled paired benchmarks.

[yuhezhang-ai]

Scoped real GPU CI passed for the current PR head.

• Parent pipeline: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/pipelines/55757619
• Automodel downstream: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/pipelines/55757786
• VLM child pipeline: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/pipelines/55757811

Passed jobs:

• gemma3_vl_4b_medpix: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/jobs/347503025
• gemma4_26b_a4b_moe: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/jobs/347503026
• gemma4_4b: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/jobs/347503027
• qwen3_5_4b: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/jobs/347503028

This was correctness coverage only; I am not treating this as a performance conclusion for replicate.

[akoumpa]

@yuhezhang-ai that's awesome; i found the mistral recipe 62a9898#diff-c5a488bfa4afe12e07697c2c0e708633517dcee3a2da8ff14b0d6f491b4123f0R147 to be affected by this, would it be ok to include it in the testing? 🙏 -- I'd very much revert 2744 and keep this. Thank you for jumping on this right away!

[athitten]

Thank you @yuhezhang-ai , can we also make sure convergence is not impacted by this. Ideally it shouldn't. If you like I can help check for gemma4s whenever the PR is ready, lmk

[yuhezhang-ai]

Added Alex's requested Mistral coverage. Scoped real GPU CI passed for ministral3_3b_squad on the current PR image/SHA.

• Parent pipeline: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/pipelines/55761174
• Automodel downstream: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/pipelines/55761183
• LLM child pipeline: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/pipelines/55761194
• ministral3_3b_squad: https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/jobs/347531221

Note: this recipe is a text-only LLM recipe that loads the Mistral3 FP8 VLM wrapper. Its vision tower is skipped by the data path but is not explicitly frozen by the recipe, so this validates that the current default behavior plus the existing FSDP accumulated-grad guard still works. The new frozen_multimodal_sharding=replicate knob would only replace that guard for this recipe if we first made the vision tower actually frozen before sharding.

[yuhezhang-ai]

While this PR is trying to add a opt in knob to replicate the frozen vision/audio tower, I found a bigger inconsistency in how we treat multimodal tower. See details below:

The issue is not really frozen vs trainable. It is whether a separately wrapped vision/audio FSDP unit actually runs for a batch.

If a batch skips that modality, the tower’s FSDP forward hook may never run, so PyTorch may not create _unsharded_param. During gradient accumulation, FSDP deferred post-backward can still touch that skipped unit and crash.

Today we effectively have two ways to avoid the problem:

1. Root-shard the skipped tower
   Do not wrap the tower separately. Let the outer/root FSDP unit own it. Since the root runs every batch, the lazy FSDP state exists. This is conservative, but coarse.

2. Per-layer/per-module shard the tower and patch PyTorch
   Wrap the tower normally. If the tower is skipped, its FSDP units do not all-gather. The PyTorch guard makes deferred post-backward tolerate skipped units with no _unsharded_param.

The code currently uses both approaches inconsistently. MoE root-shards frozen audio/visual towers. Dense VLMs often per-layer shard frozen vision through the generic layer policy, while frozen audio has special root-shard handling.

Proposed direction:

• Keep the user knob simple:

  frozen_multimodal_sharding: shard      # default
  # or
  frozen_multimodal_sharding: replicate

• Users should not need to choose between root-shard and per-layer-shard.
• Internally, make shard mean normal per-layer/container FSDP sharding for frozen multimodal towers, same as trainable towers. So we remove the root-shard method 1, and rely on the PyTorch skipped-FSDP guard enabled.
• Use replicate only when the user explicitly wants frozen tower params unsharded on every rank.

Risk: this relies on a narrow monkey patch against private PyTorch FSDP2 internals. The patch only handles the case where _unsharded_param was never created because the FSDP unit did not run, but it is still a compatibility workaround until PyTorch fixes this upstream.

This would remove the dense/MoE inconsistency and expose only the real user choice: shard frozen multimodal params or replicate them.

[yuhezhang-ai]

Follow-up after looking more closely at the rank-asymmetric mixed-media case: I think the proposed direction above should change slightly. This proposal is scoped to fully frozen multimodal towers; trainable conditional towers are a separate problem and are out of scope for this PR.

The accumulated-grad guard from #2744 and root sharding solve two different failures:

• If every rank skips a separately wrapped tower, there is no forward collective mismatch. The tower never creates _unsharded_param, and the fix(fsdp2): guard uninitialized accumulated grads #2744 guard prevents deferred post-backward bookkeeping from crashing.
• If some ranks run the tower and others skip it for the same microbatch, media ranks enter the tower FSDP all-gathers while text-only ranks do not. The process group can hang or desynchronize during forward, before the fix(fsdp2): guard uninitialized accumulated grads #2744 guard is relevant. The guard does not insert the missing forward collective.

For fully frozen towers, I think the user-facing policy should expose three explicit choices:

frozen_multimodal_sharding: root       # default, correctness-first
# frozen_multimodal_sharding: per_layer  # expert opt-in
# frozen_multimodal_sharding: replicate  # full copy on each rank

Proposed semantics:

1. root (default):

   • Resolve both top-level and nested multimodal towers before recursive wrapping.
   • Exclude their frozen parameters from standalone/per-layer FSDP units.
   • Let the outer/root FSDP unit own those parameters.
   • Because the root runs on every rank, the collective sequence remains aligned even when a rank skips the modality branch.

2. per_layer (expert opt-in):

   • Keep normal per-layer/container FSDP sharding.
   • Require the user to guarantee collective-uniform execution: every rank in the FSDP group must call or skip the tower the same number of times, in the same order, on every microbatch.
   • Retain fix(fsdp2): guard uninitialized accumulated grads #2744 as the defensive guard for the all-ranks-skip case.
   • This preserves the lower peak-unshard-memory and skip-all-vision-communication advantages when the data contract is known to be uniform.

3. replicate:

   • Exclude the fully frozen tower from all FSDP roots and keep a full copy on each rank.
   • This avoids tower collectives but increases per-rank parameter memory.

This intentionally changes the default topology for dense frozen vision, which is generally per-layer sharded today. Frozen audio and MoE frozen towers are already root-owned in the common paths, so making root the default makes vision/audio and dense/MoE behavior consistent. Trainable towers should remain unchanged by this policy.

wrap_outer_model=False should remain supported:

• root + wrap_outer_model=False: fail clearly because there is no outer FSDP root to own the tower.
• per_layer + wrap_outer_model=False: allowed under the explicit uniform-execution contract.
• replicate + wrap_outer_model=False: allowed.

The root policy should solve the dense frozen-vision mixed-data case under a full DP x CP shard group and remove the need for a dummy ViT forward on text-only ranks. The tradeoff is coarser unsharding and potentially higher peak memory, so root versus per-layer versus replicate should still be benchmarked.

For coverage, I suggest:

• A real distributed regression with two ranks in the same FSDP group, one receiving media and the other a text-only microbatch, with gradient accumulation enabled, validating root.
• An all-ranks-text-only case validating per_layer plus the fix(fsdp2): guard uninitialized accumulated grads #2744 skipped-unit behavior.
• Policy/config tests covering the three wrap_outer_model combinations above.

[yuhezhang-ai]

Follow-up audit and measurements for the proposed root / per_layer / replicate policy:

Decision: keep root as the default and do not add overrides to retained example configs. per_layer remains an
expert opt-in with a rank-uniform execution warning; replicate remains the collective-free option.

Retained-config benchmark

I ran the tracked examples/vlm_finetune/gemma4/gemma4_4b.yaml config for 15 train steps plus validation on 8 x H100
80 GB, using the same allocation/data order for all policies (Slurm job 14295517, all three runs completed):

policy	mean tok/s*	median tok/s*	peak memory	final train loss	val loss
root	4,305	4,249	36.91 GiB	2.1486	2.3141
per_layer	4,325	4,218	36.07 GiB	2.1486	2.3141
replicate	4,330	4,385	36.70 GiB	2.1486	2.3141

* excluding steps 0-1.

The complete 15-step train-loss sequence and validation loss were identical across policies. Relative to per_layer,
root was -0.47% by mean throughput, +0.71% by median throughput, and +0.84 GiB (+2.33%) peak memory. This is small
enough to prefer the safer root topology.

I also completed a supplemental Qwen2.5-VL 3B matrix before noting that config is slated for deprecation; it reached
the same decision, but I am treating Gemma4 as the relevant retained-config evidence.

Config audit

• 76 tracked frozen-multimodal FSDP2 example configs found.
• 45 dense-vision configs change effective topology (44 after excluding the Qwen2.5-VL config slated for
  deprecation).
• 27 MoE configs already had root ownership.
• 4 dense-audio configs already had root ownership.
• No retained config needs an explicit override based on the Gemma4 measurements.

The only two VLM-labeled configs with wrap_outer_model: false are LoRA benchmark configs (kimi25vl_lora.yaml,
mistral4_lora.yaml). Their trainable adapters exist when FSDP is applied, so they are not fully frozen at the policy
decision point and are not rejected by the narrow root validation.

Regression coverage

A real two-rank FSDP2 regression now passes (job 14295827) for:

• root with rank-asymmetric modality execution and two accumulated microbatches;
• replicate with rank-asymmetric modality execution;
• per_layer when all ranks uniformly skip the tower (the fix(fsdp2): guard uninitialized accumulated grads #2744 lazy-state guard case);
• root/per-layer optimizer-update parity under uniform modality execution.

Focused unit coverage is also green: 473 passed, 1 skipped; Ruff and git diff --check pass.

[github-actions]

🌿 Preview your docs: https://nvidia-preview-yuhez-fix-gemma4-frozen-vlm-fsdp.docs.buildwithfern.com/nemo/automodel

[yuhezhang-ai]

Post-rebase audit and benchmark update for root / per_layer / replicate:

Decision: keep root as the default and do not add an explicit policy to tracked configs. replicate did not
produce a reproducible, meaningfully-greater-than-5% throughput win. per_layer remains an expert opt-in with the
rank-uniform execution warning.

Current-config measurements

All runs below used the rebased signed head 360085678, real tracked configs/checkpoints, and paired policy runs on
the same allocation/data order. Thirty-step means exclude steps 0-4.

Config	GPUs	replicate vs root mean throughput	Memory result	Conclusion
Qwen3-VL 4B	2	+3.89%	replicate -0.38 GiB peak	below threshold
Qwen3-VL 4B	8	+4.70% first run	replicate -0.07 GiB peak	borderline; repeated
Qwen3-VL 4B, reversed order	8	+0.68%	replicate -0.07 GiB peak	not reproduced
Gemma3 4B	8	+0.92%	replicate -0.08 GiB steady peak	below threshold
Gemma4 26B-A4B MoE, EP8	8	+4.72%	replicate -0.28 GiB steady peak	below threshold

The two Qwen3 8-GPU run means together give +2.65% for replication. The earlier retained Gemma4 E4B matrix measured
only +0.57%. The deprecated Qwen2.5-VL result remains supplemental only.

Gemma3 and both Gemma4 matrices produced identical train-loss sequences across policies. Gemma4 MoE per_layer was
+5.70% by mean throughput versus root, but it requires every rank to execute/skip the multimodal units identically
on every microbatch. I do not think that speedup justifies encoding the unsafe assumption in the example; the new
warning was emitted as intended.

Jobs: 14297661 (Qwen3 2 GPU), 14297677 (Qwen3 8 GPU), 14299169 (Qwen3 reversed confirmation), 14297753
(Gemma3), and 14297767 (Gemma4 MoE).

Post-rebase config audit

• 85 tracked frozen-multimodal FSDP2 configs.
• 52 dense-vision configs change effective topology; excluding four Qwen2.5 configs slated for deprecation leaves
  48 current configs.
• 30 common-MoE configs and 3 dense-audio configs already used root ownership.
• The only VLM-labeled wrap_outer_model: false examples are LoRA benchmarks. Their tower adapters are trainable
  when FSDP is applied, so the narrow fully-frozen root rejection does not affect them.

Regression and validation

• Rebased real two-rank FSDP2 regression: job 14297501, 1 passed.
• Covers asymmetric modality execution for root and replicate, all-ranks-skip for per_layer, and
  root/per-layer optimizer-update parity.
• Focused unit suites: 489 passed, 1 skipped.
• Ruff and git diff --check pass.
• The rebased commit is DCO-signed; the GitHub DCO check passes.

[yuhezhang-ai]

/ok to test 3600856

[yuhezhang-ai]

/ok to test dd7a867

[jgerh]

Important: This PR was assigned to me on 7/27 and was merged on 7/28 before the technical publications review of the markdown files could be completed. Please note that I require a 48-hour window to review PRs, and a review of all customer-facing content is a mandatory step prior to release. This is not optional. Please open a new PR containing the markdown files and assign it to me for review. Thank you.

[yuhezhang-ai]

Important: This PR was assigned to me on 7/27 and was merged on 7/28 before the technical publications review of the markdown files could be completed. Please note that I require a 48-hour window to review PRs, and a review of all customer-facing content is a mandatory step prior to release. This is not optional. Please open a new PR containing the markdown files and assign it to me for review. Thank you.

Thanks, the new PR is opened here for the doc:
#3272