From c1d2ed4256797acf738d201227fb633f118184eb Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Thu, 9 Jul 2026 10:40:34 +0800
Subject: [PATCH 1/5] [models] fix: load checkpoints for trainable components
 only

_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.
---
 .agents/knowledge/constraints.md |  3 +++
 src/flow_factory/models/abc.py   | 13 +++++++++++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/.agents/knowledge/constraints.md b/.agents/knowledge/constraints.md
index 415f7a56..ab5ca31d 100644
--- a/.agents/knowledge/constraints.md
+++ b/.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.
 
diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py
index bbe7586e..e4cd7221 100644
--- a/src/flow_factory/models/abc.py
+++ b/src/flow_factory/models/abc.py
@@ -1617,7 +1617,13 @@ def load_sharded_checkpoint(checkpoint_dir: str, index_file: str) -> Dict[str, t
 
     def _load_lora(self, path: str) -> None:
         """Load LoRA adapters for target components with auto-format detection."""
-        for comp_name in self.model_args.target_components:
+        # Iterate only components that actually carry target modules. Frozen
+        # bundled members (e.g. Wan2.2's `transformer_2`, kept in
+        # `target_components` solely to be FSDP-sharded for memory) map to None
+        # in `target_module_map` and are skipped by `save_checkpoint`; iterating
+        # `target_components` here would instead look for a `.../transformer_2/`
+        # subdir that was never written and log a spurious error on every resume.
+        for comp_name in self.trainable_component_names:
             if not hasattr(self, comp_name):
                 logger.warning(f"Component {comp_name} not found, skipping")
                 continue
@@ -1716,7 +1722,10 @@ def _load_lora(self, path: str) -> None:
 
     def _load_full_model(self, path: str, strict: bool = True) -> None:
         """Load full model weights for target components."""
-        for comp_name in self.model_args.target_components:
+        # Match `save_checkpoint`: only components with target modules are
+        # written, so frozen bundled members (`target_module_map[name] is None`)
+        # must be skipped here rather than iterating all `target_components`.
+        for comp_name in self.trainable_component_names:
             if not hasattr(self, comp_name):
                 logger.warning(f"Component {comp_name} not found, skipping")
                 continue

From 46f8c763d0ef4f25fb4ae8599eef02d80605bcc9 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Thu, 9 Jul 2026 19:26:33 +0800
Subject: [PATCH 2/5] [models] fix: expose _no_split_modules on ModelBundle for
 FSDP auto-wrap

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.
---
 src/flow_factory/models/model_bundle.py | 27 +++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/src/flow_factory/models/model_bundle.py b/src/flow_factory/models/model_bundle.py
index 76c452b6..d959b716 100644
--- a/src/flow_factory/models/model_bundle.py
+++ b/src/flow_factory/models/model_bundle.py
@@ -65,6 +65,33 @@ def __init__(self, members: Dict[str, nn.Module]):
                 )
         self.members = nn.ModuleDict(members)
 
