[models,precision,examples] fix: trainable-only checkpoint load + FSDP2 support for bundled models

.agents/knowledge/constraints.md

@@ -46,6 +46,9 @@ All target components (trainable **and** frozen-but-shardable) are bundled into
 ### 9a. Sampler Geometric Constraints
 `DistributedKRepeatSampler` and `GroupContiguousSampler` require `M * K ≡ 0 (mod W * B * G)` where M=unique_sample_num, K=group_size, W=world_size, B=per_device_batch_size, G=gradient_step_per_epoch — **unless** `gradient_accumulation_steps` is set manually, in which case the constraint reduces to `M * K ≡ 0 (mod W * B)`. **GroupContiguousSampler** adds: `M ≡ 0 (mod W)`. **GroupDistributedSampler** (DGPO) requires: `K % W == 0` and `(W * B) % K == 0`; auto-aligned by `_align_for_group_distributed`. See `topics/samplers.md` for full details.
 
+### 9b. Checkpoint Save/Load Symmetry Under the Bundle
+Checkpoints are written and read for **trainable members only** — components whose `target_module_map[name]` is non-empty (`adapter.trainable_component_names`). Frozen-but-shardable bundle members (e.g. Wan2.2's `transformer_2`, kept in `target_components` only to be FSDP-sharded for memory; see #9) map to `None` and are skipped by both `save_checkpoint` and `_load_lora`/`_load_full_model`. Loaders MUST iterate `trainable_component_names`, not `target_components`, or resume logs a spurious error for a per-component subdir that was never written. `resume_type='state'` restores via `accelerator.load_state` into the prepared bundle root and is therefore keyed to bundle membership — resuming into a different `target_components` / bundle composition will mismatch.
+
 ### 10. DeepSpeed ZeRO-3 Is Unsupported
 Reward model sharding under ZeRO-3 is broken even with `GatherParameter` context manager (see the ZeRO-3 guard comment in `trainers/abc.py`). Only ZeRO-1 and ZeRO-2 are safe. Document this if users ask.
 
@@ -120,7 +123,7 @@ When using FSDP with CPU offloading, frozen components (text encoder, VAE) may b
 The adapter sets inference dtype for frozen components and training dtype for trainable parameters in `_mix_precision()`. Autocast context is configured in `BaseTrainer.__init__`. Do not manually cast tensors unless you understand the precision boundary. Details: `topics/dtype_precision.md`.
 
 ### 20a. Autocast Weight Cache Must Not Span a Forward
-`torch.autocast`'s weight cache (keyed by tensor `data_ptr`) serves **stale** casts after any in-place weight change — `optimizer.step()` or a `use_ref/ema/named_parameters` swap (`param.data.copy_`). So wrap **each** forward (and its KL) in its own `with self.autocast():`; never one autocast around the optimize loop. Active for fp32 trainable weights (`master_weight_dtype: fp32`), dormant for the bf16 default, LoRA `disable_adapter()` safe. Details + DDP/DeepSpeed caveat: `topics/autocast_param_swap.md`.
+`torch.autocast`'s weight cache (keyed by tensor `data_ptr`) serves **stale** casts after any in-place weight change — `optimizer.step()` or a `use_ref/ema/named_parameters` swap (`param.data.copy_`). So wrap **each** forward (and its KL) in its own `with self.autocast():`; never one autocast around the optimize loop. Active for fp32 trainable weights (`trainable_parameters_dtype: fp32`), dormant for the bf16 default, LoRA `disable_adapter()` safe. Details + DDP/DeepSpeed caveat: `topics/autocast_param_swap.md`.
 
 ---
 

.agents/knowledge/topics/autocast_param_swap.md

@@ -11,7 +11,7 @@
 - **Ref/EMA/named swap** — `copy_ema_to` does `param.data.copy_` (same `data_ptr`); a ref forward after the policy forward in one region reuses the policy cast → KL ≈ 0.
 - **`optimizer.step()`** — updates weights in place; if the region spans steps, later forwards reuse the pre-step cast → training frozen (loss flat).
 
-**Bites only** for fp32 trainable weights (`master_weight_dtype: fp32`); dormant for the bf16 default (nothing cached); LoRA's `disable_adapter()` ref path is safe (no `.data.copy_`).
+**Bites only** for fp32 trainable weights (`trainable_parameters_dtype: fp32`); dormant for the bf16 default (nothing cached); LoRA's `disable_adapter()` ref path is safe (no `.data.copy_`).
 
 **Rule**: wrap **each** forward (and its KL math) in its own `with self.autocast():`; never one autocast around the optimize loop. Per-forward (not `cache_enabled=False`) keeps the legit intra-forward reuse of two-pass-CFG adapters (e.g. `qwen_image_edit_plus.py` calls the transformer twice per forward). Precompute / sampling regions keep their outer autocast (weights constant).
 

.agents/knowledge/topics/dtype_precision.md

@@ -8,8 +8,8 @@
 
 | Component | Runtime dtype | Why |
 |-----------|--------------|-----|
-| Transformer (frozen) | `inference_dtype` (bf16/fp16) | Memory savings for frozen params |
-| Transformer (trainable) | `master_weight_dtype` (fp32/bf16) | Gradient precision |
+| Frozen params/buffers (frozen transformer base, VAE, text encoders) | `frozen_parameters_dtype` — default `None` **preserves each component's `from_pretrained` dtype** (no downcast) | Released checkpoints ship components in different dtypes (e.g. Z-Image: transformer fp32, text encoder bf16); set an explicit dtype to force one / save memory |
+| Transformer (trainable) | `trainable_parameters_dtype` (fp32/bf16) | Gradient precision |
 | Scheduler math | `float32` always | `1/sigma` amplification (see below) |
 | Latent storage (trajectory) | `latent_storage_dtype` (configurable) | Memory vs. precision tradeoff |
 | Advantage computation | `float64` (numpy) | Normalization stability |

examples/grpo/full/wan22/i2v.yaml

@@ -1,9 +1,11 @@
 # Environment Configuration
 launcher: "accelerate"  # Options: accelerate
 # Distributed sharding backend. For Wan2.2 you may set `target_components` to both `transformer`
-# and `transformer_2`; they are bundled into one prepared root, so DeepSpeed (ZeRO-1/2) and FSDP
-# (v1/v2) can all train and shard the two transformers together. FSDP full shard
-# (`config/accelerate_configs/fsdp_full_shard.yaml`) also shards model params; ZeRO-3 unsupported.
+# and `transformer_2`; they are bundled into one prepared root, so DeepSpeed (ZeRO-1/2) shards them
+# together. For FSDP prefer FSDP2 (`config/accelerate_configs/fsdp2.yaml`), NOT FSDP1
+# (`fsdp_full_shard.yaml`): Wan2.2 routes each timestep to only ONE transformer (boundary_ratio),
+# so the bundled two-transformer + per-block wrap trips FSDP1's single-root lazy-init (`_is_root`
+# assertion) on the first forward. ZeRO-3 unsupported.
 config_file: config/accelerate_configs/fsdp2.yaml
 num_processes: 8  # Number of processes to launch (overrides config file)
 main_process_port: 29500

examples/grpo/full/wan22/t2v.yaml

@@ -1,9 +1,11 @@
 # Environment Configuration
 launcher: "accelerate"  # Options: accelerate
 # Distributed sharding backend. For Wan2.2 you may set `target_components` to both `transformer`
-# and `transformer_2`; they are bundled into one prepared root, so DeepSpeed (ZeRO-1/2) and FSDP
-# (v1/v2) can all train and shard the two transformers together. FSDP full shard
-# (`config/accelerate_configs/fsdp_full_shard.yaml`) also shards model params; ZeRO-3 unsupported.
+# and `transformer_2`; they are bundled into one prepared root, so DeepSpeed (ZeRO-1/2) shards them
+# together. For FSDP prefer FSDP2 (`config/accelerate_configs/fsdp2.yaml`), NOT FSDP1
+# (`fsdp_full_shard.yaml`): Wan2.2 routes each timestep to only ONE transformer (boundary_ratio),
+# so the bundled two-transformer + per-block wrap trips FSDP1's single-root lazy-init (`_is_root`
+# assertion) on the first forward. ZeRO-3 unsupported.
 config_file: config/accelerate_configs/fsdp2.yaml
 num_processes: 8  # Number of processes to launch (overrides config file)
 main_process_port: 29500

examples/grpo/full/z_image/default.yaml

@@ -23,7 +23,10 @@ data:
 # Model Configuration
 model:
   finetune_type: 'full' # Options: full, lora
-  master_weight_dtype: "fp32" # Z-Image requires fp32 master weight for better performance
+  trainable_parameters_dtype: "fp32" # Z-Image trains the transformer in fp32 for numerical stability
+  # frozen_parameters_dtype unset -> preserve each frozen component's from_pretrained dtype (no downcast):
+  # Z-Image ships the transformer in fp32 but the text encoder in bf16, so forcing one frozen dtype
+  # would override that. Set an explicit dtype only to force a cast (e.g. bf16 to save memory).
   target_modules: "default" # Options: all, default, or list of module names like ["to_k", "to_q", "to_v", "to_out.0"]
   model_name_or_path: "Tongyi-MAI/Z-Image"  # Options: Tongyi-MAI/Z-Image, Tongyi-MAI/Z-Image-Turbo
   model_type: "z-image"

examples/grpo/full/z_image_turbo/default.yaml

@@ -23,7 +23,10 @@ data:
 # Model Configuration
 model:
   finetune_type: 'full' # Options: full, lora
-  master_weight_dtype: "fp32" # Z-Image-Turbo requires fp32 master weight for better performance
+  trainable_parameters_dtype: "fp32" # Z-Image-Turbo trains the transformer in fp32 for numerical stability
+  # frozen_parameters_dtype unset -> preserve each frozen component's from_pretrained dtype (no downcast):
+  # Z-Image-Turbo ships the transformer in fp32 but the text encoder in bf16, so forcing one frozen dtype
+  # would override that. Set an explicit dtype only to force a cast (e.g. bf16 to save memory).
   target_modules: "default" # Options: all, default, or list of module names like ["to_k", "to_q", "to_v", "to_out.0"]
   model_name_or_path: "Tongyi-MAI/Z-Image-Turbo"  # HuggingFace model ID or local path
   model_type: "z-image"

examples/grpo/lora/wan22/i2v.yaml

@@ -1,9 +1,11 @@
 # Environment Configuration
 launcher: "accelerate"  # Options: accelerate
 # Distributed sharding backend. For Wan2.2 you may set `target_components` to both `transformer`
-# and `transformer_2`; they are bundled into one prepared root, so DeepSpeed (ZeRO-1/2) and FSDP
-# (v1/v2) can all train and shard the two transformers together. FSDP full shard
-# (`config/accelerate_configs/fsdp_full_shard.yaml`) also shards model params; ZeRO-3 unsupported.
+# and `transformer_2`; they are bundled into one prepared root, so DeepSpeed (ZeRO-1/2) shards them
+# together. For FSDP prefer FSDP2 (`config/accelerate_configs/fsdp2.yaml`), NOT FSDP1
+# (`fsdp_full_shard.yaml`): Wan2.2 routes each timestep to only ONE transformer (boundary_ratio),
+# so the bundled two-transformer + per-block wrap trips FSDP1's single-root lazy-init (`_is_root`
+# assertion) on the first forward. ZeRO-3 unsupported.
 config_file: config/deepspeed/deepspeed_zero2.yaml
 num_processes: 8  # Number of processes to launch (overrides config file)
 main_process_port: 29500

examples/grpo/lora/wan22/t2v.yaml

@@ -1,9 +1,11 @@
 # Environment Configuration
 launcher: "accelerate"  # Options: accelerate
 # Distributed sharding backend. For Wan2.2 you may set `target_components` to both `transformer`
-# and `transformer_2`; they are bundled into one prepared root, so DeepSpeed (ZeRO-1/2) and FSDP
-# (v1/v2) can all train and shard the two transformers together. FSDP full shard
-# (`config/accelerate_configs/fsdp_full_shard.yaml`) also shards model params; ZeRO-3 unsupported.
+# and `transformer_2`; they are bundled into one prepared root, so DeepSpeed (ZeRO-1/2) shards them
+# together. For FSDP prefer FSDP2 (`config/accelerate_configs/fsdp2.yaml`), NOT FSDP1
+# (`fsdp_full_shard.yaml`): Wan2.2 routes each timestep to only ONE transformer (boundary_ratio),
+# so the bundled two-transformer + per-block wrap trips FSDP1's single-root lazy-init (`_is_root`
+# assertion) on the first forward. ZeRO-3 unsupported.
 config_file: config/deepspeed/deepspeed_zero2.yaml
 num_processes: 8  # Number of processes to launch (overrides config file)
 main_process_port: 29500

examples/grpo/lora/z_image/default.yaml

@@ -25,7 +25,10 @@ model:
   finetune_type: 'lora' # Options: full, lora
   lora_rank : 64
   lora_alpha : 128
-  master_weight_dtype: "fp32" # Z-Image requires fp32 master weight for better performance
+  trainable_parameters_dtype: "fp32" # Z-Image trains the transformer in fp32 for numerical stability
+  # frozen_parameters_dtype unset -> preserve each frozen component's from_pretrained dtype (no downcast):
+  # Z-Image ships the transformer in fp32 but the text encoder in bf16, so forcing one frozen dtype
+  # would override that. Set an explicit dtype only to force a cast (e.g. bf16 to save memory).
   target_modules: "default" # Options: all, default, or list of module names like ["to_k", "to_q", "to_v", "to_out.0"]
   model_name_or_path: "Tongyi-MAI/Z-Image"  # Options: Tongyi-MAI/Z-Image, Tongyi-MAI/Z-Image-Turbo
   model_type: "z-image"

examples/grpo/lora/z_image_turbo/default.yaml

@@ -25,7 +25,10 @@ model:
   finetune_type: 'lora' # Options: full, lora
   lora_rank : 64
   lora_alpha : 128
-  master_weight_dtype: "fp32" # Z-Image requires fp32 master weight for better performance
+  trainable_parameters_dtype: "fp32" # Z-Image-Turbo trains the transformer in fp32 for numerical stability
+  # frozen_parameters_dtype unset -> preserve each frozen component's from_pretrained dtype (no downcast):
+  # Z-Image-Turbo ships the transformer in fp32 but the text encoder in bf16, so forcing one frozen dtype
+  # would override that. Set an explicit dtype only to force a cast (e.g. bf16 to save memory).
   target_modules: "default" # Options: all, default, or list of module names like ["to_k", "to_q", "to_v", "to_out.0"]
   model_name_or_path: "Tongyi-MAI/Z-Image-Turbo"  # HuggingFace model ID or local path
   model_type: "z-image"  # Options: flux1, flux1-kontext, flux2, qwenimage, qwenimage-edit, z-image

examples/nft/full/z_image/default.yaml

@@ -23,7 +23,10 @@ data:
 # Model Configuration
 model:
   finetune_type: 'full' # Options: full, lora
-  master_weight_dtype: "fp32" # Z-Image requires fp32 master weight for better performance
+  trainable_parameters_dtype: "fp32" # Z-Image trains the transformer in fp32 for numerical stability
+  # frozen_parameters_dtype unset -> preserve each frozen component's from_pretrained dtype (no downcast):
+  # Z-Image ships the transformer in fp32 but the text encoder in bf16, so forcing one frozen dtype
+  # would override that. Set an explicit dtype only to force a cast (e.g. bf16 to save memory).
   target_modules: "default" # Options: all, default, or list of module names like ["to_k", "to_q", "to_v", "to_out.0"]
   model_name_or_path: "Tongyi-MAI/Z-Image"  # Options: Tongyi-MAI/Z-Image, Tongyi-MAI/Z-Image-Turbo
   model_type: "z-image"

examples/nft/full/z_image_turbo/default.yaml

@@ -23,7 +23,10 @@ data:
 # Model Configuration
 model:
   finetune_type: 'full' # Options: full, lora
-  master_weight_dtype: "fp32" # Z-Image-Turbo requires fp32 master weight for better performance
+  trainable_parameters_dtype: "fp32" # Z-Image-Turbo trains the transformer in fp32 for numerical stability
+  # frozen_parameters_dtype unset -> preserve each frozen component's from_pretrained dtype (no downcast):
+  # Z-Image-Turbo ships the transformer in fp32 but the text encoder in bf16, so forcing one frozen dtype
+  # would override that. Set an explicit dtype only to force a cast (e.g. bf16 to save memory).
   target_modules: "default" # Options: all, default, or list of module names like ["to_k", "to_q", "to_v", "to_out.0"]
   model_name_or_path: "Tongyi-MAI/Z-Image-Turbo"  # Options: Tongyi-MAI/Z-Image, Tongyi-MAI/Z-Image-Turbo
   model_type: "z-image"

examples/nft/lora/z_image/default.yaml

@@ -25,7 +25,10 @@ model:
   finetune_type: 'lora' # Options: full, lora
   lora_rank : 64
   lora_alpha : 128
-  master_weight_dtype: "fp32" # Z-Image requires fp32 master weight for better performance
+  trainable_parameters_dtype: "fp32" # Z-Image trains the transformer in fp32 for numerical stability
+  # frozen_parameters_dtype unset -> preserve each frozen component's from_pretrained dtype (no downcast):
+  # Z-Image ships the transformer in fp32 but the text encoder in bf16, so forcing one frozen dtype
+  # would override that. Set an explicit dtype only to force a cast (e.g. bf16 to save memory).
   target_modules: "default" # Options: all, default, or list of module names like ["to_k", "to_q", "to_v", "to_out.0"]
   model_name_or_path: "Tongyi-MAI/Z-Image"  # Options: Tongyi-MAI/Z-Image, Tongyi-MAI/Z-Image-Turbo
   model_type: "z-image"

src/flow_factory/hparams/model_args.py

@@ -46,11 +46,24 @@ class ModelArguments(ArgABC):
         metadata={"help": "Fine-tuning type. Options are ['full', 'lora']"}
     )
 
-    master_weight_dtype : Union[Literal['fp32', 'bf16', 'fp16', 'float16', 'bfloat16', 'float32'], torch.dtype] = field(
+    trainable_parameters_dtype : Union[Literal['fp32', 'bf16', 'fp16', 'float16', 'bfloat16', 'float32'], torch.dtype] = field(
         default='bfloat16',
         metadata={
-            "help": "Torch dtype for all trainable parameters (`requires_grad=True`). "
-                    "Non-trainable weights and floating-point buffers use the model inference dtype when they differ."
+            "help": "Torch dtype for all trainable parameters (`requires_grad=True`) -- i.e. the "
+                    "optimizer 'master weight' precision. (Renamed from the misnamed "
+                    "`master_weight_dtype`, which despite its name only ever set the *trainable* "
+                    "parameter dtype.)"
+        },
+    )
+    frozen_parameters_dtype : Optional[Union[Literal['fp32', 'bf16', 'fp16', 'float16', 'bfloat16', 'float32'], torch.dtype]] = field(
+        default=None,
+        metadata={
+            "help": "Torch dtype for frozen (`requires_grad=False`) parameters and floating-point "
+                    "buffers. `None` (default) preserves each frozen component's original "
+                    "`from_pretrained` dtype and never downcasts -- released checkpoints deliberately "
+                    "ship components in different dtypes (e.g. Z-Image: transformer fp32, text encoder "
+                    "bf16), so forcing one uniform frozen dtype would override those choices. Set an "
+                    "explicit dtype (e.g. 'bf16') to cast all frozen params to it, e.g. to save memory."
         },
     )
 
@@ -118,8 +131,10 @@ class ModelArguments(ArgABC):
     )
 
     def __post_init__(self):
-        if isinstance(self.master_weight_dtype, str):
-            self.master_weight_dtype = dtype_map[self.master_weight_dtype]
+        if isinstance(self.trainable_parameters_dtype, str):
+            self.trainable_parameters_dtype = dtype_map[self.trainable_parameters_dtype]
+        if isinstance(self.frozen_parameters_dtype, str):
+            self.frozen_parameters_dtype = dtype_map[self.frozen_parameters_dtype]
 
         # Normalize target_components to list
         if isinstance(self.target_components, str):
@@ -137,7 +152,9 @@ def __post_init__(self):
 
     def to_dict(self) -> dict[str, Any]:
         d = super().to_dict()
-        d['master_weight_dtype'] = str(self.master_weight_dtype).split('.')[-1]
+        d['trainable_parameters_dtype'] = str(self.trainable_parameters_dtype).split('.')[-1]
+        if self.frozen_parameters_dtype is not None:
+            d['frozen_parameters_dtype'] = str(self.frozen_parameters_dtype).split('.')[-1]
         return d
 
     def __str__(self) -> str: