diff --git a/README.md b/README.md index fa0527a..8462273 100755 --- a/README.md +++ b/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`. diff --git a/README_ja-JP.md b/README_ja-JP.md index 1425ccd..a2ac321 100755 --- a/README_ja-JP.md +++ b/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`の使用をお勧めします。 diff --git a/README_zh-CN.md b/README_zh-CN.md index 456ce3d..196b0af 100755 --- a/README_zh-CN.md +++ b/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，速度最快，显存需求最高。 diff --git a/examples/minimax_h3/predict_i2v.py b/examples/minimax_h3/predict_i2v.py index 3a1daa0..454deb9 100644 --- a/examples/minimax_h3/predict_i2v.py +++ b/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. diff --git a/examples/minimax_h3/predict_ref2va.py b/examples/minimax_h3/predict_ref2va.py index c5c947f..987fe36 100644 --- a/examples/minimax_h3/predict_ref2va.py +++ b/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. diff --git a/examples/minimax_h3/predict_t2v.py b/examples/minimax_h3/predict_t2v.py index 410d9ee..5d1fe62 100644 --- a/examples/minimax_h3/predict_t2v.py +++ b/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. diff --git a/examples/minimax_h3_fun/predict_v2v_control.py b/examples/minimax_h3_fun/predict_v2v_control.py index 1d4594a..3392278 100644 --- a/examples/minimax_h3_fun/predict_v2v_control.py +++ b/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. diff --git a/examples/minimax_h3_fun/predict_v2v_control_inpaint.py b/examples/minimax_h3_fun/predict_v2v_control_inpaint.py index e258182..4411c13 100644 --- a/examples/minimax_h3_fun/predict_v2v_control_inpaint.py +++ b/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. diff --git a/videox_fun/__init__.py b/videox_fun/__init__.py index 6bef682..30048ff 100644 --- a/videox_fun/__init__.py +++ b/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() diff --git a/videox_fun/utils/checkpoint_utils.py b/videox_fun/utils/checkpoint_utils.py new file mode 100644 index 0000000..7d73403 --- /dev/null +++ b/videox_fun/utils/checkpoint_utils.py @@ -0,0 +1,1423 @@ +"""Model weight I/O and transactional Accelerate training checkpoints.""" +from contextlib import contextmanager +from dataclasses import dataclass, field +import functools +import gc +import hashlib +import json +import os +from pathlib import Path +import random +import re +import shutil +import tempfile +import types +import uuid +import warnings +from weakref import WeakKeyDictionary, ref + +import numpy as np +import torch +import torch.distributed as dist +from accelerate import init_empty_weights +from safetensors import safe_open +from safetensors.torch import load_file + + +def _remove_prefix(text, prefix): + """str.removeprefix backport for Python < 3.9.""" + return text[len(prefix):] if prefix and text.startswith(prefix) else text + + +def checkpoint_fingerprint(path): + """Content identity for a frozen model source, without loading tensor payloads. + + Include the index and config when present. Stream the actual weights so an + in-place replacement is detected even if filenames and sizes are unchanged. + Paths are relative, allowing an identical checkpoint to be moved. + """ + path = Path(path) + if path.is_dir(): + files = safetensor_checkpoint_files(path) + if files is None: + index = path / "diffusion_pytorch_model.bin.index.json" + files = ([path / name for name in sorted(set(json.loads(index.read_text())["weight_map"].values()))] + if index.exists() else [path / "diffusion_pytorch_model.bin"]) + files = [Path(name) for name in files] + files += [file for file in (path / "config.json", path / "diffusion_pytorch_model.safetensors.index.json", + path / "diffusion_pytorch_model.bin.index.json") if file.is_file()] + else: + files = [path] + result = [] + for file in sorted(set(files)): + before = file.stat() + digest = hashlib.sha256() + with file.open("rb") as stream: + for chunk in iter(lambda: stream.read(8 * 1024**2), b""): + digest.update(chunk) + after = file.stat() + if (before.st_size, before.st_mtime_ns, before.st_ctime_ns) != (after.st_size, after.st_mtime_ns, after.st_ctime_ns): + raise RuntimeError(f"Checkpoint source changed while hashing: {file}") + result.append({"name": file.relative_to(path).as_posix() if path.is_dir() else "weights", + "size": after.st_size, "sha256": digest.hexdigest()}) + return result + + +def save_diffusers_state_dict(state_dict, model, output_dir, max_shard_size=5 * 1024**3, *, dtype=None): + """Save weights in bounded-size standard Diffusers shards. + + Consume the gathered dictionary as shards are written, avoiding a second + full-model serialization buffer. Floating dtypes are preserved unless + ``dtype`` is supplied; integer and boolean buffers are never cast. + This consumes ``state_dict``. All ranks must gather before rank zero calls it. + """ + from safetensors.torch import save_file + if not state_dict: + raise ValueError('Cannot save an empty full transformer checkpoint') + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + groups, group_bytes, total_bytes = [[]], 0, 0 + for key, value in state_dict.items(): + size = value.numel() * (torch.empty((), dtype=dtype).element_size() if dtype is not None and value.is_floating_point() else value.element_size()) + if size > max_shard_size: + raise ValueError(f"{key} exceeds the checkpoint shard size") + if groups[-1] and group_bytes + size > max_shard_size: + groups.append([]) + group_bytes = 0 + groups[-1].append(key) + group_bytes += size + total_bytes += size + weight_map = {} + for index, keys in enumerate(groups, 1): + filename = f"diffusion_pytorch_model-{index:05d}-of-{len(groups):05d}.safetensors" + tensors = {} + for key in keys: + value = state_dict.pop(key) + tensors[key] = value.detach().to(device='cpu', dtype=dtype if dtype is not None and value.is_floating_point() else value.dtype).contiguous() + weight_map[key] = filename + save_file(tensors, output_dir / filename, metadata={'format': 'pt'}) + del tensors, value + (output_dir / 'diffusion_pytorch_model.safetensors.index.json').write_text(json.dumps( + {'metadata': {'total_size': total_bytes}, 'weight_map': weight_map}, indent=2)) + model.save_config(output_dir) + for path in output_dir.glob('diffusion_pytorch_model*.safetensors'): + if path.name not in set(weight_map.values()): + path.unlink() + + +def read_diffusers_state_dict(directory): + """Read a standard single-file or sharded safetensors checkpoint on CPU.""" + directory = Path(directory) + index = directory / 'diffusion_pytorch_model.safetensors.index.json' + weight_map = json.loads(index.read_text())['weight_map'] if index.exists() else None + names = sorted(set(weight_map.values())) if weight_map is not None else ['diffusion_pytorch_model.safetensors'] + state = {} + for name in names: + shard = load_file(directory / name) + if state.keys() & shard.keys(): + raise ValueError(f'Duplicate tensors in checkpoint shard {name}') + if weight_map is not None and any(weight_map.get(key) != name for key in shard): + raise ValueError(f'Checkpoint index does not match shard {name}') + state.update(shard) + if not state or (weight_map is not None and set(state) != set(weight_map)): + raise ValueError(f'Incomplete checkpoint index: {directory}') + return state + + +def load_diffusers_checkpoint(model_cls, model_name, checkpoint, torch_dtype, *, subfolder=None): + """Validate all tensors before loading one model; partial/LoRA files are rejected.""" + if checkpoint is None: + raise ValueError("A full model checkpoint is required") + checkpoint = Path(checkpoint) + config = model_cls.load_config(model_name, subfolder=subfolder) + with init_empty_weights(include_buffers=True): + expected_model = model_cls.from_config(config) + expected = {key: tuple(value.shape) for key, value in expected_model.state_dict().items()} + del expected_model + weight_map = None + if checkpoint.is_dir(): + if not (checkpoint / "config.json").is_file(): + raise ValueError(f"Missing transformer config: {checkpoint}") + index = checkpoint / "diffusion_pytorch_model.safetensors.index.json" + if index.is_file(): + weight_map = json.loads(index.read_text())["weight_map"] + files = [checkpoint / name for name in sorted(set(weight_map.values()))] + else: + files = [checkpoint / "diffusion_pytorch_model.safetensors"] + else: + files = [checkpoint] + found = {} + state = None + for path in files: + if not path.is_file(): + raise FileNotFoundError(f"Missing checkpoint weights: {path}") + if path.suffix == ".safetensors": + with safe_open(path, framework="pt", device="cpu") as reader: + shapes = {key: tuple(reader.get_slice(key).get_shape()) for key in reader.keys()} + else: + state = torch.load(path, map_location="cpu", weights_only=True) + state = state.get("state_dict", state) + shapes = {key: tuple(value.shape) for key, value in state.items()} + for key, shape in shapes.items(): + if key in found: + raise ValueError(f"Duplicate checkpoint tensor: {key}") + if weight_map is not None and weight_map.get(key) != path.name: + raise ValueError(f"Checkpoint index does not match shard for {key}") + found[key] = shape + missing = sorted(expected.keys() - found.keys()) + unexpected = sorted(found.keys() - expected.keys()) + mismatched = [key for key in expected.keys() & found.keys() if expected[key] != found[key]] + if missing or unexpected or mismatched or (weight_map is not None and set(weight_map) != set(found)): + raise ValueError( + f"Incomplete or incompatible full checkpoint {checkpoint}: " + f"missing={missing[:5]}, unexpected={unexpected[:5]}, shape_mismatch={mismatched[:5]}" + ) + if checkpoint.is_dir(): + model, info = model_cls.from_pretrained( + checkpoint, low_cpu_mem_usage=True, torch_dtype=torch_dtype, output_loading_info=True, + ) + if any(info.get(key) for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")): + raise ValueError(f"Checkpoint could not be loaded completely: {info}") + return model + # A standalone file needs the base config/precision contract. Only this path + # loads base weights; directory checkpoints never instantiate a base model. + model = model_cls.from_pretrained( + model_name, subfolder=subfolder, low_cpu_mem_usage=True, torch_dtype=torch_dtype, + ) + if state is None: + state = load_file(str(checkpoint)) + model.load_state_dict(state, strict=True) + return model + + +def safetensor_checkpoint_files(path): + path = os.fspath(path) + if not os.path.isdir(path): + return [path] if path.endswith(".safetensors") else None + index_path = os.path.join(path, "diffusion_pytorch_model.safetensors.index.json") + if os.path.isfile(index_path): + with open(index_path) as reader: + filenames = sorted(set(json.load(reader)["weight_map"].values())) + return [os.path.join(path, filename) for filename in filenames] + single_file = os.path.join(path, "diffusion_pytorch_model.safetensors") + if os.path.isfile(single_file): + return [single_file] + files = sorted(os.path.join(path, name) for name in os.listdir(path) if name.endswith(".safetensors")) + return files or None + + +def is_complete_diffusers_checkpoint(model_cls, path, base_config): + from accelerate import init_empty_weights + from safetensors import safe_open + + if not os.path.isdir(path) or not os.path.isfile(os.path.join(path, "config.json")): + return False + files = safetensor_checkpoint_files(path) + if files is None: + return False + with init_empty_weights(): + expected_model = model_cls.from_config(base_config) + candidate_model = model_cls.from_config(model_cls.load_config(path)) + configs = [dict(model.config) for model in (expected_model, candidate_model)] + if json.dumps({k: v for k, v in configs[0].items() if not k.startswith("_")}, sort_keys=True) != json.dumps( + {k: v for k, v in configs[1].items() if not k.startswith("_")}, sort_keys=True + ): + return False + expected = {key: tuple(value.shape) for key, value in expected_model.state_dict().items()} + found = {} + for filename in files: + with safe_open(filename, framework="pt", device="cpu") as reader: + for key in reader.keys(): + if key in found: + return False + found[key] = tuple(reader.get_slice(key).get_shape()) + return found == expected + + +def copy_diffusers_checkpoint(model, path, weight_dtype, *, strip_prefix=None, state_dict_key=None): + from safetensors import safe_open + + files = safetensor_checkpoint_files(path) + if files is None: + if os.path.isdir(path): + candidate = type(model).from_pretrained(path, low_cpu_mem_usage=True, torch_dtype=weight_dtype) + state_dict = candidate.state_dict() + del candidate + else: + state_dict = torch.load(path, map_location="cpu", weights_only=True) + state_dict = state_dict.get("state_dict", state_dict) + if state_dict_key is not None: + state_dict = state_dict.get(state_dict_key, state_dict) + if strip_prefix: + state_dict = {_remove_prefix(key, strip_prefix): value for key, value in state_dict.items()} + missing, unexpected = model.load_state_dict(state_dict, strict=True) + del state_dict + else: + destination = model.state_dict() + found_shapes = {} + for filename in files: + with safe_open(filename, framework="pt", device="cpu") as reader: + for source_key in reader.keys(): + key = _remove_prefix(source_key, strip_prefix) if strip_prefix else source_key + if key in found_shapes: + raise ValueError(f"Duplicate transformer key: {key}") + found_shapes[key] = tuple(reader.get_slice(source_key).get_shape()) + if found_shapes != {key: tuple(value.shape) for key, value in destination.items()}: + raise ValueError(f"Incomplete or incompatible full transformer checkpoint: {path}") + seen, unexpected = set(), [] + with torch.no_grad(): + for filename in files: + with safe_open(filename, framework="pt", device="cpu") as reader: + for source_key in reader.keys(): + key = _remove_prefix(source_key, strip_prefix) if strip_prefix else source_key + if key in seen: + raise ValueError(f"Duplicate transformer key: {key}") + seen.add(key) + if key not in destination: + unexpected.append(key) + continue + value = reader.get_tensor(source_key) + if value.shape != destination[key].shape: + raise ValueError(f"{key}: checkpoint shape {value.shape} != model shape {destination[key].shape}") + destination[key].copy_(value) + del value + missing = sorted(set(destination) - seen) + if unexpected: + raise ValueError(f"Unexpected transformer keys in {path}: {unexpected[:3]}") + gc.collect() + + +def load_model_checkpoint(model, checkpoint): + """Load a full local checkpoint into an existing model, one shard at a time. + + Validate all keys, shapes and index entries before modifying parameters. + Supports standalone PyTorch/safetensors files and Diffusers directories. + Uses model.load_state_dict to preserve model-specific loading hooks and + returns (missing_keys, unexpected_keys), as expected by training callers. + """ + checkpoint = Path(checkpoint) + weight_map = None + if checkpoint.is_dir(): + index = checkpoint / "diffusion_pytorch_model.safetensors.index.json" + if index.is_file(): + weight_map = json.loads(index.read_text())["weight_map"] + if not weight_map: + raise ValueError(f"Empty checkpoint weight map: {index}") + files = [checkpoint / name for name in sorted(set(weight_map.values()))] + else: + files = [checkpoint / "diffusion_pytorch_model.safetensors"] + else: + files = [checkpoint] + for path in files: + if not path.is_file(): + raise FileNotFoundError(f"Missing transformer checkpoint shard: {path}") + expected_shapes = {key: tuple(value.shape) for key, value in model.state_dict().items()} + found_shapes = {} + for path in files: + if path.suffix == ".safetensors": + from safetensors import safe_open + with safe_open(str(path), framework="pt", device="cpu") as reader: + shapes = {key: tuple(reader.get_slice(key).get_shape()) for key in reader.keys()} + else: + state = torch.load(path, map_location="cpu", weights_only=True) + state = state.get("state_dict", state) + shapes = {key: tuple(value.shape) for key, value in state.items()} + del state + for key, shape in shapes.items(): + if key in found_shapes: + raise ValueError(f"Duplicate checkpoint tensor: {key}") + if weight_map is not None and weight_map.get(key) != path.name: + raise ValueError(f"Checkpoint index does not match shard for {key}") + found_shapes[key] = shape + if found_shapes != expected_shapes or (weight_map is not None and set(weight_map) != set(found_shapes)): + raise ValueError(f"Incomplete or incompatible full transformer checkpoint: {checkpoint}") + expected = set(expected_shapes) + loaded = set() + unexpected = set() + for path in files: + if path.suffix == ".safetensors": + from safetensors.torch import load_file + state = load_file(str(path)) + else: + state = torch.load(path, map_location="cpu", weights_only=True) + state = state.get("state_dict", state) + if weight_map is not None: + shard_keys = {key for key, name in weight_map.items() if name == path.name} + if not shard_keys.issubset(state): + raise ValueError(f"Checkpoint index references missing tensors in {path}") + _, extra = model.load_state_dict(state, strict=False) + loaded.update(state) + unexpected.update(extra) + del state + return sorted(expected - loaded), sorted(unexpected) + + +MANIFEST = "checkpoint_manifest.json" +PROTOCOL = 1 +_CONFIGS = WeakKeyDictionary() +_CHECKPOINT_NAME = re.compile(r"checkpoint[-_](\d+)\Z") + + +@dataclass +class CheckpointConfig: + model_ids: WeakKeyDictionary = field(default_factory=WeakKeyDictionary) + allow_legacy: bool = True + check_finite: bool = False + checksum: str = "metadata" + private_states: dict = field(default_factory=dict) + validators: dict = field(default_factory=dict) + run_id: str = None + metadata: dict = field(default_factory=dict) + legacy_metadata_keys: tuple = () + legacy_metadata_groups: tuple = () + step_guards: list = field(default_factory=list) + + +def checkpoint_config(owner): + if owner not in _CONFIGS: + _CONFIGS[owner] = CheckpointConfig( + check_finite=os.environ.get("VIDEOX_CHECKPOINT_FINITE", "0") == "1", + ) + return _CONFIGS[owner] + + +def configure_checkpointing(owner, *, model_ids=None, allow_legacy=None, check_finite=None, checksum=None, metadata=None, + legacy_metadata_keys=(), legacy_metadata_groups=()): + """Configure stable model roles and explicit compatibility/integrity options. + + ``model_ids`` maps role names to prepared models. It is required for multiple + models of the same class when saving new checkpoints. Legacy checkpoints + are detected automatically; ``allow_legacy=False`` opts into manifest-only + restore. Missing required training state is an error in either format. + ``checksum='full'`` adds streaming hashes of large payloads; the default + hashes control files and records the size of every payload. + ``legacy_metadata_keys`` names fields newly added by the caller. Older + manifests missing that entire set can load with a warning. A partially + missing set is invalid; existing values must still match. No model, + optimizer or topology checks are relaxed. + ``legacy_metadata_groups`` allows separate generations of fields to be + migrated independently, each subject to the same all-or-none check. + """ + config = checkpoint_config(owner) + if getattr(owner, "_models", None): + validate_checkpoint_owner(owner) + if model_ids is not None: + if not all(isinstance(name, str) and name for name in model_ids): + raise ValueError("Checkpoint model IDs must be nonempty strings") + if len(set(map(id, model_ids.values()))) != len(model_ids): + raise ValueError("Each checkpoint model must have exactly one ID") + config.model_ids = WeakKeyDictionary((model, name) for name, model in model_ids.items()) + if metadata is not None: + config.metadata = _plain(metadata) + config.legacy_metadata_keys = tuple(legacy_metadata_keys) + config.legacy_metadata_groups = tuple(tuple(group) for group in legacy_metadata_groups) + if any(key not in config.metadata for group in (config.legacy_metadata_keys, *config.legacy_metadata_groups) for key in group): + raise ValueError("Legacy metadata keys must exist in the current training metadata") + if allow_legacy is not None: + config.allow_legacy = bool(allow_legacy) + if check_finite is not None: + config.check_finite = bool(check_finite) + if checksum is not None: + if checksum not in ("metadata", "full"): + raise ValueError("Checkpoint checksum must be 'metadata' or 'full'") + config.checksum = checksum + return owner + + +def checkpoint_phase(function): + """Synchronize errors from a local phase, including native DCP exceptions.""" + result, failure = None, None + try: + result = function() + except BaseException as exc: + # DCP CheckpointException derives directly from BaseException. + failure = f"{type(exc).__name__}: {exc}" + failures = [failure] + if dist.is_initialized(): + failures = [None] * dist.get_world_size() + dist.all_gather_object(failures, failure) + if any(failures): + detail = "; ".join(f"rank {rank}: {error}" for rank, error in enumerate(failures) if error) + raise RuntimeError("Training checkpoint failed: " + detail) + return result + + +def _agree(value): + if dist.is_initialized(): + values = [None] * dist.get_world_size() + dist.all_gather_object(values, value) + if any(other != values[0] for other in values[1:]): + raise RuntimeError("Checkpoint configuration differs between ranks") + + +def _main_value(function): + value = checkpoint_phase(lambda: function() if not dist.is_initialized() or dist.get_rank() == 0 else None) + if dist.is_initialized(): + values = [value] + dist.broadcast_object_list(values, src=0) + value = values[0] + return value + + +def checkpoint_source_identity(path): + """Hash frozen checkpoint inputs once on rank zero and share their identity.""" + + _agree(os.fspath(path)) + return _main_value(lambda: checkpoint_fingerprint(path)) + + +def _class_name(value): + return type(value).__module__ + "." + type(value).__qualname__ + + +def _plain(value): + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, (tuple, list)): + return [_plain(item) for item in value] + if isinstance(value, dict): + return {str(key): _plain(item) for key, item in sorted(value.items(), key=lambda pair: str(pair[0]))} + if isinstance(value, torch.Tensor) and value.numel() == 1: + return value.item() + raise ValueError(f"Unsupported checkpoint configuration value: {_class_name(value)}") + + +def _callable_signature(function): + if isinstance(function, functools.partial): + return {"partial": _callable_signature(function.func), "args": _plain(function.args), + "kwargs": _plain(function.keywords)} + if isinstance(function, types.FunctionType): + return {"function": function.__name__, "defaults": _plain(function.__defaults__), + "keywords": _plain(function.__kwdefaults__), + "closure": [_plain(cell.cell_contents) for cell in function.__closure__ or ()]} + # Callable scheduler objects may carry arbitrary mutable behavior. Require + # an explicit implementation instead of claiming generic compatibility. + raise ValueError("Unsupported scheduler callable; use a supported function or functools.partial") + + +def _normalize_scheduler_callable(signature): + """Read older signatures without tying compatibility to module or bytecode.""" + signature = dict(signature) + if "partial" in signature: + signature["partial"] = _normalize_scheduler_callable(signature["partial"]) + if "function" in signature: + signature["function"] = signature["function"].rsplit(".", 1)[-1] + signature.pop("code", None) + return signature + + +def _hook_counts(owners): + # The synchronized wrapper invokes one collective per registered hook. + # Agree on their counts across ranks, not their Python implementation. + return {name: [len(owner._save_model_state_pre_hook), len(owner._load_model_state_pre_hook)] + for name, owner in owners.items()} + + +def _raw_optimizer(optimizer): + return getattr(optimizer, "optimizer", optimizer) + + +def validate_checkpoint_owner(owner): + from .. import _supported_accelerate_versions + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy + + if not _supported_accelerate_versions(): + raise ValueError("Protected checkpoints require Accelerate 1.7.0 / PyTorch 2.7.0") + mode = owner.distributed_type.name + if mode not in ("FSDP", "NO", "MULTI_GPU"): + raise ValueError(f"Protected checkpoint backend is not validated: {mode}") + if not owner._models or len(owner._models) != len(owner._optimizers): + raise ValueError("Protected checkpoints require one optimizer per prepared model") + if owner.project_configuration.save_on_each_node: + raise ValueError("Protected checkpoints require one shared checkpoint filesystem") + if any(type(_raw_optimizer(optimizer)) not in (torch.optim.Adam, torch.optim.AdamW, torch.optim.SGD) + for optimizer in owner._optimizers): + raise ValueError("Protected checkpoints support native Adam, AdamW and SGD; this optimizer is not validated") + if mode == "FSDP": + plugin = owner.state.fsdp_plugin + if getattr(plugin, "fsdp_version", 1) != 1 or plugin.state_dict_type.name != "SHARDED_STATE_DICT": + raise ValueError("Protected FSDP checkpoints require FSDP1 SHARDED_STATE_DICT") + for model in owner._models: + for unit in FSDP.fsdp_modules(model): + if (unit.process_group is not model.process_group or unit._device_mesh is not None + or dist.get_world_size(unit.process_group) != owner.num_processes + or unit.sharding_strategy not in (ShardingStrategy.FULL_SHARD, ShardingStrategy.SHARD_GRAD_OP)): + raise ValueError("Only ordinary FSDP1 sharding over the world group is validated") + if unit.cpu_offload.offload_params and not unit._use_orig_params: + raise ValueError("CPU offload checkpoints require use_orig_params=True for reliable continuation") + + +def _descriptor(owner, *, require_model_ids=True): + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + + validate_checkpoint_owner(owner) + config = checkpoint_config(owner) + models, optimizers = [], [] + if config.model_ids and set(config.model_ids) != set(owner._models): + raise ValueError("Checkpoint model_ids must identify every currently prepared model") + classes = [_class_name(owner.unwrap_model(model, keep_torch_compile=False)) for model in owner._models] + for index, (model, wrapped_optimizer) in enumerate(zip(owner._models, owner._optimizers)): + optimizer = _raw_optimizer(wrapped_optimizer) + if owner.distributed_type.name == "FSDP": + groups, shapes = _optimizer_layout(model, optimizer) + from torch.distributed.fsdp._common_utils import _get_param_to_fqns + names = _get_param_to_fqns(model) + units = [{"original_params": unit._use_orig_params, "cpu_offload": unit.cpu_offload.offload_params, + "sharding": unit.sharding_strategy.name} for unit in FSDP.fsdp_modules(model)] + else: + names = {parameter: [name] for name, parameter in model.named_parameters()} + shapes = {name: parameter.shape for name, parameter in model.named_parameters()} + groups = [[name for parameter in group["params"] for name in names[parameter]] for group in optimizer.param_groups] + units = [] + role = config.model_ids.get(model) + if role is None: + if classes.count(classes[index]) > 1 and require_model_ids: + raise ValueError("Multiple models of the same class require configure_checkpointing(model_ids={role: model})") + role = (f"legacy_model_{index}" if classes.count(classes[index]) > 1 + else "model" if len(classes) == 1 else classes[index]) + parameters = {name: {"shape": list(shapes[name]), "dtype": str(parameter.dtype), + "requires_grad": parameter.requires_grad} + for parameter, fqns in names.items() for name in fqns} + unwrapped = owner.unwrap_model(model, keep_torch_compile=False) + model_config = getattr(unwrapped, "config", {}) + model_config = {key: value for key, value in dict(model_config).items() if not key.startswith("_")} if hasattr(model_config, "items") else {} + models.append({"id": role, "class": classes[index], "parameters": parameters, + "config": _plain(model_config), "fsdp": units}) + optimizers.append({"class": _class_name(optimizer), "groups": groups, + "execution": [{key: _plain(group.get(key, False)) for key in + ("capturable", "fused", "differentiable")} for group in optimizer.param_groups]}) + if len({model["id"] for model in models}) != len(models): + raise ValueError("Checkpoint model IDs must be unique") + schedulers = [] + for wrapped in owner._schedulers: + scheduler = getattr(wrapped, "scheduler", wrapped) + index = next((i for i, optimizer in enumerate(owner._optimizers) + if scheduler.optimizer in (optimizer, _raw_optimizer(optimizer))), None) + if index is None: + raise ValueError("Scheduler is not bound to a prepared optimizer") + row = {"class": type(scheduler).__name__, "optimizer": index} + if hasattr(scheduler, "lr_lambdas"): + row["functions"] = [_callable_signature(function) for function in scheduler.lr_lambdas] + schedulers.append(row) + return {"world_size": owner.num_processes, "backend": owner.distributed_type.name, + "mixed_precision": owner.mixed_precision, "accumulation_steps": owner.gradient_accumulation_steps, + "models": models, "optimizers": optimizers, "schedulers": schedulers, + "custom_objects": [_class_name(obj) for obj in owner._custom_objects], + "scaler": _class_name(owner.scaler) if owner.scaler else None, + "private_states": config.private_states, "training": config.metadata, + "dataloaders": [{"class": _class_name(loader), + "dataset": _class_name(getattr(loader.dataset, "checkpoint_dataset", loader.dataset)), + "workers": loader.num_workers, "stateful": getattr(loader, "use_stateful_dataloader", False)} + for loader in owner._dataloaders]} + + +def _relative(name): + path = Path(name) + if path.is_absolute() or ".." in path.parts: + raise ValueError("Checkpoint paths must be relative and remain inside the checkpoint") + return path + + +def _participants(owners): + owners = owners if isinstance(owners, dict) else {"": owners} + result = {} + for name, owner in sorted(owners.items()): + path = _relative(name) + key = "" if str(path) == "." else path.as_posix() + if key in result or owner in result.values(): + raise ValueError("Checkpoint components and Accelerators must be unique") + if any(other and (path.is_relative_to(other) or Path(other).is_relative_to(path)) for other in result): + raise ValueError("Non-root checkpoint components must not overlap") + result[key] = owner + if not result: + raise ValueError("At least one checkpoint participant is required") + return result + + +def _sha(path): + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _metadata_schema(path): + from torch.distributed.checkpoint import FileSystemReader, TensorStorageMetadata + + metadata = FileSystemReader(path.parent).read_metadata() + rows = {} + for key, value in metadata.state_dict_metadata.items(): + row = {"path": list(metadata.planner_data[key])} + if isinstance(value, TensorStorageMetadata): + row.update(shape=list(value.size), dtype=str(value.properties.dtype)) + rows[key] = row + return rows + + +def _inventory(directory, checksum): + result = {} + for path in sorted(directory.rglob("*")): + if path.is_symlink(): + raise ValueError(f"Checkpoint cannot contain symbolic links: {path}") + if not path.is_file() or path.name == MANIFEST: + continue + row = {"size": path.stat().st_size} + # Large native payloads retain bounded I/O by default. All control files, + # including DCP metadata and RNG/scheduler/custom states, are hashed. + if checksum == "full" or path.suffix not in (".distcp", ".safetensors", ".bin") or path.stat().st_size <= 16 * 1024 * 1024: + row["sha256"] = _sha(path) + if path.name == ".metadata": + row["schema"] = _metadata_schema(path) + result[path.relative_to(directory).as_posix()] = row + return result + + +def _verify_inventory(directory, manifest, *, full=True): + if manifest.get("protocol") != PROTOCOL or manifest.get("status") != "complete": + raise ValueError("Checkpoint is not a committed supported checkpoint") + if not all(isinstance(manifest.get(key), str) and manifest[key] for key in ("run_id", "transaction_id")): + raise ValueError("Checkpoint identity is missing") + for name, expected in manifest["files"].items(): + path = directory / _relative(name) + if path.is_symlink() or not path.is_file() or path.stat().st_size != expected["size"]: + raise ValueError(f"Checkpoint file missing or changed: {name}") + if "sha256" in expected and (full or "schema" in expected or expected["size"] <= 16 * 1024 * 1024): + if _sha(path) != expected["sha256"]: + raise ValueError(f"Checkpoint checksum mismatch: {name}") + if "schema" in expected and _metadata_schema(path) != expected["schema"]: + raise ValueError(f"Checkpoint state schema changed: {name}") + + +def latest_training_checkpoint(output_dir, owners=None): + """Select a complete checkpoint; pass owners to also inspect legacy formats. + + Legacy files cannot describe required models, components or private state, + so their requirements come from the prepared Accelerator(s). Inspect every + rank's files so all callers select the same checkpoint. This function uses + no collectives; distributed callers synchronize errors at the call site. + """ + participants = _participants(owners) if owners is not None else None + root = Path(output_dir) + if not root.is_dir(): + return None + candidates = [] + for path in root.iterdir(): + match = _CHECKPOINT_NAME.fullmatch(path.name) + if match and path.is_dir(): + candidates.append((int(match.group(1)), path)) + candidates.sort() + for _, path in reversed(candidates): + try: + manifest_path = path / MANIFEST + if manifest_path.exists(): + manifest = json.loads(manifest_path.read_text()) + _verify_inventory(path, manifest, full=False) + if participants is not None and set(manifest["components"]) != set(participants): + raise ValueError("Checkpoint participants differ") + elif participants is None or not all(checkpoint_config(owner).allow_legacy for owner in participants.values()): + continue + if participants is not None: + for name, owner in participants.items(): + _required_files(owner, path / name, all_ranks=True) + return str(path) + except Exception: + # A broken candidate must not prevent trying the preceding one. + # Direct load still reports the actual validation error. + continue + if candidates: + raise ValueError("No complete checkpoint matching the required training state") + return None + + +def _rng_state(path): + state = torch.load(path, map_location="cpu", weights_only=False) + required = {"step", "random_state", "numpy_random_seed", "torch_manual_seed"} + if torch.cuda.is_available(): + required.add("torch_cuda_manual_seed") + if len(state.get("torch_cuda_manual_seed", ())) != torch.cuda.device_count(): + raise ValueError("Checkpoint CUDA RNG topology does not match this process") + if required - state.keys(): + raise ValueError("Checkpoint RNG state is incomplete") + # Validate using the public setters, then restore the caller's RNG exactly. + previous = {"random_state": random.getstate(), "numpy_random_seed": np.random.get_state(), + "torch_manual_seed": torch.get_rng_state()} + if torch.cuda.is_available(): + previous["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all() + try: + _apply_rng(state) + finally: + _apply_rng(previous) + return state + + +def _apply_rng(state): + random.setstate(state["random_state"]) + np.random.set_state(state["numpy_random_seed"]) + torch.set_rng_state(state["torch_manual_seed"]) + if "torch_cuda_manual_seed" in state: + torch.cuda.set_rng_state_all(state["torch_cuda_manual_seed"]) + + +def _required_files(owner, directory, *, all_ranks=False): + if owner.distributed_type.name == "FSDP": + for index in range(len(owner._models)): + for stem in ("pytorch_model_fsdp", "optimizer"): + path = directory / f"{stem}_{index}" / ".metadata" + if not path.is_file(): + raise ValueError(f"Missing native FSDP state: {path}") + from torch.distributed.checkpoint import FileSystemReader + metadata = FileSystemReader(path.parent).read_metadata() + for storage in metadata.storage_data.values(): + payload = path.parent / _relative(storage.relative_path) + if not payload.is_file() or payload.stat().st_size < storage.offset + storage.length: + raise ValueError(f"Missing or truncated DCP payload: {payload}") + for index in range(len(owner._schedulers)): + name = "scheduler.bin" if index == 0 else f"scheduler_{index}.bin" + if not (directory / name).is_file(): + raise ValueError(f"Missing scheduler state: {name}") + if owner.scaler is not None and not (directory / "scaler.pt").is_file(): + raise ValueError("Missing scaler state") + for index in range(len(owner._custom_objects)): + if not (directory / f"custom_checkpoint_{index}.pkl").is_file(): + raise ValueError("Missing registered checkpoint state") + ranks = range(owner.num_processes) if all_ranks else (owner.process_index,) + current_rng = None + for rank in ranks: + for validator in checkpoint_config(owner).validators.values(): + validator(directory, rank=rank) + rng = _rng_state(directory / f"random_states_{rank}.pkl") + if rank == owner.process_index: + current_rng = rng + return current_rng + + +def _tensor_values(value): + if isinstance(value, torch.Tensor): + yield value + elif isinstance(value, dict): + for item in value.values(): + yield from _tensor_values(item) + elif isinstance(value, (tuple, list)): + for item in value: + yield from _tensor_values(item) + + +def check_finite_state(value): + """Bound temporary validation buffers while checking local tensor state.""" + from torch.distributed._shard.sharded_tensor import ShardedTensor + + for tensor in _tensor_values(value): + tensors = [shard.tensor for shard in tensor.local_shards()] if isinstance(tensor, ShardedTensor) else [tensor] + for local in tensors: + if not local.is_floating_point() and not local.is_complex(): + continue + for chunk in local.detach().reshape(-1).split(1024 * 1024): + if not torch.isfinite(chunk).all().item(): + raise ValueError("Checkpoint contains NaN or Inf") + + +def _set_failed(owners, failed): + from .. import _failed_checkpoint_owners + + for owner in owners: + config = checkpoint_config(owner) + if failed: + _failed_checkpoint_owners.add(owner) + if not config.step_guards: + owner_ref = ref(owner) + def guard(optimizer, args, kwargs, owner_ref=owner_ref): + if owner_ref() in _failed_checkpoint_owners: + raise RuntimeError("Checkpoint restore failed; reload successfully before updating parameters") + for optimizer in owner._optimizers: + config.step_guards.append(_raw_optimizer(optimizer).register_step_pre_hook(guard)) + else: + _failed_checkpoint_owners.discard(owner) + for handle in config.step_guards: + handle.remove() + config.step_guards.clear() + + +def _native(owner, method, *args, **kwargs): + function = getattr(type(owner), method) + function = getattr(function, "_videox_native_checkpoint", function) + return function(owner, *args, **kwargs) + + +@contextmanager +def _synchronized_hooks(owner): + """Make local hook failures visible before the next native phase.""" + changes = [] + for hooks in (owner._save_model_state_pre_hook, owner._load_model_state_pre_hook): + for key, function in list(hooks.items()): + def synchronized(*args, function=function): + return checkpoint_phase(lambda: function(*args)) + changes.append((hooks, key, function)) + hooks[key] = synchronized + try: + yield + finally: + for hooks, key, function in changes: + hooks[key] = function + + +@contextmanager +def _optimizer_adapter(owner): + from .. import (enable_accelerate_patch, _accelerate_patch_owners, + _accelerate_patch_disabled_owners) + + explicit = owner in _accelerate_patch_owners + disabled = owner in _accelerate_patch_disabled_owners + if owner.distributed_type.name == "FSDP": + enable_accelerate_patch(owner) + try: + yield + finally: + if not explicit: + _accelerate_patch_owners.discard(owner) + if disabled: + _accelerate_patch_disabled_owners.add(owner) + + +def _save_native(owner, path, safe_serialization, kwargs): + project = owner.project_configuration + automatic = project.automatic_checkpoint_naming + project.automatic_checkpoint_naming = False + try: + with _synchronized_hooks(owner): + return _native(owner, "save_state", str(path), safe_serialization=safe_serialization, **kwargs) + finally: + project.automatic_checkpoint_naming = automatic + + +def _save_preflight(owners): + from .. import _failed_checkpoint_owners + + descriptions = {} + for name, owner in owners.items(): + if owner in _failed_checkpoint_owners: + raise ValueError("Cannot save after a failed restore; reload successfully first") + descriptions[name] = _descriptor(owner) + if not owner.sync_gradients or any(parameter.grad is not None for model in owner._models for parameter in model.parameters()): + raise ValueError("Save checkpoints after an optimizer update and zero_grad(set_to_none=True)") + if checkpoint_config(owner).check_finite: + check_finite_state([list(model.parameters()) for model in owner._models]) + check_finite_state([_raw_optimizer(optimizer).state for optimizer in owner._optimizers]) + return descriptions + + +def save_training_checkpoint(owners, output_dir=None, *, checkpoint_limit=None, safe_serialization=True, **kwargs): + """Atomically publish one Accelerator or a named group of Accelerators.""" + owners = checkpoint_phase(lambda: _participants(owners)) + primary = next(iter(owners.values())) + if primary.project_configuration.automatic_checkpoint_naming and output_dir is None: + output_dir = Path(primary.project_dir) / "checkpoints" / f"checkpoint_{primary.save_iteration}" + if checkpoint_limit is None: + checkpoint_limit = primary.project_configuration.total_limit + path = checkpoint_phase(lambda: Path(output_dir).expanduser().absolute()) + descriptions = checkpoint_phase(lambda: _save_preflight(owners)) + policies = {name: {"finite": checkpoint_config(owner).check_finite, "checksum": checkpoint_config(owner).checksum} + for name, owner in owners.items()} + _agree([str(path), descriptions, policies, checkpoint_limit, _hook_counts(owners)]) + if checkpoint_limit is not None and checkpoint_limit < 1: + raise ValueError("checkpoint_limit must be at least 1") + run_ids = {checkpoint_config(owner).run_id for owner in owners.values()} - {None} + _agree(sorted(run_ids)) + if len(run_ids) > 1: + raise ValueError("Checkpoint participants belong to different training runs") + run_id = _main_value(lambda: next(iter(run_ids), None) or uuid.uuid4().hex) + lock, stage = None, None + iterations = {owner: owner.project_configuration.iteration for owner in owners.values()} + + def reserve(): + path.parent.mkdir(parents=True, exist_ok=True) + lock_path = path.parent / ".training-checkpoints.lock" + descriptor = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600) + os.close(descriptor) + try: + if path.exists(): + raise ValueError(f"Refusing to overwrite an existing checkpoint: {path}") + for sibling in path.parent.iterdir(): + if _CHECKPOINT_NAME.fullmatch(sibling.name) and (sibling / MANIFEST).is_file(): + other = json.loads((sibling / MANIFEST).read_text()) + if other.get("run_id") != run_id: + raise ValueError("Output directory contains checkpoints from another run; resume that run or use a new output directory") + stage_path = tempfile.mkdtemp(prefix=f".{path.name}.staging-", dir=path.parent) + return str(lock_path), stage_path + except BaseException: + lock_path.unlink() + raise + + lock, stage = _main_value(reserve) + try: + for name, owner in owners.items(): + checkpoint_phase(lambda owner=owner, name=name: _save_native(owner, Path(stage) / name, safe_serialization, kwargs)) + checkpoint_phase(lambda owner=owner, name=name: _required_files(owner, Path(stage) / name)) + + def publish(): + directory = Path(stage) + checksum = "full" if any(checkpoint_config(owner).checksum == "full" for owner in owners.values()) else "metadata" + manifest = {"protocol": PROTOCOL, "status": "complete", "run_id": run_id, + "transaction_id": uuid.uuid4().hex, "components": descriptions, + "versions": {"accelerate": "1.7.0", "torch": "2.7.0"}, + "integrity": checksum, "files": _inventory(directory, checksum)} + with (directory / MANIFEST).open("x") as stream: + json.dump(manifest, stream, sort_keys=True) + stream.flush() + os.fsync(stream.fileno()) + if path.exists(): + raise ValueError("Checkpoint destination appeared while saving") + os.rename(directory, path) + descriptor = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + _main_value(publish) + for owner in owners.values(): + checkpoint_config(owner).run_id = run_id + if checkpoint_limit is not None: + def prune(): + checkpoints = [] + for candidate in path.parent.iterdir(): + match = _CHECKPOINT_NAME.fullmatch(candidate.name) + if match and (candidate / MANIFEST).is_file(): + checkpoints.append((int(match.group(1)), candidate)) + checkpoints.sort() + for _, candidate in checkpoints[:-checkpoint_limit]: + manifest = json.loads((candidate / MANIFEST).read_text()) + if candidate != path and manifest.get("run_id") == run_id: + shutil.rmtree(candidate) + _main_value(prune) + except BaseException: + # Native save increments this counter even before a later component fails. + if not path.exists(): + for owner, iteration in iterations.items(): + owner.project_configuration.iteration = iteration + raise + finally: + def cleanup(): + if stage and Path(stage).exists(): + shutil.rmtree(stage) + if lock: + try: + Path(lock).unlink() + except FileNotFoundError: + pass + _main_value(cleanup) + return path + + +def _compare_descriptor(saved, current, *, legacy_metadata_keys=(), legacy_metadata_groups=()): + # Preserve the previously supported no-scheduler group reordering. Group + # membership changes remain invalid; scheduler-bound group order is exact. + saved, current = json.loads(json.dumps(saved)), json.loads(json.dumps(current)) + for descriptor in (saved, current): + descriptor.pop("hooks", None) # Legacy manifests recorded Python hook identities. + for state in descriptor.get("private_states", {}).values(): + # The existing RNG/progress schema predates explicit versioning. + state.setdefault("schema_version", 1) + for scheduler in descriptor["schedulers"]: + scheduler["class"] = scheduler["class"].rsplit(".", 1)[-1] + if "functions" in scheduler: + scheduler["functions"] = [_normalize_scheduler_callable(fn) for fn in scheduler["functions"]] + missing = [] + for group in (legacy_metadata_keys, *legacy_metadata_groups): + absent = sorted(set(group) - saved["training"].keys()) + if absent and set(group) & saved["training"].keys(): + raise ValueError("Checkpoint training metadata is incomplete; refusing partial legacy metadata migration") + missing.extend(absent) + for key in missing: + saved["training"][key] = current["training"][key] + if saved["backend"] == current["backend"] == "FSDP" and not saved["schedulers"] and not current["schedulers"]: + for descriptor in (saved, current): + for optimizer in descriptor["optimizers"]: + pairs = sorted(zip(optimizer["groups"], optimizer["execution"]), key=lambda pair: sorted(pair[0])) + optimizer["groups"] = [sorted(group) for group, _ in pairs] + optimizer["execution"] = [execution for _, execution in pairs] + if saved != current: + raise ValueError("Checkpoint training configuration differs: verify model roles, optimizer type/groups, scheduler, precision and topology") + if missing and (not dist.is_initialized() or dist.get_rank() == 0): + warnings.warn("Older checkpoint lacks training metadata for " + ", ".join(missing) + + "; using the current configuration for these fields. Historical values cannot be verified. " + "The next save will record them.", RuntimeWarning, stacklevel=2) + + +def load_training_checkpoint(owners, input_dir=None, **kwargs): + """Validate every participant before mutating any model; stop training on failure.""" + owners = checkpoint_phase(lambda: _participants(owners)) + primary = next(iter(owners.values())) + if input_dir is None and primary.project_configuration.automatic_checkpoint_naming: + input_dir = _main_value(lambda: latest_training_checkpoint(Path(primary.project_dir) / "checkpoints", owners)) + path = checkpoint_phase(lambda: Path(input_dir).expanduser().absolute()) + def read_manifest(): + manifest_path = path / MANIFEST + if not manifest_path.exists(): + return None + value = json.loads(manifest_path.read_text()) + if not isinstance(value, dict): + raise ValueError("Invalid checkpoint manifest") + return value + manifest = checkpoint_phase(read_manifest) + descriptions = checkpoint_phase(lambda: { + name: _descriptor(owner, require_model_ids=manifest is not None) for name, owner in owners.items() + }) + _agree([str(path), descriptions, _hook_counts(owners), {name: {"legacy": checkpoint_config(owner).allow_legacy, + "finite": checkpoint_config(owner).check_finite} + for name, owner in owners.items()}]) + + def preflight(): + if any(parameter.grad is not None for owner in owners.values() + for model in owner._models for parameter in model.parameters()): + raise ValueError("Clear accumulated gradients with zero_grad(set_to_none=True) before restoring") + if manifest is None: + if not all(checkpoint_config(owner).allow_legacy for owner in owners.values()): + raise ValueError("Checkpoint has no completion manifest; allow_legacy=False requires a new-format checkpoint") + else: + if set(manifest["components"]) != set(owners): + raise ValueError("Checkpoint participants differ; restore the complete named group") + if manifest.get("versions") != {"accelerate": "1.7.0", "torch": "2.7.0"}: + raise ValueError("Checkpoint library versions differ") + _verify_inventory(path, manifest) + for name, current in descriptions.items(): + config = checkpoint_config(owners[name]) + _compare_descriptor(manifest["components"][name], current, + legacy_metadata_keys=config.legacy_metadata_keys if config.allow_legacy else (), + legacy_metadata_groups=config.legacy_metadata_groups if config.allow_legacy else ()) + rng_states = {name: _required_files(owner, path / name) for name, owner in owners.items()} + return manifest, rng_states + + manifest, rng_states = checkpoint_phase(preflight) + _agree(hashlib.sha256(json.dumps(manifest, sort_keys=True).encode()).hexdigest()) + if manifest is None: + def warn_legacy(): + if primary.is_main_process: + warnings.warn("Loading a legacy checkpoint using native component order; required files were checked, " + "but historical component identity and atomic completion cannot be verified. " + "The next save will include a completion manifest.", RuntimeWarning, stacklevel=2) + checkpoint_phase(warn_legacy) + try: + for name, owner in owners.items(): + def restore(owner=owner, name=name): + with _optimizer_adapter(owner), _synchronized_hooks(owner): + _native(owner, "load_state", str(path / name), **kwargs) + checkpoint_phase(restore) + owner.step = rng_states[name]["step"] + if checkpoint_config(owner).check_finite: + checkpoint_phase(lambda: check_finite_state([list(model.parameters()) for model in owner._models])) + if owner.distributed_type.name != "FSDP": + checkpoint_phase(lambda: check_finite_state([_raw_optimizer(optimizer).state for optimizer in owner._optimizers])) + # Native Accelerate swallows RNG errors. Apply the prevalidated state + # explicitly after all hooks/components, and propagate any failure. + checkpoint_phase(lambda: _apply_rng(next(reversed(rng_states.values())))) + except BaseException: + _set_failed(owners.values(), True) + raise + _set_failed(owners.values(), False) + run_id = manifest["run_id"] if manifest else _main_value(lambda: uuid.uuid4().hex) + for owner in owners.values(): + checkpoint_config(owner).run_id = run_id + if owner.project_configuration.automatic_checkpoint_naming: + match = _CHECKPOINT_NAME.fullmatch(path.name) + if match: + owner.project_configuration.iteration = max(owner.save_iteration, int(match.group(1)) + 1) + + +def dispatch_checkpoint(owner, method, original, *args, **kwargs): + """Agree on FSDP dispatch before choosing native or protected checkpoints.""" + from .. import (_accelerate_patch_requested, _accelerate_patch_owners, + _accelerate_patch_disabled_owners, _supported_accelerate_versions, + _eligible_fsdp_checkpoint) + + enabled = ((owner in _accelerate_patch_owners or _accelerate_patch_requested()) + and owner not in _accelerate_patch_disabled_owners) + plugin = owner.state.fsdp_plugin + _agree([enabled, _supported_accelerate_versions(), owner.distributed_type.name, + getattr(plugin, "fsdp_version", 1), getattr(plugin.state_dict_type, "name", None)]) + if not enabled or not _eligible_fsdp_checkpoint(plugin): + return original(owner, *args, **kwargs) + if method == "save_state": + return save_training_checkpoint(owner, *args, **kwargs) + return load_training_checkpoint(owner, *args, **kwargs) + + +def patch_accelerate_fsdp(accelerator): + """Explicitly enable FSDP1 sharded restore for an instance. + + Repository imports enable automatic dispatch by default. This opt-in API + also works with VIDEOX_ACCELERATE_PATCHES=0 and checks supported versions. + """ + from .. import enable_accelerate_patch + + return enable_accelerate_patch(accelerator) + + +def unpatch_accelerate_fsdp(accelerator): + """Disable automatic and explicit patch dispatch for this instance.""" + from .. import disable_accelerate_patch + + disable_accelerate_patch(accelerator) + + +def _checkpoint_phase(function, process_group): + """Propagate local preflight failures before entering FSDP collectives.""" + import torch.distributed as dist + + result, error = None, None + try: + result = function() + except BaseException as exc: + error = f"{type(exc).__name__}: {exc}" + errors = [None] * dist.get_world_size(process_group) + dist.all_gather_object(errors, error, group=process_group) + failures = [f"rank {rank}: {error}" for rank, error in enumerate(errors) if error] + if failures: + raise ValueError("FSDP optimizer restore failed: " + "; ".join(failures)) + return result + + +def _optimizer_layout(model, optimizer): + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from torch.distributed.fsdp._common_utils import _get_param_to_fqns + + names = _get_param_to_fqns(model) + groups = [[name for parameter in group["params"] for name in names[parameter]] + for group in optimizer.param_groups] + shapes = {name: parameter.shape for parameter, fqns in names.items() for name in fqns} + for unit in FSDP.fsdp_modules(model): + if unit._handle is None: + continue + flat = unit._handle.flat_param + if unit._use_orig_params: + for parameter, shape in zip(flat._params, flat._shapes): + for name in names[parameter]: + shapes[name] = shape + else: + shapes.update(zip(names[flat], flat._shapes)) + return groups, shapes + + +def _optimizer_load_destination(metadata, shapes, process_group): + """Allocate only this rank's CPU chunks, directly from saved metadata. + + Do not derive the schema from live optimizer state: it may be empty, partial, + or use different options. PyTorch's optimizer-specific DCP loader allocates + full CUDA tensors before sharding on NCCL, which is too large for CPU offload. + These native ShardedTensors use the same dim-0 layout without that allocation. + """ + import math + import torch.distributed as dist + from torch.distributed.checkpoint import TensorStorageMetadata + from torch.distributed.checkpoint._nested_dict import unflatten_state_dict + from torch.distributed._shard.sharded_tensor import Shard, ShardedTensor + from torch.distributed._shard.sharded_tensor.metadata import ShardedTensorMetadata, TensorProperties + from torch.distributed._shard.metadata import ShardMetadata + + rank, world_size = dist.get_rank(process_group), dist.get_world_size(process_group) + flat = {} + for key, stored in metadata.state_dict_metadata.items(): + path = metadata.planner_data[key] + if not path or path[0] != "optimizer": + raise ValueError(f"Unexpected optimizer checkpoint entry: {key}") + if not isinstance(stored, TensorStorageMetadata): + flat[key] = None + continue + size = stored.size + if len(path) >= 3 and path[1] == "state": + if path[2] not in shapes: + raise ValueError(f"Unknown optimizer parameter: {path[2]}") + is_moment = path[-1] in ("exp_avg", "exp_avg_sq", "max_exp_avg_sq") + if (len(size) or is_moment) and tuple(size) != tuple(shapes[path[2]]): + raise ValueError(f"Optimizer checkpoint tensor shape mismatch: {key}") + if not len(size) or path[1] != "state": + flat[key] = torch.empty(size, dtype=stored.properties.dtype, device="cpu") + continue + chunk = max(1, math.ceil(size[0] / world_size)) + shards, local = [], [] + for shard_rank, offset in enumerate(range(0, size[0], chunk)): + offsets = [offset] + [0] * (len(size) - 1) + lengths = [min(chunk, size[0] - offset)] + list(size[1:]) + shard = ShardMetadata(offsets, lengths, + f"rank:{dist.get_global_rank(process_group, shard_rank)}/cpu") + shards.append(shard) + if shard_rank == rank: + local.append(Shard(torch.empty(lengths, dtype=stored.properties.dtype, device="cpu"), shard)) + properties = TensorProperties(dtype=stored.properties.dtype, requires_grad=False) + flat[key] = ShardedTensor._init_from_local_shards_and_global_metadata( + local, ShardedTensorMetadata(shards, size, properties), process_group=process_group, + ) + return unflatten_state_dict(flat, metadata.planner_data) + + +def _align_optimizer_groups(state, expected_groups, optimizer, *, allow_reorder=True): + """Match groups by parameter names, including reordered and empty groups.""" + # DCP omits empty dictionaries (e.g. a checkpoint before the first step). + state.setdefault("state", {}) + saved = state["param_groups"] + if len(saved) != len(expected_groups): + raise ValueError("Optimizer checkpoint parameter group count differs") + if not allow_reorder and [group["params"] for group in saved] != expected_groups: + raise ValueError("Parameter group reordering with a scheduler is not supported") + remaining, aligned = list(saved), [] + for expected in expected_groups: + matches = [index for index, group in enumerate(remaining) + if len(group["params"]) == len(expected) and set(group["params"]) == set(expected)] + if not matches: + raise ValueError("Optimizer checkpoint parameter groups differ from the prepared optimizer") + group = dict(remaining.pop(matches[0]), params=list(expected)) + required_options = set(optimizer.defaults) + if required_options - group.keys(): + raise ValueError("Optimizer checkpoint options do not match this optimizer type") + current = optimizer.param_groups[len(aligned)] + # Tensor learning rates may live on CUDA (notably capturable Adam). + # DCP loads on CPU; retain the runtime device when restoring options. + for key, value in group.items(): + if isinstance(value, torch.Tensor): + if isinstance(current.get(key), torch.Tensor): + group[key] = value.to(current[key].device) + elif key == "lr" and group.get("capturable", False) and current["params"]: + group[key] = value.to(current["params"][0].device) + aligned.append(group) + expected_names = {name for group in expected_groups for name in group} + if set(state["state"]) - expected_names: + raise ValueError("Optimizer checkpoint contains state outside its parameter groups") + # A parameter may legitimately have no state before its first update. Once + # initialized, Adam needs all its moments, including the saved AMSGrad option. + if type(optimizer) in (torch.optim.Adam, torch.optim.AdamW): + for group in aligned: + required = {"step", "exp_avg", "exp_avg_sq"} + if group.get("amsgrad", False): + required.add("max_exp_avg_sq") + for name in group["params"]: + if name not in state["state"]: + continue + values = state["state"][name] + if required - values.keys(): + raise ValueError(f"Incomplete Adam optimizer state for {name}") + if any(not isinstance(values[key], torch.Tensor) for key in required): + raise ValueError(f"Invalid Adam optimizer tensor state for {name}") + if values["step"].ndim != 0: + raise ValueError(f"Invalid Adam optimizer step for {name}") + state["param_groups"] = aligned + + +def _load_sharded_fsdp_optimizer(accelerator, optimizer, model, input_dir, optimizer_index): + from torch.distributed.checkpoint import FileSystemReader, load + from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + from .fsdp_utils import release_completed_fsdp_gradients + + optimizer = getattr(optimizer, "optimizer", optimizer) + process_group = model.process_group + _checkpoint_phase(lambda: validate_checkpoint_owner(accelerator), process_group) + checkpoint_dir = Path(input_dir) / f"optimizer_{optimizer_index}" + + def prepare(): + if len(accelerator._optimizers) != len(accelerator._models): + raise ValueError("FSDP checkpoint restore requires one optimizer per prepared model") + groups, shapes = _optimizer_layout(model, optimizer) + metadata = FileSystemReader(checkpoint_dir).read_metadata() + if getattr(accelerator.state.fsdp_plugin.cpu_offload, "offload_params", False): + # Model restore creates idle pinned gradients. Free these before + # allocating moments, selecting each model independently: a single + # Accelerator can contain both original and flat parameter layouts. + original_param_models = [prepared for prepared in accelerator._models + if all(unit._use_orig_params for unit in FSDP.fsdp_modules(prepared))] + release_completed_fsdp_gradients(*original_param_models) + return groups, shapes, metadata + + expected_groups, shapes, metadata = _checkpoint_phase(prepare, process_group) + state = _checkpoint_phase( + lambda: _optimizer_load_destination(metadata, shapes, process_group), process_group, + ) + # These are already one-dimensional shards. The nested-shard flattener + # drops empty local shards and can mis-detect the checkpoint schema. + load(state, checkpoint_id=checkpoint_dir, process_group=process_group, + planner=DefaultLoadPlanner(flatten_sharded_tensors=False)) + + _checkpoint_phase( + lambda: _align_optimizer_groups(state["optimizer"], expected_groups, optimizer, + allow_reorder=not accelerator._schedulers), process_group, + ) + if checkpoint_config(accelerator).check_finite: + _checkpoint_phase(lambda: check_finite_state(state["optimizer"]), process_group) + _restore_fsdp_optimizer_state(model, optimizer, state["optimizer"]) + + +@torch.no_grad() +def _restore_fsdp_optimizer_state(model, optimizer, state): + """Release superseded moments before native FSDP conversion. + + All saved leaves and parameter groups have already been read and validated. + Custom optimizers and load pre-hooks retain access to the original state. + No dummy optimizer step is needed, even when restoring an empty optimizer. + """ + from torch.distributed.fsdp import ( + FullyShardedDataParallel as FSDP, StateDictType, + ShardedStateDictConfig, ShardedOptimStateDictConfig, + ) + + if (type(optimizer) in (torch.optim.Adam, torch.optim.AdamW, torch.optim.SGD) + and not optimizer._optimizer_load_state_dict_pre_hooks): + optimizer.state.clear() + with FSDP.state_dict_type( + model, StateDictType.SHARDED_STATE_DICT, + ShardedStateDictConfig(offload_to_cpu=True), + ShardedOptimStateDictConfig(offload_to_cpu=True), + ): + local = _checkpoint_phase(lambda: FSDP.optim_state_dict_to_load(model, optimizer, state), model.process_group) + _checkpoint_phase(lambda: optimizer.load_state_dict(local), model.process_group) + + +def register_checkpoint_state( + accelerator, generators, progress, *, filename="training_state_{rank}.pt", progress_keys=None, +): + """Save private noise streams and consumed batches separately for each rank.""" + + template = filename + filename = filename.format(rank=accelerator.process_index) + progress_keys = tuple(progress) if progress_keys is None else tuple(progress_keys) + config = checkpoint_config(accelerator) + config.private_states[template] = { + "schema_version": 1, + "generators": {name: generator.device.type for name, generator in generators.items()}, + "progress_keys": list(progress_keys), + } + + def read_state(directory, *, rank=None): + path = Path(directory) / template.format(rank=accelerator.process_index if rank is None else rank) + state = torch.load(path, map_location="cpu", weights_only=True) + if set(state["generators"]) != set(generators) or set(progress_keys) - state["progress"].keys(): + raise ValueError("Checkpoint private RNG/progress schema differs") + for name, generator in generators.items(): + previous = generator.get_state() + try: + generator.set_state(state["generators"][name]) + finally: + generator.set_state(previous) + return state + + config.validators[template] = read_state + + def save_state(models, weights, output_dir): + path = Path(output_dir) / filename + path.parent.mkdir(parents=True, exist_ok=True) + torch.save({"generators": {name: g.get_state() for name, g in generators.items()}, + "progress": {key: progress[key] for key in progress_keys}}, path) + + def load_state(models, input_dir): + state = read_state(input_dir) + if state is None: + return + for name, generator in generators.items(): + generator.set_state(state["generators"][name]) + progress.update(state["progress"], restored=True) + + return (accelerator.register_save_state_pre_hook(save_state), + accelerator.register_load_state_pre_hook(load_state)) + + +def check_finite_training_value(value, name): + if value is not None and not torch.isfinite(torch.as_tensor(value)).all(): + raise FloatingPointError(f"Non-finite {name}; refusing to update/save corrupted training weights") diff --git a/videox_fun/utils/fsdp_utils.py b/videox_fun/utils/fsdp_utils.py index d85e724..66dc803 100644 --- a/videox_fun/utils/fsdp_utils.py +++ b/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) diff --git a/videox_fun/utils/gpu_memory_mode.py b/videox_fun/utils/gpu_memory_mode.py index 4d465b2..521589e 100644 --- a/videox_fun/utils/gpu_memory_mode.py +++ b/videox_fun/utils/gpu_memory_mode.py @@ -21,6 +21,8 @@ and every script picks it up by writing ``model_cpu_offload_and_``. 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: