[FSDP2] Add additional_modules_to_shard for sharding custom modules

[yuxinyuan]

What does this PR do?

Adds an opt-in additional_modules_to_shard option to FullyShardedDataParallelPlugin (FSDP2) that lets users explicitly shard extra modules into their own FSDP2 units, on top of what auto_wrap_policy selects.

Motivation

FSDP2 preparation shards the modules matched by auto_wrap_policy, then uses a transformers-specific heuristic to carve the input/output embeddings and the final norm into their own units. Any module that is neither matched by the policy nor discovered by that heuristic falls into the root FSDP2 unit, which is not resharded after forward.

This is common for composite / multimodal / omni models that stitch together several transformer stacks. Take a Qwen3-TTS-style model as a concrete example: it contains a speaker encoder, a standard LM decoder, and a code predictor (another transformer). The top-level composite also doesn't expose a single get_input_embeddings/get_output_embeddings/final-norm through the standard transformers interface, so the built-in heuristic can't carve those out.

additional_modules_to_shard lets users close that gap explicitly — shard the non-standard modules into their own units, and optionally give specific submodules different sharding behavior (e.g. keep a precision-sensitive module in fp32, or skip resharding a custom tail)

API

additional_modules_to_shard is a selector -> overrides mapping:

• selector (key): an exact fully-qualified module name or an fnmatch glob (e.g. "model.norm", "model.layers.*.mlp"), or a tuple of them. A str shards each matched module as its own unit; a tuple groups the referenced modules into a single unit (like the built-in norm+lm_head tail). Names are canonicalized before matching, so selectors written against the original module names still match after torch.compile / activation checkpointing.
• overrides (value): a dict overriding the per-module sharding kwargs — reshard_after_forward and/or mp_policy; unset keys inherit the model-wide values (offload_policy is intentionally not supported).

import torch
from torch.distributed.fsdp import MixedPrecisionPolicy
from accelerate import Accelerator, FullyShardedDataParallelPlugin

plugin = FullyShardedDataParallelPlugin(
    fsdp_version=2,
    additional_modules_to_shard={
        # keep the tied final norm + lm_head in one unit and skip resharding after forward
        ("model.norm", "lm_head"): {"reshard_after_forward": False},
        # keep a precision-sensitive module in fp32
        "vision_tower": {"mp_policy": MixedPrecisionPolicy(param_dtype=torch.float32)},
        # shard every per-layer custom mlp with the model-wide defaults
        "model.layers.*.custom_mlp": {},
    },
)
accelerator = Accelerator(fsdp_plugin=plugin)

Tests

Adds CPU unit tests mirroring the existing mocked fsdp2_prepare_model tests (fully_shard spied), covering: single/glob/grouped selection, per-module kwarg overrides, precedence over the auto wrap policy, name canonicalization for torch.compile/activation checkpointing, group-vs-traversal ordering, config validation, the FSDP1 no-op, and the setter.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR. @SunMarc