fix(fsdp): resolve fp32 master-weight compute dtype per parameter (#3328)

* fix(fsdp): resolve fp32 master-weight compute dtype per parameter

Custom MoE models (e.g. Qwen3.5/3.6-MoE) never record `_hf_compute_dtype`, so
every parameter falls to the compute-dtype fallback in `_make_compute_dtype_fn`.
The fallback keyed on a module-global `storage_is_uniform` flag, which a single
native-bf16 sibling (`shared_expert_gate`) flips false -- misclassifying the
layer's fp32 master weights as fp32 compute. With #3255's
`cast_forward_inputs=False`, the misclassified fp32 unit no longer casts its
bf16 input and the forward crashes (`bf16 @ fp32`).

Resolve the fallback per parameter: unhinted fp32 storage under a lower-precision
policy is an fp32 master weight and computes in the policy dtype; any other
storage keeps its own dtype. A single bf16 sibling no longer forces the layer's
fp32 master weights into fp32 compute. The `(storage, compute)` grouping still
keeps each FSDP unit storage-uniform, so `_init_mp_dtypes` is unaffected.

Fixes #3327

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: NancyFyong <88076188+NancyFyong@users.noreply.github.com>

* fix(qwen3_5_moe): pin _fp32_params compute, build shared_expert_gate in dtype

Follow-up to the per-parameter compute-dtype fallback (#3328). Two gaps that
the resolver change exposes on Qwen3.5-MoE:

- The MoE parallelizer keys the fp32 compute pin on
  `_keep_in_fp32_modules_strict`, which the model never set, so the SSM-gating
  A_log/dt_bias fell to the policy dtype (bf16) instead of fp32. Add the strict
  entry, matching the inkling/kimi_k3 siblings.
- `shared_expert_gate` was built at the `initialize_linear_module` default
  (bf16) even under fp32 master-weight storage, making it the lone bf16 sibling
  the #3327 fallback has to tolerate. Build it in `config.dtype`.

Add a compute-contract regression test for the mixed-storage, no-HF-hint layout
(fp32 masters + a genuinely-bf16 shared_expert_gate + the fp32 holder): ordinary
weights and the gate compute bf16; only the pinned A_log/dt_bias stay fp32.

Fixes #3327

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: NancyFyong <88076188+NancyFyong@users.noreply.github.com>

* fix(qwen3_next): pin _fp32_params fp32 compute via strict contract

Qwen3-Next uses the same intrinsic-fp32 `_fp32_params` holder but declared it
only in `_keep_in_fp32_modules`. The MoE parallelizer keys the fp32 compute pin
on `_keep_in_fp32_modules_strict`, so under the generalized fallback its
A_log/dt_bias would compute in bf16. Add the strict entry, matching Qwen3.5-MoE.

Add a model-owned assertion guarding the strict declaration on every model with
the holder, so a future model that forgets the contract fails a test rather than
silently dropping to bf16.

Fixes #3327

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: NancyFyong <88076188+NancyFyong@users.noreply.github.com>

---------

Signed-off-by: NancyFyong <88076188+NancyFyong@users.noreply.github.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>