Skip to content

refactor(fsdp)(5/5): compile per-model precision plans onto the dtype patch - #106

Open
Rockdu wants to merge 13 commits into
radixark:mainfrom
Rockdu:refactor/fsdp-5-precision-plans
Open

refactor(fsdp)(5/5): compile per-model precision plans onto the dtype patch#106
Rockdu wants to merge 13 commits into
radixark:mainfrom
Rockdu:refactor/fsdp-5-precision-plans

Conversation

@Rockdu

@Rockdu Rockdu commented Aug 8, 2026

Copy link
Copy Markdown
Collaborator

What

The declaration-to-mechanism compiler for the per-parameter dtype patch that landed in #100, plus the first per-model declarations:

  • mixed_precision.compile_param_dtype_maps: expands root-relative FQN glob patterns ({"*.norm2.*": "fp32"}) against the model, then re-keys each matched parameter to the FQN local to its owning fully_shard wrap — child wraps get module-relative maps, root-owned parameters stay root-relative (diagram in the module docstring). Zero-match patterns and doubly-matched parameters are rejected at compile time.
  • FSDPParallelPlan(no_split_modules, param_dtype_patterns): one per-model declaration consumed by apply_fsdp2. Model backends resolve it — native packages declare FSDP_PARALLEL_PLAN (LTX migrated from FSDP_NO_SPLIT_MODULES), diffusers families get a models/diffusers/<family>/parallel_plan.py dispatched by the family key, falling back to the model's _no_split_modules.
  • Wan2.2 declares scale_shift_table / time_embedder / norm2 → fp32 — diffusers' own _keep_in_fp32_modules minus the paramless entries (rope, norm1, norm3 have no weights). scale_shift_table is a bare nn.Parameter, which module-level wrapping could not express; the per-parameter patch covers it. SD3 / Qwen-Image declare empty plans.
  • apply_fsdp2 compiles the patterns and applies the FSDP2 patch only when at least one override survives (a pattern equal to the run dtype compiles to nothing), so plan-less runs take the unpatched path bit-for-bit.

Why

#100 merged the mechanism (ParamDtypeMixedPrecisionPolicy + patch) without a producer: nothing in-tree built a param_dtype_map. This PR is the missing translation layer — families state intent as root FQN patterns; the compiler owns the root-vs-child FQN bookkeeping that the raw policy API pushes onto its caller.

Files

  • miles/backends/fsdp_utils/mixed_precision.py — the compiler
  • miles/backends/fsdp_utils/models/parallel_plan.pyFSDPParallelPlan
  • miles/backends/fsdp_utils/model_backend.pyfsdp_parallel_plan() resolution per backend
  • miles/backends/fsdp_utils/models/diffusers/{wan2_2,sd3,qwen_image}/parallel_plan.py — family declarations
  • miles/backends/fsdp_utils/models/ltx/parallel_plan.pyFSDP_NO_SPLIT_MODULESFSDP_PARALLEL_PLAN
  • miles/backends/fsdp_utils/actor.py — wire the plan into apply_fsdp2
  • miles/backends/fsdp_utils/configs/train_pipeline_config.py — registry stamps model_family for plan dispatch
  • tests/fast/backends/fsdp_utils/test_mixed_precision.py — compiler unit tests (stage-a-cpu)
  • tests/fast-gpu/backends/fsdp_utils/_param_dtype_map_integration_worker.py — end-to-end compile→wrap→forward on GPU
  • tests/fast/backends/fsdp_utils/{test_model_backend,models/ltx/test_ltx}.py — plan resolution coverage

Pending GPU validation

  • Wan2.2 numerics change by design (fp32 gathers for the declared parameters); the 5-GPU Wan recipe has no recorded e2e standard, so CI gates only the SD3/LTX standards, which this PR leaves bit-identical (empty plans).
  • tests/fast-gpu param-dtype-map suites on an H200 runner.

Checklist

  • pre-commit run --all-files passes
  • Added/updated tests for new behaviour
  • pytest -x is green — not run locally (no local torch); stage-a-cpu covers the new unit tests
  • If launch flags changed, python3 train.py --help still parses — no flags changed
  • If a public flag was added, it appears in the CLI reference docs — no flag added
  • If an example was added, it has a real walkthrough — no example added

@Rockdu
Rockdu force-pushed the refactor/fsdp-5-precision-plans branch from f5584d4 to 495b064 Compare August 8, 2026 08:18
@Rockdu
Rockdu force-pushed the refactor/fsdp-5-precision-plans branch from 495b064 to 704cef1 Compare August 8, 2026 08:47
@Rockdu Rockdu added the run-ci-fsdp Run FSDP domain tests on this PR label Aug 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci-fsdp Run FSDP domain tests on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant