refactor(fsdp)(5/5): compile per-model precision plans onto the dtype patch - #106
Open
Rockdu wants to merge 13 commits into
Open
refactor(fsdp)(5/5): compile per-model precision plans onto the dtype patch#106Rockdu wants to merge 13 commits into
Rockdu wants to merge 13 commits into
Conversation
Rockdu
force-pushed
the
refactor/fsdp-5-precision-plans
branch
from
August 8, 2026 08:18
f5584d4 to
495b064
Compare
Translate model-root FQN patterns into exact module-local policy maps and install the per-parameter patch only when apply_fsdp2 has effective overrides.
Apply the repository import ordering and Black layout to the compiler integration.
Rockdu
force-pushed
the
refactor/fsdp-5-precision-plans
branch
from
August 8, 2026 08:47
495b064 to
704cef1
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 owningfully_shardwrap — 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 byapply_fsdp2. Model backends resolve it — native packages declareFSDP_PARALLEL_PLAN(LTX migrated fromFSDP_NO_SPLIT_MODULES), diffusers families get amodels/diffusers/<family>/parallel_plan.pydispatched by the family key, falling back to the model's_no_split_modules.scale_shift_table/time_embedder/norm2→ fp32 — diffusers' own_keep_in_fp32_modulesminus the paramless entries (rope,norm1,norm3have no weights).scale_shift_tableis a barenn.Parameter, which module-level wrapping could not express; the per-parameter patch covers it. SD3 / Qwen-Image declare empty plans.apply_fsdp2compiles 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 aparam_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 compilermiles/backends/fsdp_utils/models/parallel_plan.py—FSDPParallelPlanmiles/backends/fsdp_utils/model_backend.py—fsdp_parallel_plan()resolution per backendmiles/backends/fsdp_utils/models/diffusers/{wan2_2,sd3,qwen_image}/parallel_plan.py— family declarationsmiles/backends/fsdp_utils/models/ltx/parallel_plan.py—FSDP_NO_SPLIT_MODULES→FSDP_PARALLEL_PLANmiles/backends/fsdp_utils/actor.py— wire the plan intoapply_fsdp2miles/backends/fsdp_utils/configs/train_pipeline_config.py— registry stampsmodel_familyfor plan dispatchtests/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 GPUtests/fast/backends/fsdp_utils/{test_model_backend,models/ltx/test_ltx}.py— plan resolution coveragePending GPU validation
tests/fast-gpuparam-dtype-map suites on an H200 runner.Checklist
pre-commit run --all-filespassespytest -xis green — not run locally (no local torch); stage-a-cpu covers the new unit testspython3 train.py --helpstill parses — no flags changed