Update text_encoder_offload and accelerate saving hook

README.md

@@ -141,11 +141,14 @@ Table: inference entry selection
 
 ### 2.2 GPU Memory Saving Options
 
-Since Wan2.1 has a very large number of parameters, we need to consider memory optimization strategies to adapt to consumer-grade GPUs. We provide `GPU_memory_mode` for each prediction file, allowing you to choose between `model_cpu_offload`, `model_cpu_offload_and_qfloat8`, and `sequential_cpu_offload`. This solution is also applicable to CogVideoX-Fun generation.
+Since Wan2.1 has a very large number of parameters, we need to consider memory optimization strategies to adapt to consumer-grade GPUs. We provide `GPU_memory_mode` for each prediction file, allowing you to choose between `model_cpu_offload`, `model_cpu_offload_and_qfloat8`, `model_group_offload`, `sequential_cpu_offload`, `text_encoder_offload`, and `model_full_load_and_qfloat8`. This solution is also applicable to CogVideoX-Fun generation.
 
 - `model_cpu_offload`: The entire model is moved to the CPU after use, saving some GPU memory.
 - `model_cpu_offload_and_qfloat8`: The entire model is moved to the CPU after use, and the transformer model is quantized to float8, saving more GPU memory.
 - `sequential_cpu_offload`: Each layer of the model is moved to the CPU after use. It is slower but saves a significant amount of GPU memory.
+- `model_group_offload`: Weights are transferred between the CPU and GPU at the leaf level, with streams used for asynchronous prefetching, balancing speed and GPU memory.
+- `text_encoder_offload`: Only the text encoder is kept on the CPU and streamed to the GPU module by module for a single encoding pass, while the transformer and VAE stay resident on the GPU. Text encoding runs only once at the beginning of each request, so the transfer cost is paid once instead of at every step. An `_and_qfloat8` suffix can be added.
+- `model_full_load_and_qfloat8`: The model stays resident on the GPU and only the transformer is quantized to float8. A good choice when GPU memory is borderline but higher speed is required.
 
 `qfloat8` may slightly reduce model performance but saves more GPU memory. If you have sufficient GPU memory, it is recommended to use `model_cpu_offload`.
 

README_ja-JP.md

@@ -141,11 +141,14 @@ mkdir models/Personalized_Model
 
 ### 2.2 顕存節約方案
 
-Wan2.1のパラメータが非常に大きいため、GPUメモリを節約し、コンシューマー向けGPUに適応させる必要があります。各予測ファイルには`GPU_memory_mode`を提供しており、`model_cpu_offload`、`model_cpu_offload_and_qfloat8`、`sequential_cpu_offload`の中から選択できます。この方法はCogVideoX-Funの生成にも適用されます。
+Wan2.1のパラメータが非常に大きいため、GPUメモリを節約し、コンシューマー向けGPUに適応させる必要があります。各予測ファイルには`GPU_memory_mode`を提供しており、`model_cpu_offload`、`model_cpu_offload_and_qfloat8`、`model_group_offload`、`sequential_cpu_offload`、`text_encoder_offload`、`model_full_load_and_qfloat8`の中から選択できます。この方法はCogVideoX-Funの生成にも適用されます。
 
 - `model_cpu_offload`: モデル全体が使用後にCPUに移動し、一部のGPUメモリを節約します。
 - `model_cpu_offload_and_qfloat8`: モデル全体が使用後にCPUに移動し、Transformerモデルに対してfloat8の量子化を行い、より多くのGPUメモリを節約します。
 - `sequential_cpu_offload`: モデルの各層が使用後にCPUに移動します。速度は遅くなりますが、大量のGPUメモリを節約します。
+- `model_group_offload`: 重みをleaf単位でCPUとGPUの間で転送し、streamによる非同期プリフェッチを利用して、速度とGPUメモリのバランスを取ります。
+- `text_encoder_offload`: テキストエンコーダーのみをCPUに保持し、モジュール単位でGPUにストリーミングして一度だけエンコードを行い、TransformerとVAEはGPUに常駐します。テキストエンコードは各リクエストの先頭で一度だけ実行されるため、転送コストは毎ステップではなく一度だけ発生します。`_and_qfloat8`サフィックスを追加できます。
+- `model_full_load_and_qfloat8`: モデルはGPUに常駐し、Transformerのみをfloat8に量子化します。GPUメモリがぎりぎりで、かつ高い速度が求められる場合に適した選択肢です。
 
 `qfloat8`はモデルの性能を部分的に低下させる可能性がありますが、より多くのGPUメモリを節約できます。十分なGPUメモリがある場合は、`model_cpu_offload`の使用をお勧めします。
 

