fix(distributed): control frozen multimodal FSDP sharding

docs/breaking-changes.mdx

@@ -5,6 +5,26 @@ position: 6
 ---
 ## 0.5.0
 
+### Frozen Multimodal FSDP2 Modules Default to Root Sharding
+
+Fully frozen vision/audio towers and multimodal projectors now default to the `root` FSDP2 policy. Previously, dense
+frozen vision towers commonly used per-layer FSDP units. Root ownership keeps the collective sequence aligned when
+some ranks execute a modality branch and other ranks skip it, but it can increase peak unshard memory.
+
+To retain the previous dense vision topology, opt in explicitly:
+
+```yaml
+distributed:
+  strategy: fsdp2
+  multimodal:
+    frozen_sharding: per_layer
+```
+
+`per_layer` requires every rank in the FSDP group to execute or skip the frozen multimodal units the same number of
+times and in the same order on every microbatch. Use `replicate` instead to avoid those collectives while keeping a
+full copy of the frozen parameters on every rank. See [YAML Configuration](/development/configuration) for all three
+policies and the `wrap_outer_model` constraint.
+
 ### FSDP2 Default `reduce_dtype` Is Now `float32`
 
 The default [`MixedPrecisionPolicy`](https://docs.pytorch.org/docs/stable/distributed.fsdp.fully_shard.html) built by `FSDP2Config` now uses `reduce_dtype=torch.float32` instead of `torch.bfloat16`. Forward/backward compute still uses `param_dtype=torch.bfloat16`, but gradient reduction now accumulates in fp32 to reduce communication-rounding error at larger data-parallel world sizes.

docs/guides/configuration.mdx

@@ -103,6 +103,25 @@ This option is intended for diagnosed checkpoint defects where rare token IDs pr
 
 For diagnosis, safety behavior, verification steps, and performance impact, see [Repair Damaged Input-Embedding Rows](/model-coverage/troubleshooting#repair-damaged-input-embedding-rows).
 
+### Frozen Multimodal FSDP Policy
+
+FSDP2 uses the correctness-first `root` policy for fully frozen vision/audio towers and multimodal projectors:
+
+```yaml
+distributed:
+  strategy: fsdp2
+  multimodal:
+    frozen_sharding: root
+```
+
+The supported policies are:
+
+- `root` (default): the always-run outer FSDP root owns frozen multimodal parameters. This keeps collective ordering aligned when some ranks execute a modality branch and others skip it.
+- `per_layer`: layers inside the frozen multimodal module get their normal FSDP units. This can reduce peak unshard memory, but it is an expert option: every rank in the FSDP group must execute or skip those units the same number of times and in the same order on every microbatch.
+- `replicate`: frozen multimodal parameters are excluded from FSDP and copied on every rank. This removes their collectives at the cost of higher per-rank parameter memory.
+
+This setting does not change modules that contain any trainable parameters. For nested MoE/VLM models, `root` requires `distributed.moe.wrap_outer_model: true` when a fully frozen multimodal module is outside the inner text-model FSDP root. Use `per_layer` or `replicate` if the outer model must remain unwrapped.
+
 ## Interpolate Environment Variables in YAML
 
 NeMo AutoModel supports env var interpolation inside YAML **string values**.

examples/long_context_validation/gemma4_31B/gemma4_31b_base_coderforge_cp8_64k_1e5_800steps.yaml

@@ -55,6 +55,9 @@ checkpoint:
 
 distributed:
   strategy: fsdp2
+  multimodal:
+    # Every rank receives text-only batches, so the frozen multimodal modules are skipped uniformly.
+    frozen_sharding: per_layer
   dp_size: none               # auto: 128 / (tp1 * cp8) = 16
   tp_size: 1
   cp_size: 8                  # 65536 / (2*8) = 4096 -> divisible; 8192 tokens/rank

examples/vlm_finetune/gemma4/gemma4_26b_a4b_moe.yaml

@@ -65,6 +65,8 @@ checkpoint:
 
 distributed:
   strategy: fsdp2
+  multimodal:
+    frozen_sharding: replicate
   dp_size: none
   tp_size: 1
   cp_size: 1

examples/vlm_finetune/gemma4/gemma4_31b_tulu3_text_cp8_16k.yaml

@@ -56,6 +56,9 @@ checkpoint:
 
 distributed:
   strategy: fsdp2
+  multimodal:
+    # Every rank receives text-only batches, so the frozen multimodal modules are skipped uniformly.
+    frozen_sharding: per_layer
   dp_size: none
   tp_size: 1
   cp_size: 8

examples/vlm_finetune/gemma4/gemma4_e4b_tulu3_text_cp16_64k.yaml

@@ -77,6 +77,9 @@ checkpoint:
 
 distributed:
   strategy: fsdp2
+  multimodal:
+    # Every rank receives text-only batches, so the frozen multimodal modules are skipped uniformly.
+    frozen_sharding: per_layer
   dp_size: none
   tp_size: 1
   cp_size: 16

examples/vlm_finetune/gemma4_joint_drafter/gemma4_31b_joint_drafter_tulu_magicoder_mix.yaml

@@ -83,6 +83,9 @@ checkpoint:
 
 distributed:
   strategy: fsdp2
+  multimodal:
+    # Every rank receives text-only batches, so the frozen vision tower is skipped uniformly.
+    frozen_sharding: per_layer
   dp_size: none
   # Joint drafter is TP/PP-unsafe (the drafter cross-attends the base's shared KV);
   # keep tp_size/cp_size/pp_size at 1 and shard the 31B base with FSDP2 instead.

examples/vlm_finetune/gemma4_joint_drafter/gemma4_4b_joint_drafter_tulu_magicoder_mix.yaml

@@ -74,6 +74,9 @@ checkpoint:
 
 distributed:
   strategy: fsdp2
+  multimodal:
+    # Every rank receives text-only batches, so the frozen vision tower is skipped uniformly.
+    frozen_sharding: per_layer
   dp_size: none
   tp_size: 1
   cp_size: 1

nemo_automodel/_transformers/infrastructure.py

@@ -413,6 +413,10 @@ def instantiate_infrastructure(
                 "enable_async_tensor_parallel",
                 bool(model_wrapper.enable_async_tensor_parallel),
             )
+            moe_kwargs.setdefault(
+                "frozen_multimodal_sharding",
+                model_wrapper.frozen_multimodal_sharding,
+            )
         parallelize_fn = partial(
             parallelize_model,
             activation_checkpointing=activation_checkpointing,

nemo_automodel/components/distributed/__init__.py

@@ -18,6 +18,7 @@
     FSDP2Config,
     MegatronFSDPConfig,
     MoEParallelizerConfig,
+    MultimodalDistributedConfig,
 )
 from nemo_automodel.components.distributed.init_utils import DistInfo, initialize_distributed
 from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes
@@ -31,6 +32,7 @@
     "MegatronFSDPConfig",
     "MeshContext",
     "MoEParallelizerConfig",
+    "MultimodalDistributedConfig",
     "ParallelismSizes",
     "PipelineConfig",
     "initialize_distributed",

nemo_automodel/components/distributed/config.py

@@ -43,6 +43,8 @@
 import torch
 from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy
 
+from nemo_automodel.shared.multimodal_fsdp import FrozenMultimodalSharding, normalize_frozen_multimodal_sharding
+
 if TYPE_CHECKING:
     from nemo_automodel.components.distributed.mesh import MeshContext, ParallelismSizes
     from nemo_automodel.components.distributed.pipelining.config import PipelineConfig
@@ -181,6 +183,29 @@ def to_dict(self) -> Dict[str, Any]:
         return {f.name: getattr(self, f.name) for f in fields(self)}
 
 
+@dataclass
+class MultimodalDistributedConfig:
+    """Distributed policies for resolved multimodal modules.
+
+    Attributes:
+        frozen_sharding: Controls fully frozen multimodal modules such as
+            vision/audio towers and projectors. ``"root"`` (default) keeps
+            their parameters in an always-run outer FSDP root, which is safe
+            when modality execution differs across ranks. ``"per_layer"``
+            uses normal layer/container FSDP units and requires every rank in
+            the FSDP group to execute or skip the module identically on every
+            microbatch. ``"replicate"`` excludes the frozen parameters from
+            FSDP roots so each rank keeps a full copy. Modules with any
+            trainable parameters use normal layer/container sharding
+            regardless of this setting.
+    """
+
+    frozen_sharding: FrozenMultimodalSharding = "root"
+
+    def __post_init__(self) -> None:
+        self.frozen_sharding = normalize_frozen_multimodal_sharding(self.frozen_sharding)
+
+
 @dataclass
 class FSDP2Config:
     """
@@ -250,6 +275,8 @@ class FSDP2Config:
             memory at a small throughput cost.  Default ``2``.
         fsdp2_forward_prefetch_depth (int): Number of FSDP units to prefetch during
             forward pass.  Default ``1``.
+        multimodal: Policies for resolved multimodal modules that use the text
+            model's distributed axes.
     """
 
     sequence_parallel: bool = False
@@ -274,8 +301,9 @@ class FSDP2Config:
     enable_fsdp2_prefetch: bool = False
     fsdp2_backward_prefetch_depth: int = 2
     fsdp2_forward_prefetch_depth: int = 1
+    multimodal: MultimodalDistributedConfig = field(default_factory=MultimodalDistributedConfig)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.mp_policy is None:
             # FSDP2 default: bf16 compute and fp32 gradient reduction. Pair with
             # ``model.torch_dtype: float32`` for fp32 optimizer state. See
@@ -289,6 +317,8 @@ def __post_init__(self):
         self.activation_checkpointing_scope = normalize_activation_checkpointing_scope(
             self.activation_checkpointing_scope
         )
+        if not isinstance(self.multimodal, MultimodalDistributedConfig):
+            raise TypeError("FSDP2Config.multimodal must be a MultimodalDistributedConfig instance.")
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert config to dictionary (shallow, preserves policy objects)."""

nemo_automodel/components/distributed/fsdp2.py

@@ -111,6 +111,7 @@ def __init__(
         self.enable_fsdp2_prefetch = config.enable_fsdp2_prefetch
         self.fsdp2_backward_prefetch_depth = config.fsdp2_backward_prefetch_depth
         self.fsdp2_forward_prefetch_depth = config.fsdp2_forward_prefetch_depth
+        self.frozen_multimodal_sharding = config.multimodal.frozen_sharding
 
     def parallelize(self, model):
         """
@@ -169,6 +170,7 @@ def parallelize(self, model):
             fsdp2_forward_prefetch_depth=self.fsdp2_forward_prefetch_depth,
             reshard_after_forward=self.reshard_after_forward,
             activation_checkpointing_scope=self.activation_checkpointing_scope,
+            frozen_multimodal_sharding=self.frozen_multimodal_sharding,
         )
 
         return model