+    @property
+    def _no_split_modules(self):
+        """Aggregate the members' ``_no_split_modules`` so accelerate's FSDP
+        ``TRANSFORMER_BASED_WRAP`` can discover the transformer block class(es).
+
+        accelerate's ``set_auto_wrap_policy`` reads ``getattr(root, "_no_split_modules")``
+        off the single root passed to ``prepare`` (this bundle). Without this the
+        list is empty, ``transformer_layer_cls`` becomes an empty set, and the whole
+        bundle is wrapped as ONE FSDP unit -> one monolithic flat param (the full
+        unsharded model materialized on every rank) -> OOM at init for large models
+        (e.g. Wan2.2 A14B: ~53GB flat param). We surface the block-class *names* that
+        the underlying model classes already declare (e.g. diffusers
+        ``WanTransformer3DModel._no_split_modules == ['WanTransformerBlock']``);
+        accelerate then resolves each name to its class via ``get_module_class_from_name``
+        over this root's submodule tree and wraps per-block. Returns ``None`` when no
+        member declares any (accelerate keeps its prior fallback unchanged).
+        """
+        names: list[str] = []
+        # Walk members' submodule tree (NOT self, to avoid recursing on this property);
+        # collect every ``_no_split_modules`` a nested module declares.
+        for module in self.members.modules():
+            nsm = getattr(module, "_no_split_modules", None)
+            if nsm:
+                names.extend(nsm)
+        deduped = list(dict.fromkeys(names))
+        return deduped or None
+
     def forward(self, component_name: str, /, *args, **kwargs):
         """Dispatch the call to ``members[component_name]``.
 

From 6d95785169df2a13fcf7a954baa8dabc74b464bd Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Thu, 9 Jul 2026 19:27:18 +0800
Subject: [PATCH 3/5] [precision/fsdp2] rename
 master_weight_dtype->trainable_parameters_dtype, 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.
---
 .agents/knowledge/constraints.md              |  2 +-
 .../knowledge/topics/autocast_param_swap.md   |  2 +-
 .agents/knowledge/topics/dtype_precision.md   |  2 +-
 examples/grpo/full/z_image/default.yaml       |  3 +-
 examples/grpo/full/z_image_turbo/default.yaml |  3 +-
 examples/grpo/lora/z_image/default.yaml       |  3 +-
 examples/grpo/lora/z_image_turbo/default.yaml |  3 +-
 examples/nft/full/z_image/default.yaml        |  3 +-
 examples/nft/full/z_image_turbo/default.yaml  |  3 +-
 examples/nft/lora/z_image/default.yaml        |  3 +-
 src/flow_factory/hparams/model_args.py        | 26 ++++--
 src/flow_factory/models/abc.py                | 91 ++++++++++---------
 src/flow_factory/trainers/abc.py              |  2 +
 13 files changed, 85 insertions(+), 61 deletions(-)

diff --git a/.agents/knowledge/constraints.md b/.agents/knowledge/constraints.md
index ab5ca31d..d3e32d58 100644
--- a/.agents/knowledge/constraints.md
+++ b/.agents/knowledge/constraints.md
@@ -123,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`.
 
 ---
 
diff --git a/.agents/knowledge/topics/autocast_param_swap.md b/.agents/knowledge/topics/autocast_param_swap.md
index d9436397..6daea426 100644
--- a/.agents/knowledge/topics/autocast_param_swap.md
+++ b/.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).
 
diff --git a/.agents/knowledge/topics/dtype_precision.md b/.agents/knowledge/topics/dtype_precision.md
index f1264f12..4992916c 100644
--- a/.agents/knowledge/topics/dtype_precision.md
+++ b/.agents/knowledge/topics/dtype_precision.md
@@ -9,7 +9,7 @@
 | Component | Runtime dtype | Why |
 |-----------|--------------|-----|
 | Transformer (frozen) | `inference_dtype` (bf16/fp16) | Memory savings for frozen params |
-| Transformer (trainable) | `master_weight_dtype` (fp32/bf16) | Gradient precision |
+| 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 |
diff --git a/examples/grpo/full/z_image/default.yaml b/examples/grpo/full/z_image/default.yaml
index 1b06aa21..198a2715 100644
--- a/examples/grpo/full/z_image/default.yaml
+++ b/examples/grpo/full/z_image/default.yaml
@@ -23,7 +23,8 @@ 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 needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
+  frozen_parameters_dtype: "fp32"    # Z-Image also requires frozen params (VAE / text encoder / frozen base) in fp32
   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"
diff --git a/examples/grpo/full/z_image_turbo/default.yaml b/examples/grpo/full/z_image_turbo/default.yaml
index 9bc068d0..d74abfdc 100644
--- a/examples/grpo/full/z_image_turbo/default.yaml
+++ b/examples/grpo/full/z_image_turbo/default.yaml
@@ -23,7 +23,8 @@ 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 needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
+  frozen_parameters_dtype: "fp32"    # Z-Image-Turbo also requires frozen params (VAE / text encoder / frozen base) in fp32
   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"
diff --git a/examples/grpo/lora/z_image/default.yaml b/examples/grpo/lora/z_image/default.yaml
index ae7f72c7..7c77c09c 100644
--- a/examples/grpo/lora/z_image/default.yaml
+++ b/examples/grpo/lora/z_image/default.yaml
@@ -25,7 +25,8 @@ 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 needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
+  frozen_parameters_dtype: "fp32"    # Z-Image also requires frozen params (VAE / text encoder / frozen base) in fp32
   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"
diff --git a/examples/grpo/lora/z_image_turbo/default.yaml b/examples/grpo/lora/z_image_turbo/default.yaml
index 129c4f5d..6df260dd 100644
--- a/examples/grpo/lora/z_image_turbo/default.yaml
+++ b/examples/grpo/lora/z_image_turbo/default.yaml
@@ -25,7 +25,8 @@ 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 needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
+  frozen_parameters_dtype: "fp32"    # Z-Image-Turbo also requires frozen params (VAE / text encoder / frozen base) in fp32
   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
diff --git a/examples/nft/full/z_image/default.yaml b/examples/nft/full/z_image/default.yaml
index 30ce76f8..14156370 100644
--- a/examples/nft/full/z_image/default.yaml
+++ b/examples/nft/full/z_image/default.yaml
@@ -23,7 +23,8 @@ 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 needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
+  frozen_parameters_dtype: "fp32"    # Z-Image also requires frozen params (VAE / text encoder / frozen base) in fp32
   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"
diff --git a/examples/nft/full/z_image_turbo/default.yaml b/examples/nft/full/z_image_turbo/default.yaml
index a21607f4..e3e40069 100644
--- a/examples/nft/full/z_image_turbo/default.yaml
+++ b/examples/nft/full/z_image_turbo/default.yaml
@@ -23,7 +23,8 @@ 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 needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
+  frozen_parameters_dtype: "fp32"    # Z-Image-Turbo also requires frozen params (VAE / text encoder / frozen base) in fp32
   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"
diff --git a/examples/nft/lora/z_image/default.yaml b/examples/nft/lora/z_image/default.yaml
index 6c00b9cc..010a0757 100644
--- a/examples/nft/lora/z_image/default.yaml
+++ b/examples/nft/lora/z_image/default.yaml
@@ -25,7 +25,8 @@ 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 needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
+  frozen_parameters_dtype: "fp32"    # Z-Image also requires frozen params (VAE / text encoder / frozen base) in fp32
   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"
diff --git a/src/flow_factory/hparams/model_args.py b/src/flow_factory/hparams/model_args.py
index e70bccc1..ef094c44 100644
--- a/src/flow_factory/hparams/model_args.py
+++ b/src/flow_factory/hparams/model_args.py
@@ -46,11 +46,21 @@ 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) falls back to the model inference dtype implied by "
+                    "`mixed_precision`, preserving prior behavior."
         },
     )
 
@@ -118,8 +128,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 +149,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:
diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py
index e4cd7221..c2486720 100644
--- a/src/flow_factory/models/abc.py
+++ b/src/flow_factory/models/abc.py
@@ -893,37 +893,54 @@ def _cast_module_mixed_precision(
         return n_trainable
 
     def _mix_precision(self):
-        """Apply mixed precision to default pipeline modules plus any extra ``target_components`` names."""
-        # Get inference and master dtypes
+        """Set trainable params to ``trainable_parameters_dtype`` and frozen params + floating-point
+        buffers to ``frozen_parameters_dtype`` (falling back to the inference dtype when unset).
+
+        This is the single place that decides every parameter's *original* dtype before
+        ``accelerator.prepare``; the trainer only bundles + prepares.
+
+        FSDP2 caveat: FSDP2 shards each unit with ONE original dtype, and accelerate upcasts the
+        trainable params to an fp32 master when ``mixed_precision != 'no'``. So a trained component
+        that also bundles frozen members (e.g. Wan2.2 trains ``transformer`` while ``transformer_2``
+        is frozen-but-sharded) would otherwise mix fp32/low-precision within a unit and trip FSDP2's
+        uniform-dtype assert. We therefore force the TRAINED components to a uniform fp32 original
+        dtype here (compute stays low-precision via accelerate's ``MixedPrecisionPolicy``); untrained
+        components (text encoder / VAE, not sharded) keep ``frozen_parameters_dtype``.
+        """
         inference_dtype = self._inference_dtype
-        master_dtype = self.model_args.master_weight_dtype
+        train_dtype = self.model_args.trainable_parameters_dtype
+        frozen_dtype = self.model_args.frozen_parameters_dtype or inference_dtype
 
-        # Get target components and all component names
         target_set = frozenset(self.model_args.target_components)
         component_names = self._resolve_component_names(None)
         merged_names = list(dict.fromkeys([*component_names, *self.model_args.target_components]))
 
-        # If master dtype is the same as inference dtype, cast all components to inference dtype
-        if master_dtype == inference_dtype:
-            # Cast all components to inference dtype
+        # FSDP2: trained (sharded) components need a uniform fp32 original dtype.
+        if self._is_fsdp2() and self.accelerator.mixed_precision != "no":
+            for name in merged_names:
+                self.get_component(name).to(dtype=torch.float32 if name in target_set else frozen_dtype)
+            logger.info(f"FSDP2: trained components -> fp32 (uniform orig dtype); other components -> {frozen_dtype}")
+            return
+
+        # Uniform dtype -> a single cast of every component suffices.
+        if train_dtype == frozen_dtype:
             for name in merged_names:
-                self.get_component(name).to(dtype=inference_dtype)
+                self.get_component(name).to(dtype=frozen_dtype)
             return
 
+        # Split: trainable -> train_dtype, frozen -> frozen_dtype (within trained components too).
         trainable_count = 0
         for name in merged_names:
             component = self.get_component(name)
             if name in target_set:
-                # Cast trainable parameters to master dtype
-                trainable_count += self._cast_module_mixed_precision(
-                    component, master_dtype, inference_dtype
-                )
+                trainable_count += self._cast_module_mixed_precision(component, train_dtype, frozen_dtype)
             else:
-                # Cast frozen parameters to inference dtype
-                component.to(dtype=inference_dtype)
+                component.to(dtype=frozen_dtype)
 
         if trainable_count > 0:
-            logger.info(f"Set {trainable_count} trainable parameters to {master_dtype}")
+            logger.info(
+                f"Set {trainable_count} trainable parameters to {train_dtype}, frozen params to {frozen_dtype}"
+            )
 
     # ============================== LoRA Management ==============================
     def apply_lora(
@@ -1181,37 +1198,21 @@ def is_param_match_key(name, keys, strict=True):
 
                 state_dict = clone_tensors_for_torch_save(self._unwrap(model).state_dict())
         elif self.accelerator.is_fsdp2:
-            # FSDP/FSDP2
+            # FSDP2: gather the full (unsharded) params to rank0 via the DTensor-aware API.
+            # NOTE: the previous `state_dict_keys` path toggled `requires_grad` at runtime to
+            # sub-select params, but FSDP2's `ignore_frozen_params` is keyed off the trainability
+            # captured at `fully_shard` time -- the runtime toggle is a no-op, so it yielded an
+            # EMPTY adapter. Gather straight through (LoRA params are exactly the trainable subset,
+            # so `ignore_frozen_params=True` returns them) and let the shared key-filter below
+            # narrow to `state_dict_keys` when provided.
             from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
-            if state_dict_keys is not None:
-                # Temporarily mark unwanted params as frozen
-                # This `requires_grad` trick does not work correctly. Don't know why.
-                original_state = {}
-                
-                # Freeze unwanted params
-                for name, param in model.named_parameters():
-                    original_state[name] = param.requires_grad
-                    param.requires_grad = is_param_match_key(name, state_dict_keys)
-                
-                options = StateDictOptions(
-                    full_state_dict=True,
-                    broadcast_from_rank0=True,
-                    cpu_offload=True,
-                    ignore_frozen_params=True,
-                )
-                state_dict = get_model_state_dict(model, options=options)
-                
-                # Restore original state
-                for name, param in model.named_parameters():
-                    param.requires_grad = original_state[name]
-            else:
-                options = StateDictOptions(
-                    full_state_dict=True, 
-                    broadcast_from_rank0=True, 
-                    cpu_offload=True, 
-                    ignore_frozen_params=ignore_frozen_params
-                )
-                state_dict = get_model_state_dict(model, options=options)
+            options = StateDictOptions(
+                full_state_dict=True,
+                broadcast_from_rank0=True,
+                cpu_offload=True,
+                ignore_frozen_params=ignore_frozen_params,
+            )
+            state_dict = get_model_state_dict(model, options=options)
         elif self.accelerator.distributed_type == DistributedType.FSDP:
             from torch.distributed.fsdp import FullStateDictConfig, StateDictType
             from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py
index 91bc1890..fe8602dc 100644
--- a/src/flow_factory/trainers/abc.py
+++ b/src/flow_factory/trainers/abc.py
@@ -346,6 +346,8 @@ def _initialization(self):
         eval_dataloader_list = [eval_dataloaders[n] for n in eval_dataloader_names]
 
         # One prepare call -> one DDP/FSDP/DeepSpeed root for the whole bundle.
+        # (Parameter dtypes -- incl. the FSDP2 uniform-fp32 requirement for sharded trained
+        # components -- are already handled in the adapter's `_mix_precision`.)
         prepared = self.accelerator.prepare(model_bundle, self.optimizer, *eval_dataloader_list)
         self.model_bundle = prepared[0]
         self.optimizer = prepared[1]

From 38ca25ea53a15f2230f0a01bec54cffd307c27d0 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Thu, 9 Jul 2026 19:27:18 +0800
Subject: [PATCH 4/5] [examples] docs: recommend FSDP2 over FSDP1 for Wan2.2

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.
---
 examples/grpo/full/wan22/i2v.yaml | 8 +++++---
 examples/grpo/full/wan22/t2v.yaml | 8 +++++---
 examples/grpo/lora/wan22/i2v.yaml | 8 +++++---
 examples/grpo/lora/wan22/t2v.yaml | 8 +++++---
 4 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/examples/grpo/full/wan22/i2v.yaml b/examples/grpo/full/wan22/i2v.yaml
index f4b02774..f8b89f6f 100644
--- a/examples/grpo/full/wan22/i2v.yaml
+++ b/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
diff --git a/examples/grpo/full/wan22/t2v.yaml b/examples/grpo/full/wan22/t2v.yaml
index 8635c175..0dc30916 100644
--- a/examples/grpo/full/wan22/t2v.yaml
+++ b/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
diff --git a/examples/grpo/lora/wan22/i2v.yaml b/examples/grpo/lora/wan22/i2v.yaml
index 10548eb4..b8446e98 100644
--- a/examples/grpo/lora/wan22/i2v.yaml
+++ b/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
diff --git a/examples/grpo/lora/wan22/t2v.yaml b/examples/grpo/lora/wan22/t2v.yaml
index 12718431..5cc66966 100644
--- a/examples/grpo/lora/wan22/t2v.yaml
+++ b/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

From 5cb276a0b7e90e60b68954a15089c9f0aabe7b47 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Thu, 9 Jul 2026 19:52:57 +0800
Subject: [PATCH 5/5] [models,examples] fix: preserve frozen params' loaded
 dtype by default

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.
---
 .agents/knowledge/topics/dtype_precision.md   |  2 +-
 examples/grpo/full/z_image/default.yaml       |  6 ++-
 examples/grpo/full/z_image_turbo/default.yaml |  6 ++-
 examples/grpo/lora/z_image/default.yaml       |  6 ++-
 examples/grpo/lora/z_image_turbo/default.yaml |  6 ++-
 examples/nft/full/z_image/default.yaml        |  6 ++-
 examples/nft/full/z_image_turbo/default.yaml  |  6 ++-
 examples/nft/lora/z_image/default.yaml        |  6 ++-
 src/flow_factory/hparams/model_args.py        |  7 ++-
 src/flow_factory/models/abc.py                | 49 ++++++++++++-------
 10 files changed, 66 insertions(+), 34 deletions(-)

diff --git a/.agents/knowledge/topics/dtype_precision.md b/.agents/knowledge/topics/dtype_precision.md
index 4992916c..e861275f 100644
--- a/.agents/knowledge/topics/dtype_precision.md
+++ b/.agents/knowledge/topics/dtype_precision.md
@@ -8,7 +8,7 @@
 
 | Component | Runtime dtype | Why |
 |-----------|--------------|-----|
-| Transformer (frozen) | `inference_dtype` (bf16/fp16) | Memory savings for frozen params |
+| 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 |
diff --git a/examples/grpo/full/z_image/default.yaml b/examples/grpo/full/z_image/default.yaml
index 198a2715..8b30e3ed 100644
--- a/examples/grpo/full/z_image/default.yaml
+++ b/examples/grpo/full/z_image/default.yaml
@@ -23,8 +23,10 @@ data:
 # Model Configuration
 model:
   finetune_type: 'full' # Options: full, lora
-  trainable_parameters_dtype: "fp32" # Z-Image needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
-  frozen_parameters_dtype: "fp32"    # Z-Image also requires frozen params (VAE / text encoder / frozen base) in fp32
+  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"
diff --git a/examples/grpo/full/z_image_turbo/default.yaml b/examples/grpo/full/z_image_turbo/default.yaml
index d74abfdc..08ad3659 100644
--- a/examples/grpo/full/z_image_turbo/default.yaml
+++ b/examples/grpo/full/z_image_turbo/default.yaml
@@ -23,8 +23,10 @@ data:
 # Model Configuration
 model:
   finetune_type: 'full' # Options: full, lora
-  trainable_parameters_dtype: "fp32" # Z-Image-Turbo needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
-  frozen_parameters_dtype: "fp32"    # Z-Image-Turbo also requires frozen params (VAE / text encoder / frozen base) in fp32
+  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"
diff --git a/examples/grpo/lora/z_image/default.yaml b/examples/grpo/lora/z_image/default.yaml
index 7c77c09c..7af80ec9 100644
--- a/examples/grpo/lora/z_image/default.yaml
+++ b/examples/grpo/lora/z_image/default.yaml
@@ -25,8 +25,10 @@ model:
   finetune_type: 'lora' # Options: full, lora
   lora_rank : 64
   lora_alpha : 128
-  trainable_parameters_dtype: "fp32" # Z-Image needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
-  frozen_parameters_dtype: "fp32"    # Z-Image also requires frozen params (VAE / text encoder / frozen base) in fp32
+  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"
diff --git a/examples/grpo/lora/z_image_turbo/default.yaml b/examples/grpo/lora/z_image_turbo/default.yaml
index 6df260dd..9ce85cdd 100644
--- a/examples/grpo/lora/z_image_turbo/default.yaml
+++ b/examples/grpo/lora/z_image_turbo/default.yaml
@@ -25,8 +25,10 @@ model:
   finetune_type: 'lora' # Options: full, lora
   lora_rank : 64
   lora_alpha : 128
-  trainable_parameters_dtype: "fp32" # Z-Image-Turbo needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
-  frozen_parameters_dtype: "fp32"    # Z-Image-Turbo also requires frozen params (VAE / text encoder / frozen base) in fp32
+  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
diff --git a/examples/nft/full/z_image/default.yaml b/examples/nft/full/z_image/default.yaml
index 14156370..3ee70e87 100644
--- a/examples/nft/full/z_image/default.yaml
+++ b/examples/nft/full/z_image/default.yaml
@@ -23,8 +23,10 @@ data:
 # Model Configuration
 model:
   finetune_type: 'full' # Options: full, lora
-  trainable_parameters_dtype: "fp32" # Z-Image needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
-  frozen_parameters_dtype: "fp32"    # Z-Image also requires frozen params (VAE / text encoder / frozen base) in fp32
+  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"
diff --git a/examples/nft/full/z_image_turbo/default.yaml b/examples/nft/full/z_image_turbo/default.yaml
index e3e40069..9e7d2d38 100644
--- a/examples/nft/full/z_image_turbo/default.yaml
+++ b/examples/nft/full/z_image_turbo/default.yaml
@@ -23,8 +23,10 @@ data:
 # Model Configuration
 model:
   finetune_type: 'full' # Options: full, lora
-  trainable_parameters_dtype: "fp32" # Z-Image-Turbo needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
-  frozen_parameters_dtype: "fp32"    # Z-Image-Turbo also requires frozen params (VAE / text encoder / frozen base) in fp32
+  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"
diff --git a/examples/nft/lora/z_image/default.yaml b/examples/nft/lora/z_image/default.yaml
index 010a0757..840365c1 100644
--- a/examples/nft/lora/z_image/default.yaml
+++ b/examples/nft/lora/z_image/default.yaml
@@ -25,8 +25,10 @@ model:
   finetune_type: 'lora' # Options: full, lora
   lora_rank : 64
   lora_alpha : 128
-  trainable_parameters_dtype: "fp32" # Z-Image needs fp32 for numerical stability -- BOTH trainable and frozen (see frozen_parameters_dtype below)
-  frozen_parameters_dtype: "fp32"    # Z-Image also requires frozen params (VAE / text encoder / frozen base) in fp32
+  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"
diff --git a/src/flow_factory/hparams/model_args.py b/src/flow_factory/hparams/model_args.py
index ef094c44..1f4977b7 100644
--- a/src/flow_factory/hparams/model_args.py
+++ b/src/flow_factory/hparams/model_args.py
@@ -59,8 +59,11 @@ class ModelArguments(ArgABC):
         default=None,
         metadata={
             "help": "Torch dtype for frozen (`requires_grad=False`) parameters and floating-point "
-                    "buffers. `None` (default) falls back to the model inference dtype implied by "
-                    "`mixed_precision`, preserving prior behavior."
+                    "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."
         },
     )
 
diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py
index c2486720..8332fe78 100644
--- a/src/flow_factory/models/abc.py
+++ b/src/flow_factory/models/abc.py
@@ -872,44 +872,50 @@ def _cast_module_mixed_precision(
         self,
         component: torch.nn.Module,
         train_dtype: torch.dtype,
-        frozen_dtype: torch.dtype,
+        frozen_dtype: Optional[torch.dtype],
     ) -> int:
         """
         Set floating-point parameters/buffers without a trainable round-trip through frozen_dtype.
 
-        Trainable parameters use ``train_dtype``; frozen parameters and floating-point buffers use
-        ``frozen_dtype``. Integer/bool buffers are left unchanged (same as ``Module.to``).
+        Trainable parameters use ``train_dtype``. Frozen parameters and floating-point buffers use
+        ``frozen_dtype`` when it is set, or are left at their loaded dtype when ``frozen_dtype`` is
+        ``None`` (preserve). Integer/bool buffers are left unchanged (same as ``Module.to``).
         """ 
         n_trainable = 0
         for _, param in component.named_parameters():
             if param.requires_grad:
                 param.data = param.data.to(dtype=train_dtype)
                 n_trainable += 1
-            else:
+            elif frozen_dtype is not None:
                 param.data = param.data.to(dtype=frozen_dtype)
         for _, buf in component.named_buffers():
-            if buf.is_floating_point():
+            if buf.is_floating_point() and frozen_dtype is not None:
                 buf.data = buf.data.to(dtype=frozen_dtype)
         return n_trainable
 
     def _mix_precision(self):
-        """Set trainable params to ``trainable_parameters_dtype`` and frozen params + floating-point
-        buffers to ``frozen_parameters_dtype`` (falling back to the inference dtype when unset).
+        """Set trainable params to ``trainable_parameters_dtype``; by default leave frozen params
+        and floating-point buffers at their loaded (``from_pretrained``) dtype.
 
         This is the single place that decides every parameter's *original* dtype before
         ``accelerator.prepare``; the trainer only bundles + prepares.
 
+        Frozen-dtype policy: ``frozen_parameters_dtype=None`` (default) preserves each frozen
+        parameter/buffer's original dtype and never downcasts -- a released checkpoint deliberately
+        ships components in different dtypes (e.g. Z-Image: transformer fp32, text encoder bf16), and
+        forcing one uniform frozen dtype would override those choices. Set an explicit
+        ``frozen_parameters_dtype`` to opt into casting every frozen param to that dtype.
+
         FSDP2 caveat: FSDP2 shards each unit with ONE original dtype, and accelerate upcasts the
         trainable params to an fp32 master when ``mixed_precision != 'no'``. So a trained component
         that also bundles frozen members (e.g. Wan2.2 trains ``transformer`` while ``transformer_2``
         is frozen-but-sharded) would otherwise mix fp32/low-precision within a unit and trip FSDP2's
         uniform-dtype assert. We therefore force the TRAINED components to a uniform fp32 original
         dtype here (compute stays low-precision via accelerate's ``MixedPrecisionPolicy``); untrained
-        components (text encoder / VAE, not sharded) keep ``frozen_parameters_dtype``.
+        components are cast to ``frozen_parameters_dtype`` when set, else preserved.
         """
-        inference_dtype = self._inference_dtype
         train_dtype = self.model_args.trainable_parameters_dtype
-        frozen_dtype = self.model_args.frozen_parameters_dtype or inference_dtype
+        frozen_dtype = self.model_args.frozen_parameters_dtype  # None -> preserve loaded dtype
 
         target_set = frozenset(self.model_args.target_components)
         component_names = self._resolve_component_names(None)
@@ -918,28 +924,37 @@ def _mix_precision(self):
         # FSDP2: trained (sharded) components need a uniform fp32 original dtype.
         if self._is_fsdp2() and self.accelerator.mixed_precision != "no":
             for name in merged_names:
-                self.get_component(name).to(dtype=torch.float32 if name in target_set else frozen_dtype)
-            logger.info(f"FSDP2: trained components -> fp32 (uniform orig dtype); other components -> {frozen_dtype}")
+                if name in target_set:
+                    self.get_component(name).to(dtype=torch.float32)
+                elif frozen_dtype is not None:
+                    self.get_component(name).to(dtype=frozen_dtype)
+                # else: preserve the untrained component's loaded dtype
+            logger.info(
+                f"FSDP2: trained components -> fp32 (uniform orig dtype); "
+                f"other components -> {frozen_dtype or 'preserved (loaded dtype)'}"
+            )
             return
 
-        # Uniform dtype -> a single cast of every component suffices.
-        if train_dtype == frozen_dtype:
+        # Explicit uniform dtype -> a single cast of every component suffices.
+        if frozen_dtype is not None and train_dtype == frozen_dtype:
             for name in merged_names:
                 self.get_component(name).to(dtype=frozen_dtype)
             return
 
-        # Split: trainable -> train_dtype, frozen -> frozen_dtype (within trained components too).
+        # Split: trainable -> train_dtype; frozen -> frozen_dtype, or preserved when None.
         trainable_count = 0
         for name in merged_names:
             component = self.get_component(name)
             if name in target_set:
                 trainable_count += self._cast_module_mixed_precision(component, train_dtype, frozen_dtype)
-            else:
+            elif frozen_dtype is not None:
                 component.to(dtype=frozen_dtype)
+            # else: preserve the fully-frozen component's loaded dtype
 
         if trainable_count > 0:
             logger.info(
-                f"Set {trainable_count} trainable parameters to {train_dtype}, frozen params to {frozen_dtype}"
+                f"Set {trainable_count} trainable parameters to {train_dtype}; "
+                f"frozen params -> {frozen_dtype or 'preserved (loaded dtype)'}"
             )
 
     # ============================== LoRA Management ==============================