README_zh-CN.md

@@ -155,6 +155,7 @@ cd VideoX-Fun
 - model_group_offload：以leaf层级在cpu与gpu之间搬运权重，并借助stream异步预取，兼顾速度与显存。
 - model_cpu_offload_and_qfloat8：整个模型在使用后会进入cpu，并且对transformer模型进行了float8的量化，可以节省更多的显存。
 - model_cpu_offload：整个模型在使用后会进入cpu，可以节省部分显存。
+- text_encoder_offload：只把文本编码器留在cpu，按模块流式搬上gpu做唯一一次编码，transformer与vae常驻gpu；文本编码在每个请求开头只执行一次，因此搬运只付一次而不是每步一次。可加`_and_qfloat8`后缀。
 - model_full_load_and_qfloat8：模型常驻gpu，仅对transformer做float8量化，显存临界且对速度要求较高时可选。
 - 默认（传入model_full_load或其他取值）：模型全部进入gpu，速度最快，显存需求最高。
 

examples/minimax_h3/predict_i2v.py

@@ -20,7 +20,7 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
                               convert_model_weight_to_float8, merge_lora,
                               save_videos_with_audio_grid, unmerge_lora)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload, text_encoder_offload, text_encoder_offload_and_qfloat8].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -36,6 +36,9 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
 # 
 # sequential_cpu_offload means that each layer of the model will be moved to the CPU after use, 
 # resulting in slower speeds but saving a large amount of GPU memory.
+# 
+# text_encoder_offload keeps only the text encoder in CPU RAM, streaming it to the GPU module by module for
+# its single encode pass, so the transformer stays resident (add `_and_qfloat8` for float8 storage).
 GPU_memory_mode     = "model_group_offload"
 # Multi GPUs config
 # Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used. 

examples/minimax_h3/predict_ref2va.py

@@ -21,7 +21,7 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
                               convert_model_weight_to_float8, merge_lora,
                               save_videos_with_audio_grid, unmerge_lora)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload, text_encoder_offload, text_encoder_offload_and_qfloat8].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -37,6 +37,9 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
 # 
 # sequential_cpu_offload means that each layer of the model will be moved to the CPU after use, 
 # resulting in slower speeds but saving a large amount of GPU memory.
+# 
+# text_encoder_offload keeps only the text encoder in CPU RAM, streaming it to the GPU module by module for
+# its single encode pass, so the transformer stays resident (add `_and_qfloat8` for float8 storage).
 GPU_memory_mode     = "model_group_offload"
 # Multi GPUs config
 # Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used. 

examples/minimax_h3/predict_t2v.py

@@ -19,7 +19,7 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
                               convert_model_weight_to_float8, merge_lora,
                               save_videos_with_audio_grid, unmerge_lora)
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload, text_encoder_offload, text_encoder_offload_and_qfloat8].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -34,6 +34,10 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
 # balancing memory efficiency and speed between full-module and leaf-level offloading methods.
 # 
 # sequential_cpu_offload means that each layer of the model will be moved to the CPU after use, 
+# resulting in slower speeds but saving a large amount of GPU memory.
+# 
+# text_encoder_offload keeps only the text encoder in CPU RAM, streaming it to the GPU module by module for
+# its single encode pass, so the transformer stays resident (add `_and_qfloat8` for float8 storage). 
 GPU_memory_mode     = "model_group_offload"
 # Multi GPUs config
 # Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used. 

examples/minimax_h3_fun/predict_v2v_control.py

@@ -22,7 +22,7 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
                               save_videos_with_audio_grid, unmerge_lora)
 
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload, text_encoder_offload, text_encoder_offload_and_qfloat8].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -38,6 +38,9 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
 # 
 # sequential_cpu_offload means that each layer of the model will be moved to the CPU after use, 
 # resulting in slower speeds but saving a large amount of GPU memory.
+# 
+# text_encoder_offload keeps only the text encoder in CPU RAM, streaming it to the GPU module by module for
+# its single encode pass, so the transformer stays resident (add `_and_qfloat8` for float8 storage).
 GPU_memory_mode     = "model_group_offload"
 # Multi GPUs config
 # Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used. 

examples/minimax_h3_fun/predict_v2v_control_inpaint.py

@@ -22,7 +22,7 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
                               save_videos_with_audio_grid, unmerge_lora)
 
 
-# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload].
+# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, model_group_offload, sequential_cpu_offload, text_encoder_offload, text_encoder_offload_and_qfloat8].
 # model_full_load means that the entire model will be moved to the GPU.
 # 
 # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
@@ -38,6 +38,9 @@ from videox_fun.utils import (MiniMaxH3Scheduler, apply_gpu_memory_mode,
 # 
 # sequential_cpu_offload means that each layer of the model will be moved to the CPU after use, 
 # resulting in slower speeds but saving a large amount of GPU memory.
+# 
+# text_encoder_offload keeps only the text encoder in CPU RAM, streaming it to the GPU module by module for
+# its single encode pass, so the transformer stays resident (add `_and_qfloat8` for float8 storage).
 GPU_memory_mode     = "model_group_offload"
 # Multi GPUs config
 # Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used. 

videox_fun/__init__.py

@@ -1,5 +1,8 @@
+import functools
 import importlib.util
 import os
+import warnings
+from weakref import WeakSet
 
 if importlib.util.find_spec("turbox") is not None:
     import turbox
@@ -16,3 +19,139 @@ if os.environ.get("VIDEOX_PERF", "0").strip() not in ("", "0"):
     from .utils.perf_metrics import install_training as _install_perf_training
 
     _install_perf_training()
+
+# Repository-wide FSDP optimizer restore compatibility. Keep this bootstrap here
+# so importing the package installs it without eagerly importing videox_fun.utils.
+# Preserve explicit opt-ins and opt-outs when the package is reloaded.
+if "_accelerate_patch_owners" not in globals():
+    _accelerate_patch_owners = WeakSet()
+    _accelerate_patch_disabled_owners = WeakSet()
+    _accelerate_patches_automatic = False
+    _accelerate_patch_warned_unsupported = False
+
+if "_failed_checkpoint_owners" not in globals():
+    _failed_checkpoint_owners = WeakSet()
+
+
+def _accelerate_patch_requested():
+    return os.environ.get("VIDEOX_ACCELERATE_PATCHES", "1").strip().lower() not in ("", "0", "false", "no", "off")
+
+
+def _eligible_fsdp_checkpoint(plugin):
+    return (getattr(plugin, "fsdp_version", 1) == 1
+            and getattr(getattr(plugin, "state_dict_type", None), "name", None) == "SHARDED_STATE_DICT")
+
+
+def _supported_accelerate_versions():
+    import accelerate
+    import torch
+    from packaging.version import Version
+
+    versions = (Version(accelerate.__version__), Version(torch.__version__))
+    return tuple(version.public for version in versions) == ("1.7.0", "2.7.0")
+
+
+def _install_accelerate_dispatch():
+    import accelerate.accelerator as implementation
+
+    original = implementation.load_fsdp_optimizer
+    if getattr(original, "_videox_optimizer_dispatch", False) is True:
+        return
+
+    @functools.wraps(original)
+    def load_optimizer(plugin, owner, optimizer, model, input_dir, optimizer_index=0, adapter_only=False):
+        global _accelerate_patch_warned_unsupported
+        # The adapter requires one optimizer per prepared model. Preserve native
+        # handling for other layouts unless the caller explicitly opts in.
+        automatic = (_accelerate_patches_automatic and _accelerate_patch_requested()
+                     and len(getattr(owner, "_models", ())) == len(getattr(owner, "_optimizers", ())))
+        enabled = owner in _accelerate_patch_owners or automatic
+        if enabled and owner not in _accelerate_patch_disabled_owners and not adapter_only and _eligible_fsdp_checkpoint(plugin):
+            if _supported_accelerate_versions():
+                from .utils.checkpoint_utils import \
+                    _load_sharded_fsdp_optimizer
+
+                return _load_sharded_fsdp_optimizer(owner, optimizer, model, input_dir, optimizer_index)
+            if not _accelerate_patch_warned_unsupported:
+                warnings.warn("VideoX FSDP optimizer patch supports Accelerate 1.7.0 / PyTorch 2.7.0; "
+                              "using native restore on this version", RuntimeWarning, stacklevel=2)
+                _accelerate_patch_warned_unsupported = True
+        return original(plugin, owner, optimizer, model, input_dir, optimizer_index, adapter_only)
+
+    load_optimizer._videox_optimizer_dispatch = True
+    # Accelerator imports this symbol directly; existing Accelerator class aliases
+    # and instances therefore use the same dispatch without constructor wrapping.
+    implementation.load_fsdp_optimizer = load_optimizer
+
+
+def _install_checkpoint_lifecycle():
+    from accelerate import Accelerator
+
+    if getattr(Accelerator.save_state, "_videox_native_checkpoint", None) is not None:
+        return
+    native_save, native_load, native_backward = Accelerator.save_state, Accelerator.load_state, Accelerator.backward
+
+    @functools.wraps(native_save)
+    def save_state(owner, output_dir=None, safe_serialization=True, **kwargs):
+        if getattr(owner.distributed_type, "name", None) == "FSDP":
+            from .utils.checkpoint_utils import dispatch_checkpoint
+
+            if owner.project_configuration.automatic_checkpoint_naming:
+                output_dir = None
+            return dispatch_checkpoint(owner, "save_state", native_save, output_dir,
+                                       safe_serialization=safe_serialization, **kwargs)
+        return native_save(owner, output_dir, safe_serialization=safe_serialization, **kwargs)
+
+    @functools.wraps(native_load)
+    def load_state(owner, input_dir=None, **kwargs):
+        if getattr(owner.distributed_type, "name", None) == "FSDP":
+            from .utils.checkpoint_utils import dispatch_checkpoint
+
+            return dispatch_checkpoint(owner, "load_state", native_load, input_dir, **kwargs)
+        return native_load(owner, input_dir, **kwargs)
+
+    @functools.wraps(native_backward)
+    def backward(owner, *args, **kwargs):
+        if owner in _failed_checkpoint_owners:
+            raise RuntimeError("Checkpoint restore failed; reload successfully before continuing training")
+        return native_backward(owner, *args, **kwargs)
+
+    save_state._videox_native_checkpoint = native_save
+    load_state._videox_native_checkpoint = native_load
+    backward._videox_native_checkpoint = native_backward
+    Accelerator.save_state, Accelerator.load_state, Accelerator.backward = save_state, load_state, backward
+
+
+def install_accelerate_patches():
+    """Install lazy lifecycle guards and the version-specific optimizer adapter."""
+    global _accelerate_patches_automatic
+    if importlib.util.find_spec("accelerate") is None:
+        return
+    # The guards also compare rank opt-out/version settings before FSDP I/O.
+    _install_checkpoint_lifecycle()
+    if not _accelerate_patch_requested() or not _supported_accelerate_versions():
+        return
+    _install_accelerate_dispatch()
+    _accelerate_patches_automatic = True
+
+
+def enable_accelerate_patch(accelerator):
+    plugin = getattr(accelerator.state, "fsdp_plugin", None)
+    if getattr(accelerator.distributed_type, "name", None) != "FSDP" or not _eligible_fsdp_checkpoint(plugin):
+        return accelerator
+    if not _supported_accelerate_versions():
+        raise RuntimeError("FSDP optimizer patch supports Accelerate 1.7.0 / PyTorch 2.7.0; "
+                           "validate the adapter before enabling it on another version")
+    _install_checkpoint_lifecycle()
+    _install_accelerate_dispatch()
+    _accelerate_patch_disabled_owners.discard(accelerator)
+    _accelerate_patch_owners.add(accelerator)
+    return accelerator
+
+
+def disable_accelerate_patch(accelerator):
+    _accelerate_patch_owners.discard(accelerator)
+    _accelerate_patch_disabled_owners.add(accelerator)
+
+
+install_accelerate_patches()

videox_fun/utils/fsdp_utils.py

@@ -1,10 +1,19 @@
-"""Blockwise autograd replay for memory-bounded FSDP training."""
+"""Blockwise replay and opt-in CPU-offload memory management for FSDP1."""
 
-from contextlib import nullcontext
+import gc
+from contextlib import contextmanager, nullcontext
 
 import torch
 
-__all__ = ["checkpoint_fsdp_forward"]
+__all__ = ['checkpoint_fsdp_forward', 'release_frozen_fsdp_gradients', 'release_completed_fsdp_gradients', 'release_offloaded_fsdp_gradients', 'release_fsdp_gradients_during_validation', 'release_fsdp_gradients_during_replay', 'configure_fsdp_pageable_cpu_parameters', 'configure_fsdp_precision']
+
+
+def _require_cpu_offload_runtime():
+    from packaging.version import Version
+
+    if Version(torch.__version__).release[:2] != (2, 7):
+        raise RuntimeError("FSDP CPU-offload memory patches support PyTorch 2.7; "
+                           "validate private runtime contracts before enabling on another version")
 
 
 def checkpoint_fsdp_forward(model, accumulated, *, context_fn=None, **kwargs):
@@ -105,3 +114,296 @@ def checkpoint_fsdp_forward(model, accumulated, *, context_fn=None, **kwargs):
 
     with torch.autograd.graph.save_on_cpu():
         return Replay.apply(torch.empty(0, device=device, requires_grad=True), *values)
+
+
+def release_frozen_fsdp_gradients(model):
+    """Release unused CPU gradient buffers after permanently frozen forwards.
+
+    FSDP1 allocates CPU-offloaded gradients even for frozen parameters. Clearing
+    each unit after forward bounds their lifetime; FSDP supports None and will
+    allocate the buffer again on its next forward. Teacher/EMA/text weights and
+    saved activations remain live. This must not be used on trainable models.
+    """
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+    _require_cpu_offload_runtime()
+
+    if any(parameter.requires_grad for parameter in model.parameters()):
+        raise ValueError('Frozen FSDP gradient release requires all parameters to be frozen')
+    if getattr(model, '_videox_frozen_gradient_hooks', None) is not None:
+        return model
+    hooks = []
+
+    def release(module, inputs, output):
+        handle = module._handle
+        if handle is not None and handle._offload_params:
+            if handle.flat_param.requires_grad:
+                raise RuntimeError('A permanently frozen FSDP model was made trainable')
+            handle.flat_param._cpu_grad = None
+
+    def release_after_init(handle, initialize):
+        # Root lazy initialization visits every handle before the first forward.
+        # Drop each frozen unit's unused buffer immediately so those allocations
+        # do not accumulate into another whole model during lazy initialization.
+        def initialize_frozen(*args, **kwargs):
+            if handle.flat_param.requires_grad:
+                raise RuntimeError('A permanently frozen FSDP model was made trainable')
+            result = initialize(*args, **kwargs)
+            handle.flat_param._cpu_grad = None
+            return result
+        return initialize_frozen
+
+    for module in model.modules():
+        if isinstance(module, FSDP):
+            handle = module._handle
+            if handle is not None and handle._offload_params:
+                handle.init_flat_param_attributes = release_after_init(handle, handle.init_flat_param_attributes)
+            hooks.append(module.register_forward_hook(release))
+
+    def trim(module, inputs, output):
+        torch._C._host_emptyCache()
+
+    hooks.append(model.register_forward_hook(trim))
+    model._videox_frozen_gradient_hooks = hooks
+    return model
+
+
+def release_completed_fsdp_gradients(*models):
+    """Free unused CPU gradient storage after zero_grad or before resume.
+
+    FSDP recreates these buffers on the next forward. This bounds the save
+    peak when saving multiple optimizers without discarding accumulated gradients.
+    """
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+    _require_cpu_offload_runtime()
+
+    for model in models:
+        if any(parameter.grad is not None for parameter in model.parameters()):
+            raise RuntimeError('Cannot release gradients before optimizer zero_grad(set_to_none=True)')
+    released = 0
+    for model in models:
+        for module in model.modules():
+            if not isinstance(module, FSDP):
+                continue
+            handle = module._handle
+            if handle is None or not handle._offload_params:
+                continue
+            parameter = handle.flat_param
+            gradient = getattr(parameter, '_cpu_grad', None)
+            if gradient is not None:
+                released += gradient.numel() * gradient.element_size()
+            parameter.grad = None
+            parameter._cpu_grad = None
+            del gradient
+    # Per-forward autograd Function classes can keep completed replay accumulators
+    # alive through Python cycles even after their loss/output variables are gone.
+    gc.collect()
+    torch._C._host_emptyCache()
+    return released
+
+
+def release_offloaded_fsdp_gradients(model):
+    """Drop redundant CUDA grad shards after FSDP1 has offloaded them to CPU.
+
+    PyTorch 2.7 retains ``_saved_grad_shard`` until the end of backward even
+    with CPUOffload. The optimizer consumes ``_cpu_grad`` on this path.
+    FSDP records the D2H copy's stream before this callback returns, so the
+    allocator keeps the CUDA storage alive until the asynchronous copy ends.
+    Only explicitly marked modules use this compatibility workaround.
+    """
+    import functools
+    import torch.distributed.fsdp._runtime_utils as runtime
+
+    _require_cpu_offload_runtime()
+
+    for module in model.modules():
+        module._videox_release_cpu_grad = True
+    original = runtime._post_reduce_grad_callback
+    if getattr(original, "_videox_cpu_grad_release", False):
+        return
+
+    @functools.wraps(original)
+    def post_reduce(state, handle, grad_to_offload):
+        original(state, handle, grad_to_offload)
+        if not (getattr(state, "_videox_release_cpu_grad", False)
+                and handle._offload_params and not handle._has_optim_in_backward):
+            return
+        parameter = handle.flat_param
+        cpu_grad = getattr(parameter, "_cpu_grad", None)
+        saved = getattr(parameter, "_saved_grad_shard", None)
+        if isinstance(cpu_grad, torch.Tensor) and cpu_grad.device.type == "cpu" \
+                and isinstance(saved, torch.Tensor) and saved.is_cuda:
+            del parameter._saved_grad_shard
+
+    post_reduce._videox_cpu_grad_release = True
+    runtime._post_reduce_grad_callback = post_reduce
+
+
+@contextmanager
+def release_fsdp_gradients_during_validation(model):
+    """Bound unused offloaded gradient buffers during a no-grad validation pass.
+
+    Enter only between completed training steps, with gradients cleared. FSDP1
+    allocates CPU gradient buffers even under no_grad; release them after each
+    unit instead of retaining a full model's buffers until validation ends.
+    Hooks are temporary and do not affect subsequent training/backward.
+    """
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+    units = [unit for unit in model.modules() if isinstance(unit, FSDP)
+             and unit._handle is not None and unit._handle._offload_params]
+    if not units:
+        yield
+        return
+    if torch.is_grad_enabled():
+        raise RuntimeError('Validation gradient release requires torch.no_grad()')
+    release_completed_fsdp_gradients(model)
+    hooks = []
+
+    def release(unit, inputs, output):
+        if torch.is_grad_enabled():
+            raise RuntimeError('Validation gradient release cannot run with autograd enabled')
+        unit.zero_grad(set_to_none=True)
+        unit._handle.flat_param.grad = None
+        unit._handle.flat_param._cpu_grad = None
+
+    try:
+        for unit in units:
+            hooks.append(unit.register_forward_hook(release))
+        yield
+    finally:
+        for hook in hooks:
+            hook.remove()
+        model.zero_grad(set_to_none=True)
+        release_completed_fsdp_gradients(model)
+
+
+@contextmanager
+def release_fsdp_gradients_during_replay(model):
+    """Bound idle CPU gradients during checkpoint_fsdp_forward rollouts.
+
+    The caller may have autograd enabled to record the custom replay outputs,
+    but every actual model forward must run under no_grad. Exit this context
+    before backward; the replay then recreates the native gradient buffers.
+    Pending gradients are rejected, and both hooks and initializers are restored.
+    """
+    from contextlib import ExitStack
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+    def without_idle_gradients(unit, initialize):
+        def wrapped(*args, **kwargs):
+            if torch.is_grad_enabled():
+                raise RuntimeError('FSDP replay buffer release requires no-grad model forwards')
+            result = initialize(*args, **kwargs)
+            unit.zero_grad(set_to_none=True)
+            unit._handle.flat_param.grad = None
+            unit._handle.flat_param._cpu_grad = None
+            return result
+        return wrapped
+
+    def restore_initializer(handle, had_override, previous):
+        if had_override:
+            handle.init_flat_param_attributes = previous
+        else:
+            del handle.init_flat_param_attributes
+
+    with ExitStack() as stack:
+        # Only context entry bypasses the no-grad check. The installed per-unit
+        # hooks still reject autograd-enabled forwards before clearing buffers.
+        with torch.no_grad():
+            stack.enter_context(release_fsdp_gradients_during_validation(model))
+        for unit in FSDP.fsdp_modules(model):
+            handle = unit._handle
+            if handle is None or not handle._offload_params:
+                continue
+            previous = handle.init_flat_param_attributes
+            had_override = 'init_flat_param_attributes' in handle.__dict__
+            stack.callback(restore_initializer, handle, had_override, previous)
+            handle.init_flat_param_attributes = without_idle_gradients(unit, previous)
+        yield
+
+
+def configure_fsdp_pageable_cpu_parameters(model):
+    """Keep original CPU parameter shards in low-memory FSDP1 training.
+
+    Install before lazy initialization. Native FSDP still creates pinned gradient
+    buffers for asynchronous backward copies; only the parameter shard retains
+    its existing pageable storage. H2D parameter transfers may be slower.
+    """
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+    _require_cpu_offload_runtime()
+
+    def initialize_pageable(handle, initialize):
+        def wrapped(*args, **kwargs):
+            original = handle.flat_param.data
+            result = initialize(*args, **kwargs)
+            parameter = handle.flat_param
+            if original.device.type == 'cpu' and not original.is_pinned():
+                if (parameter.data_ptr() != original.data_ptr()
+                        or parameter.shape != original.shape
+                        or parameter.dtype != original.dtype):
+                    raise RuntimeError('FSDP initialization unexpectedly replaced the CPU parameter shard')
+                parameter._local_shard = original
+            return result
+        return wrapped
+
+    handles = [unit._handle for unit in FSDP.fsdp_modules(model)
+               if unit._handle is not None and unit._handle._offload_params
+               and not getattr(unit._handle, '_videox_pageable_cpu_parameters', False)]
+    if any(hasattr(handle.flat_param, '_local_shard') for handle in handles):
+        raise RuntimeError('Configure pageable FSDP parameters before lazy initialization')
+    for handle in handles:
+        handle.init_flat_param_attributes = initialize_pageable(handle, handle.init_flat_param_attributes)
+        handle._videox_pageable_cpu_parameters = True
+    return model
+
+
+def configure_fsdp_precision(model, accelerator, *, keep_in_fp32_modules=(), buffer_dtype=torch.float32):
+    """Preserve optimizer identities and apply explicit per-module precision.
+
+    Call before prepare. ``keep_in_fp32_modules`` contains exact module paths;
+    an empty sequence works for models without special precision requirements.
+    Reuse the original auto-wrap policy when configuring multiple models.
+    """
+    from dataclasses import replace
+    from torch.distributed.fsdp import MixedPrecision
+    from torch.distributed.fsdp.wrap import CustomPolicy
+
+    plugin = accelerator.state.fsdp_plugin
+    if getattr(plugin, "fsdp_version", 1) != 1:
+        raise ValueError("Per-module FSDP precision configuration requires FSDP1")
+    if getattr(plugin.cpu_offload, "offload_params", False):
+        release_offloaded_fsdp_gradients(model)
+    plugin.use_orig_params = True
+    if not hasattr(plugin, "_videox_base_wrap_policy"):
+        plugin.set_auto_wrap_policy(model)
+        plugin._videox_base_wrap_policy = plugin.auto_wrap_policy
+    base_policy = plugin._videox_base_wrap_policy
+    pinned = {module for name, module in model.named_modules() if name in keep_in_fp32_modules}
+    precision = plugin.mixed_precision_policy or MixedPrecision()
+    plugin.mixed_precision_policy = replace(precision, buffer_dtype=buffer_dtype)
+
+    def policy(module):
+        if module in pinned and any(True for _ in module.parameters()):
+            return {"mixed_precision": MixedPrecision(
+                param_dtype=torch.float32, reduce_dtype=torch.float32,
+                buffer_dtype=torch.float32, cast_forward_inputs=True,
+            )}
+        return bool(callable(base_policy) and base_policy(
+            module=module, recurse=False, nonwrapped_numel=sum(p.numel() for p in module.parameters())
+        ))
+
+    if base_policy is not None and not callable(base_policy):
+        # CustomPolicy / ModuleWrapPolicy expose the policy protocol in FSDP1.
+        class PrecisionPolicy(CustomPolicy):
+            def _run_policy(self, root_module, ignored_modules, root_kwargs):
+                result = base_policy._run_policy(root_module, ignored_modules, root_kwargs)
+                for module, kwargs in super()._run_policy(root_module, ignored_modules, root_kwargs).items():
+                    # Preserve other overrides from the existing policy.
+                    result.setdefault(module, dict(root_kwargs))["mixed_precision"] = kwargs["mixed_precision"]
+                return result
+        plugin.auto_wrap_policy = PrecisionPolicy(policy)
+    else:
+        plugin.auto_wrap_policy = CustomPolicy(policy)

videox_fun/utils/gpu_memory_mode.py

@@ -21,6 +21,8 @@ and every script picks it up by writing ``model_cpu_offload_and_<quant>``.
 
 import logging
 
+import torch
+
 from .fp4_optimization import (convert_float4_weight_dtype_wrapper,
                                convert_model_weight_to_float4)
 from .fp8_optimization import (convert_model_weight_to_float8,
@@ -36,8 +38,13 @@ GPU_MEMORY_MODES = (
     "model_cpu_offload",
     "model_group_offload",
     "sequential_cpu_offload",
+    "text_encoder_offload",
 )
 
+# Attribute prefix of the conditioner components (`text_encoder`, plus `text_encoder_2` /
+# `text_encoder_3` on the multi-encoder families).
+TEXT_ENCODER_PREFIX = "text_encoder"
+
 # Separator between the memory part and the quantization part of a mode string.
 QUANT_SEPARATOR = "_and_"
 
@@ -153,6 +160,86 @@ def split_gpu_memory_mode(mode):
     return base, tag
 
 
+def _stream_module_parameters(model, device):
+    """Onload a module's own parameters around its forward and offload them again after it.
+
+    A whole-component onload cannot work for the big conditioners: MiniMax-H3's is 63 GiB
+    (Qwen3-VL, 31.5B parameters in bf16), so even beside a float8 transformer (33 GiB) it is
+    96 GiB on an 80 GiB card. Moving every parameter-owning module on its own keeps the peak
+    at the largest block (~0.8 GiB); the total transferred volume is the same, which is what
+    makes this a placement for a component a request touches once.
+
+    Hooking each module's `forward` also means no input tensor goes through accelerate's
+    `send_to_device`, which trips over the meta-device tensors a `low_cpu_mem_usage`
+    conditioner is built from.
+    """
+    def _onload(module, args):
+        module.to(device)
+        return None  # a non-None return would replace the module inputs
+
+    def _offload(module, args, output):
+        module.to("cpu")
+        return None  # a non-None return would replace the module output
+
+    streamed = 0
+    for module in model.modules():
+        if not any(True for _ in module.parameters(recurse=False)) and not any(
+                True for _ in module.buffers(recurse=False)):
+            continue
+        module.register_forward_pre_hook(_onload)
+        module.register_forward_hook(_offload)
+        streamed += 1
+    return streamed
+
+
+def _place_without_text_encoder(pipeline, device):
+    """Keep every component on `device` except the conditioners, which stream from CPU RAM.
+
+    `model_cpu_offload` offloads all components, the transformer included -- and the transformer
+    is the one that runs once per denoising step, so its transfer dominates the request. A
+    request touches the conditioner exactly once, before the denoising loop, so leaving only it
+    in CPU RAM turns the largest component into a single one-shot transfer.
+
+    FSDP-wrapped components are skipped: `shard_model` has already placed their flat-param
+    shards, and a `.to()` would ask every rank for the unsharded parameters.
+    """
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+    conditioners = []
+    for name, component in pipeline.components.items():
+        if not isinstance(component, torch.nn.Module):
+            continue
+        if name.startswith(TEXT_ENCODER_PREFIX):
+            conditioners.append((name, component))
+            continue
+        if isinstance(component, FSDP):
+            continue
+        component.to(device)
+
+    if not conditioners:
+        raise ValueError(
+            "GPU_memory_mode 'text_encoder_offload' needs a '%s' component to offload, and "
+            "%s has none. Use 'model_full_load' or 'model_cpu_offload' for this pipeline."
+            % (TEXT_ENCODER_PREFIX, type(pipeline).__name__))
+
+    for name, component in conditioners:
+        if isinstance(component, FSDP) or isinstance(getattr(component, "model", component), FSDP):
+            # Its shards live in FSDP's hands, not in the module's own parameters, so the per-module
+            # hooks below would find nothing to move. FSDP's CPUOffload is the sharding-side
+            # equivalent and is enabled at wrap time (`shard_model(..., offload_to_cpu=True)`).
+            raise ValueError(
+                "GPU_memory_mode 'text_encoder_offload' cannot offload '%s': it is FSDP-wrapped, "
+                "and only FSDP can move its shards. Set the script's `fsdp_text_encoder = False` "
+                "so this mode can stream the conditioner, or use 'model_cpu_offload', whose "
+                "CPU-offload path shards and offloads the conditioner together." % name)
+        component.to("cpu")
+        streamed = _stream_module_parameters(component, device)
+        print(
+            "GPU_memory_mode 'text_encoder_offload': '%s' kept in CPU RAM, %d "
+            "parameter-owning modules streamed to %s for the single encode pass."
+            % (name, streamed, device))
+
+
 def _target_transformers(pipeline, transformers):
     """Modules the quantization / sequential-offload special cases apply to.
 
@@ -254,6 +341,8 @@ def apply_gpu_memory_mode(pipeline, mode, device, weight_dtype=None, *,
             offload_type="leaf_level", use_stream=True)
     elif base == "model_cpu_offload":
         pipeline.enable_model_cpu_offload(device=device)
+    elif base == "text_encoder_offload":
+        _place_without_text_encoder(pipeline, device)
     elif base == "model_full_load":
         pipeline.to(device=device)
     elif strict: