Keep every FSDP2 group dtype-uniform
FSDP2 all-gathers a group into a single flat buffer, so it requires one
original dtype per group. A model that keeps some parameters in float32
(`_keep_in_fp32_modules_strict`, e.g. ERNIE-4.5-MoE's router) trips that
check at the first forward when it is sharded at load time:
AssertionError: FSDP expects uniform original parameter dtype but got
{torch.bfloat16, torch.float32}
The model is built uniformly in the requested dtype and the loader applies
the dtype plan only when the weights land, which is after sharding, so the
FSDP plan was being built against dtypes the model would not actually have.
LoRA is unaffected because frozen parameters are left out of the check.
Apply the loader's dtype plan to the (meta) parameters before resolving the
plans, which is free, then split any group that holds more than one dtype:
it keeps its dominant dtype and each submodule owning another dtype becomes
a group of its own, listed ahead of its parent as FSDP2 requires. Stating
the rule as the invariant FSDP2 enforces, rather than as a special case for
one flag, means groups are only split when they would actually have failed,
and any future source of dtype heterogeneity is covered.
Two consequences for the groups this carves out:
- mixed precision would all-gather them in `param_dtype` and silently undo
the upcast the flag asked for, so they are gathered in their own dtype;
- a frozen parameter whose first dimension is smaller than the mesh cannot
be sharded usefully (FSDP2 pads, so every rank stores a full row either
way) and has no gradient to reduce, so it is left unsharded. The root
group ignores the same parameters, otherwise it would pick them back up.
Tests: `tests/distributed/test_fsdp_plan.py` covers the grouping on CPU
against a meta model, since the distributed FSDP tests need several
accelerators. ERNIE-4.5-MoE joins those, and `test_fsdp2_bf16_training`
trains a step with the model sharded at load time.
