Skip to content

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

Merged
Rockdu merged 13 commits into
radixark:mainfrom
Rockdu:refactor/fsdp-5-precision-plans
Aug 8, 2026
Merged

refactor(fsdp)(5/5): compile per-model precision plans onto the dtype patch#106
Rockdu merged 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.
  • 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
@Rockdu
Rockdu merged commit 48fbd2c into radixark:main Aug 8, 2026
23 of 28 checks passed
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