refactor(distributed): unify multimodal module discovery and policy configuration

Motivation

#2763 exposed a recurring problem: several features independently decide which model submodules
are multimodal, and those private taxonomies drift into distributed-correctness bugs. The next
rank-asymmetric mixed-modality safety feature also needs richer shared metadata: modality, role,
and whether a path may be skipped.

This issue proposes:

1. one resolver for multimodal module identity;
2. one distributed.multimodal namespace for policies applied to those modules; and
3. a later, separately tested migration of freeze_config to the same addressing vocabulary.

Problem

“Which submodule is the vision/audio tower?” is currently answered independently:

Owner	Match rule
_get_model_layer_group_specs()	per-class FQN paths in language / vision / audio groups
shared/multimodal_fsdp.py	exact leaf-name taxonomy for dense and MoE FSDP
apply_parameter_freezing()	attributes plus substring matching on lowercased FQNs

Before #2763, FSDP creation and enumeration also used separate lists. That caused real bugs:

• BAGEL's vit_model existed in the MoE tuple but not the shared FSDP list.
• Gemma4's sibling embed_vision and embed_audio modules—instances of the same class—could
  receive different sharding treatment because only one name was listed.
• _iter_fsdp_modules() could miss FSDP units created by apply_fsdp(), excluding them from the
  MoE gradient-accumulation synchronization state machine.

#2763 fixes those immediate FSDP inconsistencies by sharing one taxonomy. The architectural
duplication remains: activation checkpointing still uses the layer-group table, freezing uses
substring matching, and future conditional-execution safety needs more than a leaf name.

Three planned PRs need to add policy for the same multimodal paths:

• #2763 initially adds distributed.frozen_multimodal_sharding to choose frozen-module FSDP ownership.
• #2990 adds vision frame sharding over CP as distributed.cp_vision_frame_sharding.
• A safety follow-up will add a dummy-forward policy for rank-asymmetric mixed-modality batches, as run_dummy_forward.

Rather than let each PR introduce an independent distributed config path, all three should extend
distributed.multimodal. They govern different behaviors, but each controls how resolved
multimodal modules interact with the existing parallel axes of the text model.

Finally, freeze_config is unvalidated. In the current VLM/audio/long-context example scope, 65
configs set freeze_config.freeze_embeddings, but the VLM freezing path never reads that key.

Proposal

1. One multimodal module resolver

@dataclass(frozen=True)
class MultimodalModuleSpec:
    path: str                                        # FQN relative to model root
    modality: str = "vision"                         # conventional but open
    role: str | None = None                          # conventional but open
    conditionally_executed: bool = True              # forward may be skipped

modality should remain an open string: a structure or graph encoder should be addressable without
being mislabeled as vision. The label is mainly a config key; frozen state comes from
requires_grad, and conditional execution comes from the declaration.

role should land with the resolver even though role-based freezing lands later. Existing and
planned consumers already distinguish towers, projectors, embedders, and resamplers. Keep it open
rather than using a closed enum, and use None when a compatibility fallback cannot identify the
role safely instead of silently defaulting every module to tower.

Keep generic config decoding outside the resolver and use a typed contract:

def resolve_multimodal_modules(
    model: nn.Module,
    *,
    declared_specs: Sequence[MultimodalModuleSpec] = (),
    override_specs: Sequence[MultimodalModuleSpec] = (),
) -> tuple[ResolvedMultimodalModule, ...]:
    ...

declared_specs is the canonical inventory chosen from, in order:

1. a model-class declaration (the normal path for models we own);
2. a central per-architecture table for plain Transformers classes;
3. the current leaf-name heuristic as a compatibility fallback, with a warning.

Optional override_specs take precedence for the same canonical path. A
ResolvedMultimodalModule pairs the stable spec with its live module object. This keeps the
resolver independent of an ambiguous generic config dependency while making declarations and
overrides explicit.

The table is still required because models such as Gemma3, Qwen3-VL, Llava, and SmolVLM are plain
Transformers classes and cannot declare AutoModel-specific metadata. The heuristic keeps unknown
third-party VLMs working without making the fallback silent.

