[models,precision,examples] fix: trainable-only checkpoint load + FSDP2 support for bundled models - #202
Merged
Merged
Conversation
_load_lora and _load_full_model iterated model_args.target_components, but save_checkpoint only writes components that carry target modules. A bundled but frozen member (e.g. Wan2.2 transformer_2, kept in target_components solely to be FSDP-sharded for memory) is skipped on save, so resume probed a per-component subdir that was never written and logged a spurious error. Iterate trainable_component_names on load so it mirrors save.
ModelBundle is the single root passed to accelerator.prepare but did not expose `_no_split_modules`, so accelerate's TRANSFORMER_BASED_WRAP resolved an empty transformer_layer_cls set and wrapped the whole bundle as ONE FSDP unit -> a single monolithic flat param (the full unsharded model materialized on every rank) -> OOM at init for large models (e.g. Wan2.2 A14B ~53GB). Aggregate the members' declared block classes (e.g. WanTransformerBlock) so FSDP shards per-block. Returns None when no member declares any, preserving accelerate's fallback. No effect on DDP/DeepSpeed.
…ype, add frozen_parameters_dtype, support FSDP2 for bundled models Dtype config (concept fix): - Rename `master_weight_dtype` -> `trainable_parameters_dtype` (despite the name it only ever set the TRAINABLE param dtype). Add `frozen_parameters_dtype` (default None -> model inference dtype, preserving prior behavior). `_mix_precision` is the single place that decides every parameter's original dtype. FSDP2 support (Wan2.2 A14B bundled dual-transformer): - Uniform original dtype: FSDP2 shards each unit with ONE dtype, and accelerate upcasts trainable params to an fp32 master (mixed_precision != 'no'). So a trained component that bundles frozen members (Wan2.2 trains `transformer`, `transformer_2` frozen-but-sharded) would mix fp32/bf16 in a unit and trip FSDP2's uniform-dtype assert. `_mix_precision` now forces the trained (sharded) components to a uniform fp32 original dtype (compute stays low-precision via MixedPrecisionPolicy); untrained components (text encoder / VAE) keep `frozen_parameters_dtype`. Detection uses `_is_fsdp2()` (accelerator.is_fsdp2). - Non-empty checkpoint save: get_state_dict's `state_dict_keys` path toggled `requires_grad` to sub-select params, but FSDP2 keys ignore_frozen_params off the fully_shard-time trainability -> the toggle was a no-op and produced an EMPTY adapter. Gather via `get_model_state_dict` and let the shared key-filter narrow to LoRA. Verified on Wan2.2-A14B / FSDP2 / 8xGPU: non-empty adapter + bit-exact save/load closure (max_abs_diff=0). Configs/docs: z-image examples set trainable+frozen dtype fp32 (needs fp32 for numerical stability); all example configs + knowledge docs updated to the new field name.
Wan2.2 routes each timestep to only one transformer (boundary_ratio); the bundled two-transformer root + per-block wrapping trips FSDP1's single-root lazy-init (`_is_root` assertion) on the first forward. FSDP2 (config/accelerate_configs/ fsdp2.yaml) shards the bundle correctly. Annotate the wan22 example configs.
frozen_parameters_dtype=None previously fell back to the inference dtype, force-casting every frozen component to a single dtype and silently downcasting params that from_pretrained loaded in higher precision. Released checkpoints deliberately ship components in different dtypes (e.g. Z-Image: transformer fp32, text encoder bf16), so the uniform cast overrode those choices. None now preserves each frozen param/buffer's loaded dtype (no downcast); an explicit frozen_parameters_dtype opts into forcing a cast. Drop the now-unneeded frozen_parameters_dtype: fp32 from the Z-Image configs.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR refines Flow-Factory’s post-ModelBundle training stack by (1) making checkpoint load behavior symmetric with save for bundled-but-frozen members, (2) improving FSDP2 behavior for bundled roots (auto-wrap discovery + correct state-dict gathering), and (3) clarifying/expanding dtype configuration so frozen components can preserve their loaded precision by default.
Changes:
- Make LoRA/full checkpoint loading iterate only trainable components (
trainable_component_names) to matchsave_checkpointoutput. - Add bundled-root FSDP2 support by aggregating
_no_split_modulesacross bundle members and fixing FSDP2 state-dict gathering behavior. - Rename
master_weight_dtype→trainable_parameters_dtypeand addfrozen_parameters_dtypewith default “preserve loaded dtype”; update examples and internal docs/constraints accordingly.
Reviewed changes
Copilot reviewed 18 out of 18 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| src/flow_factory/trainers/abc.py | Clarifies that parameter dtype decisions are handled in adapter _mix_precision() before preparing the bundled root. |
| src/flow_factory/models/model_bundle.py | Adds _no_split_modules aggregation on ModelBundle to enable per-block FSDP2 wrapping via accelerate auto-wrap. |
| src/flow_factory/models/abc.py | Updates mixed-precision casting semantics, fixes FSDP2 state-dict gathering behavior, and makes checkpoint loading symmetric with saving (trainable-only). |
| src/flow_factory/hparams/model_args.py | Renames the trainable dtype knob and adds frozen_parameters_dtype with None meaning “preserve loaded dtype”. |
| examples/nft/lora/z_image/default.yaml | Migrates config key to trainable_parameters_dtype and documents frozen-dtype preservation. |
| examples/nft/full/z_image/default.yaml | Migrates config key to trainable_parameters_dtype and documents frozen-dtype preservation. |
| examples/nft/full/z_image_turbo/default.yaml | Migrates config key to trainable_parameters_dtype and documents frozen-dtype preservation. |
| examples/grpo/lora/z_image/default.yaml | Migrates config key to trainable_parameters_dtype and documents frozen-dtype preservation. |
| examples/grpo/lora/z_image_turbo/default.yaml | Migrates config key to trainable_parameters_dtype and documents frozen-dtype preservation. |
| examples/grpo/lora/wan22/t2v.yaml | Updates guidance to prefer FSDP2 over FSDP1 for bundled Wan2.2 dual-transformer layouts. |
| examples/grpo/lora/wan22/i2v.yaml | Updates guidance to prefer FSDP2 over FSDP1 for bundled Wan2.2 dual-transformer layouts. |
| examples/grpo/full/z_image/default.yaml | Migrates config key to trainable_parameters_dtype and documents frozen-dtype preservation. |
| examples/grpo/full/z_image_turbo/default.yaml | Migrates config key to trainable_parameters_dtype and documents frozen-dtype preservation. |
| examples/grpo/full/wan22/t2v.yaml | Updates guidance to prefer FSDP2 over FSDP1 for bundled Wan2.2 dual-transformer layouts. |
| examples/grpo/full/wan22/i2v.yaml | Updates guidance to prefer FSDP2 over FSDP1 for bundled Wan2.2 dual-transformer layouts. |
| .agents/knowledge/topics/dtype_precision.md | Updates precision documentation to reflect new trainable/frozen dtype controls. |
| .agents/knowledge/topics/autocast_param_swap.md | Updates references to the renamed trainable dtype config key. |
| .agents/knowledge/constraints.md | Adds explicit bundle checkpoint save/load symmetry constraint and updates autocast constraint wording for renamed key. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| frozen_dtype = self.model_args.frozen_parameters_dtype # None -> preserve loaded dtype | ||
|
|
||
| # Get target components and all component names | ||
| target_set = frozenset(self.model_args.target_components) |
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.
Summary
Three related fixes on top of the single-root
ModelBundlerefactor (#190), all validated against the Wan2.2 dual-transformer "train one, shard the other" layout and Z-Image.1. Checkpoint load/save symmetry (
c1d2ed4)_load_lora/_load_full_modeliteratedmodel_args.target_components, butsave_checkpointonly writes components that carry target modules. A bundled-but-frozen member (e.g. Wan2.2transformer_2, kept intarget_componentssolely to be FSDP-sharded for memory) is skipped on save, so resume probed a per-component subdir that was never written and logged a spurious error. Load now iteratestrainable_component_namesto mirror save.2. FSDP2 support for the bundled root
_no_split_modulesonModelBundle(46f8c76): accelerate'sTRANSFORMER_BASED_WRAPreads_no_split_modulesoff the single prepared root; without it the whole bundle was wrapped as ONE FSDP unit (a monolithic flat param materialized on every rank → OOM at init for A14B). Now aggregates members' declared block classes so FSDP shards per-block. No effect on DDP/DeepSpeed.6d95785): renamemaster_weight_dtype→trainable_parameters_dtype(it only ever set the trainable dtype); addfrozen_parameters_dtype; force trained components to a uniform fp32 orig dtype under FSDP2 (compute stays bf16 viaMixedPrecisionPolicy); fixget_state_dictproducing an EMPTY adapter on FSDP2 (therequires_gradtoggle was a no-op). Verified bit-exact save/load on Wan2.2-A14B (FSDP2, 8×GPU).38ca25e): recommend FSDP2 over FSDP1 for Wan2.2.3. Frozen params preserve their loaded dtype (
5cb276a)frozen_parameters_dtype=None(default) now preserves each frozen component'sfrom_pretraineddtype instead of force-casting all frozen params to the inference dtype. Released checkpoints deliberately ship components at different precisions (Z-Image: transformer fp32, text encoder bf16); the old uniform cast silently downcast them. An explicitfrozen_parameters_dtypeopts into forcing a cast.Migration notes
master_weight_dtype→trainable_parameters_dtype. Unknown keys are shunted toextra_kwargswith only a warning, so stale configs silently fall back to thebfloat16default — update your YAMLs.mixed_precision: bf16are unchanged; setfrozen_parameters_dtype: bf16to restore the previous downcast where memory matters.Test plan
Nonepreserves per-component (fp32/bf16/fp16), explicit dtype forces, uniform fast-path.max_abs_diff=0)..cursor/plans/gpu_checkpoint_roundtrip_test.plan.md(SD3.5 + Wan2.2 train-one/shard-one + train-both).Notes
black/isortdrift inabc.pywas left untouched; changed lines follow the file's existing style.