feat(fsdp): per-module weight-precision control, and Qwen3.5 compute-dtype alignment

[Rockdu]

What

Ports miles_diffusion#91's per-parameter weight-precision control onto the FSDP backend, and uses it to close the Qwen3.5 train/rollout compute-dtype gap.

The FSDP backend already had a PrecisionPolicy with run-level param_dtype / reduce_dtype / autocast_dtype and a per-arch hook registry. This adds the missing axis: per-module gather dtype. A model declares PrecisionSpec rules (FQN glob and/or class-name glob), and a compiler lowers them onto nested fully_shard groups — one wrap unit per module whose gather dtype differs from its parent's. Everything stays inside FSDP (DTensor params, FSDP grad reduction, DCP/offload unchanged).

Three dtype axes, each with one owner:

axis	meaning	owner
master	resident dtype the optimizer steps on	--disable-fp32-master (run-level)
gather	all-gather dtype, i.e. the dtype weights carry in the forward	param_dtype, refined per module by PrecisionSpec
compute	dtype the ops run at	autocast_dtype + arch class patches

miles_diffusion's input_dtype_policy is deliberately not ported: an LLM's model-boundary inputs are integer input_ids/position_ids, so there is no float boundary cast to declare.

The Qwen3.5 finding

Dumped every nn.Module's IO dtype for Qwen3.5-4B under FSDP2 (non-intrusive Dumper + the dumper-module-report skill), in the precision-aligned policy that Qwen3 dense's true-on-policy contract uses — fp32 master, fp32 gather, bf16 autocast.

527 of 1074 dumped tensors were fp32, from two compounding causes:

1. MixedPrecisionPolicy defaults to cast_forward_inputs=True, so every wrapped decoder layer re-cast its activations up to the fp32 gather dtype at each layer boundary — while autocast ran the matmuls at bf16.
2. F.embedding is not an autocast-covered op, so the embedding output carried the gathered weight dtype and seeded the residual stream. Qwen3_5RMSNorm ends in output.type_as(x) and every residual add promotes, so that one dtype propagated through the whole activation path.

Both are fixed here, and both are load-bearing — fixing either alone leaves the stream fp32. Result: 0 of 1074 fp32; all 1072 float tensors bf16, fp32 confined to master weights, int64 confined to input_ids/position_ids.

Changes

• adaptations/precision.py: ModuleSel / Rule / PrecisionSpec, the compile_precision lowering (whole-tree coverage, minimal wrap-unit cover, zero-hit rules rejected), wrap_plan deepest-first ordering, and parse_precision_rules for the --fsdp-precision-rules operator escape hatch
• apply_fsdp2: wraps per compiled unit at that unit's param_dtype; new cast_forward_inputs parameter, which the actor sets to False exactly when a forward autocast is active
• adaptations/specs/qwen3_5_moe.py: registers the embedding gather rule for GatedDeltaNet archs (Qwen3.5/3.6/Qwen3-Next). Self-activating — a no-op unless compute and gather dtypes disagree
• debug_dump.py + tools/dump_fsdp_module_dtypes.py: guarded Dumper registration in the train step, and a harness that drives the actor's construction path on synthetic tokens to produce the dtype inventory

Default behavior is unchanged: with no autocast the policy resolves to 0 spec rules and cast_forward_inputs=True, and a default-policy Qwen3.5 dump is byte-for-byte the same dtype census as before this PR.

Tests

• tests/fast/backends/test_fsdp_precision_plan.py — 13 CPU tests for the compiler, each docstring drawing the resulting gather dtype per node
• tests/fast/backends/test_fsdp_precision_wrap.py — nested wrap units under real FSDP2 on 2 gloo ranks, asserting the param dtype each module sees at forward and that master/grad stay fp32
• Both registered to stage-a-cpu; full tests/fast/backends/ suite passes (400 tests)

Pending

• The Qwen3.5 alignment was validated on the dtype axis only. Numerical parity against an sglang Qwen3.5 rollout still needs a comparator run, which is what the dump harness is meant to feed next.
• No Qwen3.5 true-on-policy contract exists yet, so the fp32-gather + bf16-autocast policy is reachable only via the harness overrides. Landing a contract is separate work.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[Rockdu]

Closing — opened against the upstream repo by mistake. Work continues on the fork branch.