Separate canonical declarations from live module binding. The model declaration, architecture
table, and optional config overrides produce stable specs; binding those paths to module objects
happens after structural transformations such as PEFT for distributed consumers. If freezing must
bind earlier, rebind after PEFT rather than retaining stale module objects. Evaluate requires_grad
from the live parameters after freezing/PEFT and immediately before sharding; never cache frozen
state in the declaration.

Resolver contracts should include canonical unwrapped FQNs, identity deduplication for aliases,
and the ability for explicit overrides to correct metadata. Parent/child normalization must retain
a child with a distinct policy-relevant role rather than discarding it merely because its parent
was selected.

Activation-checkpointing scope filtering, both FSDP parallelizers, FSDP unit enumeration, and
future conditional-path safety should consume this resolver. activation_checkpointing_scope: multimodal should then mean every resolved non-language group rather than explicitly
“vision + audio.”

2. Group multimodal distributed policies

distributed.multimodal should hold policies for applying the text model's existing parallel axes
to resolved multimodal modules. It should not create a separate vision/audio mesh; that is topology
and belongs beside pipeline_config or moe_parallel_config.

knob	scope	axis borrowed	what is distributed
frozen_sharding	fully frozen resolved modules	DP / FSDP	parameters
frame_sharding	per modality	CP, or CP × TP	encoder compute
run_dummy_forward	per modality	DP	collective alignment

Target config shape:

# Optional override for third-party models or incorrect declarations.
# Normally absent because an AutoModel-owned model declares these specs itself.
multimodal:
  modules:
    - path: perceiver_resampler
      modality: vision
      role: resampler
      conditionally_executed: true

distributed:
  multimodal:
    frozen_sharding: root                    # root | per_layer | replicate; default for ALL
    validate_rank_uniform_execution: false   # false | <n diagnostic steps> | always
    vision:
      frame_sharding:                        # was distributed.cp_vision_frame_sharding
        enabled: true                        # default false
        mesh_dims: [cp]                      # only supported value initially
        min_tokens: 2048                     # below this, stay replicated
        cost_alpha: auto                     # int | auto | null
      run_dummy_forward: auto                # auto | true | false
    audio:
      frozen_sharding: per_layer             # per-modality override of the top-level default
      run_dummy_forward: auto

frozen_sharding is a top-level default because “fully frozen” is not modality-specific. Its
scope is nevertheless limited to modules returned by the resolver, so it cannot silently capture
a frozen language model. Per-modality sections override that default.

run_dummy_forward is per modality because only the model knows how to construct a valid dummy
input and execute that modality-specific path. The optional top-level multimodal.modules list is
an identity-metadata escape hatch; model-class declaration is the normal mechanism and travels with
the model for pip users. It does not belong under distributed because activation checkpointing
and freezing also consume the same resolved inventory and must not read another component's policy
config.

frame_sharding.mesh_dims is a serialized selection so the configuration remains self-describing
now that the policy name does not encode CP. It defaults to [cp], and [cp] is the only
supported value initially. Runtime code must resolve the selected dimensions from MeshContext
and reject dimensions that are absent or unsupported; it must not accept and ignore the field.
Extend the accepted values to [cp, tp] only when the caller, ownership and gradient constraints,
and the corresponding test matrix are implemented. The model capability remains separate: it
declares whether the encoder supports independent per-frame computation, while mesh_dims
declares where that computation is sharded.

These knobs need shared validation. Frame sharding can create rank asymmetry when a rank receives
no frames; frozen sharding and frame sharding partition the same module; and a conditionally
executed per-layer FSDP unit is unsafe unless execution is rank-uniform.

Why run_dummy_forward is needed

A trainable tower or projector may be skipped on a text-only batch. If it owns an FSDP unit, ranks
that skip it do not enter its collectives and the group can hang. Keeping it in the root avoids
that unit-level mismatch, but leaves no gradient for its parameters on the skipping ranks.

Root ownership plus set_reduce_scatter_unused_params(True) was prototyped and deadlocks with
nested FSDP groups: text-only ranks schedule the conditional root reduce-scatter first, while image
ranks schedule a nested language-layer reduce-scatter first on the same communicator.

The working remedy is to run the conditional path on a dummy input and add a differentiable
zero-weighted result. Every rank then traverses the same trainable modules without changing the
loss. This pattern already exists in llama_nemotron_vl; retrieval forwards a per-call
run_dummy_vision signal only to models that accept it.

The mechanism must remain model-owned: only the model knows its cheapest valid dummy input and how
to connect the tower plus any trainable projector/embedder/resampler to the graph. What this
proposal adds is reusable policy, so safety is not hardcoded in individual recipes.

run_dummy_forward: auto derives whether dummy execution is required from the distributed
strategy, final post-freezing/PEFT requires_grad, conditionally_executed metadata, and the
rank-uniform modality-execution declaration from the batch planner. The model-owned dummy-forward
capability determines whether the required safety action can be performed.

If those facts require dummy execution and the model supports it, enable it. If dummy execution
is required but the model lacks that capability, fail setup rather than silently proceeding
unsafely. When the planner guarantees rank-uniform modality execution, auto may leave the dummy
path disabled. Explicit overrides remain subject to the same correctness validation.

3. Validate component contracts; do not derive configs from each other

Do not make frozen_sharding: auto read dataset config. Identical distributed config should not
change meaning depending on the paired dataset.

Instead, each component declares its own fact:

• the resolver identifies conditionally executed modules;
• the data/batch planner declares whether modality execution may differ across ranks;
• distributed config selects sharding and dummy-forward policy.

Setup validation can then reject per_layer with potentially rank-asymmetric execution, a
trainable conditional FSDP unit with dummy safety disabled, or unsupported frame-sharding axes.

An optional runtime validator can check planned modality execution for the first N microbatches or
continuously. A bounded check is diagnostic only: uniformity during warmup cannot guarantee a later
batch is uniform and therefore cannot make an otherwise unsafe per_layer policy correct.
Correctness must come from a rank-uniform batching contract, dummy-forward safety, or an always-on
precheck. The collective must run at an always-executed model/recipe boundary before any
conditional tower; placing it inside a tower hook would deadlock because skipping ranks never enter
the hook.

4. Migrate freeze_config last, in a separate PR

Freezing currently uses substrings, so intent depends on naming: multi_modal_projector remains
trainable while Gemma4's embed_vision is frozen. Reuse the resolver's (modality, role)
vocabulary:

freeze_config:
  vision: {tower: true, projector: false}
  audio:  {tower: true, projector: true}
  language_model: false
  modules: [patch_generator.video_embedder]   # escape hatch

This is behavior-changing because today's substring overreach is load-bearing for some configs.
First snapshot the frozen-parameter FQN set for every tracked example, then preserve or explicitly
document each change.

Strict validation should reject the ignored VLM freeze_embeddings key rather than begin honoring
it: those 65 configs have historically trained with the key ignored.

Suggested sequencing

1. Decide whether #2990 moves into the nested namespace before merge or gets an explicit migration.
2. Keep #2763 focused on the immediate FSDP policy and shared FSDP taxonomy.
3. Land the resolver, including open role metadata, and migrate discovery consumers without
   changing freezing behavior.
4. Add the model-owned dummy-forward protocol, policy, and setup/runtime validation.
5. Migrate freeze_config separately, gated by frozen-parameter snapshots.

The YAML above is the target shape; the fields do not need to land together.

Open questions

1. Should the collator/batch planner own the planned modality-execution vector, with the dataset
   declaration used as a conservative setup-time fact?
2. Should conditionally_executed be model-declared or inferred from the forward contract?
3. Should diagnostic runtime uniformity validation default to a bounded warmup window or remain
   opt-in?
4. What is the smallest model-owned dummy protocol covering the encoder and all trainable
   projector/embedder/resampler modules?
5. Should dummy input use the full image size or a model-derived minimum?
6. Should config policy and per-call signal share a name (run_dummy_forward vs.
   run_dummy_vision)?

Already resolved in #2763

FSDP unit creation and enumeration now consume the same shared taxonomy. Narrowing creation to the
old two-name enumeration was rejected: with wrap_outer_model=False, a trainable multimodal module
on the outer model could otherwise belong to no FSDP unit and lose gradient reduction across DP
ranks.