From cb8867f7321875d657cdc3a139597b2dcebad27d Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Mon, 31 Aug 2026 14:43:53 +0800 Subject: [PATCH 01/31] feat(platform): add PyTorch MPS device backend --- lightx2v_platform/base/__init__.py | 2 ++ lightx2v_platform/base/mps.py | 24 ++++++++++++++++++++++++ 2 files changed, 26 insertions(+) create mode 100644 lightx2v_platform/base/mps.py diff --git a/lightx2v_platform/base/__init__.py b/lightx2v_platform/base/__init__.py index fc28dc1ec..aace92b57 100755 --- a/lightx2v_platform/base/__init__.py +++ b/lightx2v_platform/base/__init__.py @@ -5,6 +5,7 @@ from lightx2v_platform.base.hygon_dcu import HygonDcuDevice from lightx2v_platform.base.mthreads_musa import MusaDevice from lightx2v_platform.base.metax_cuda import MetaxDevice +from lightx2v_platform.base.mps import MpsDevice from lightx2v_platform.base.nvidia import CudaDevice from lightx2v_platform.base.ppu_cuda import PpuDevice from lightx2v_platform.base.enflame_gcu import EnflameGcuDevice @@ -25,4 +26,5 @@ "EnflameGcuDevice", "IntelXpuDevice", "IluvatarDevice", + "MpsDevice", ] diff --git a/lightx2v_platform/base/mps.py b/lightx2v_platform/base/mps.py new file mode 100644 index 000000000..458cf2d6c --- /dev/null +++ b/lightx2v_platform/base/mps.py @@ -0,0 +1,24 @@ +import torch + +from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER + + +@PLATFORM_DEVICE_REGISTER("mps") +class MpsDevice: + name = "mps" + + @staticmethod + def init_device_env(): + pass + + @staticmethod + def is_available() -> bool: + return torch.backends.mps.is_available() + + @staticmethod + def get_device() -> str: + return "mps" + + @staticmethod + def init_parallel_env(): + raise NotImplementedError("MPS backend only supports single-device execution.") From 9babeeb0d4b3e45be88d2b73a65dea534f0b3d95 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Mon, 31 Aug 2026 22:25:09 +0800 Subject: [PATCH 02/31] fix(mps): isolate CUDA-only imports --- lightx2v/common/ops/attn/__init__.py | 46 +++++++------- lightx2v/common/ops/mm/mm_weight.py | 29 ++++++--- lightx2v/common/ops/norm/layer_norm_weight.py | 6 +- lightx2v/common/ops/norm/rms_norm_weight.py | 22 ++++--- lightx2v/common/ops/rope/__init__.py | 6 +- lightx2v/pipeline.py | 62 ++++++++++--------- 6 files changed, 104 insertions(+), 67 deletions(-) diff --git a/lightx2v/common/ops/attn/__init__.py b/lightx2v/common/ops/attn/__init__.py index ac44eb0b7..99b13b9d7 100755 --- a/lightx2v/common/ops/attn/__init__.py +++ b/lightx2v/common/ops/attn/__init__.py @@ -1,22 +1,26 @@ -from .draft_attn import DraftAttnWeight -from .dynamic_sparse_attn import DynamicSparseAttnWeight -from .flash_attn import ( - FlashAttn2Weight, - FlashAttn3Weight, - FlashAttn4Weight, - SparseFlashAttn4Weight, -) -from .general_sparse_attn import GeneralSparseAttnWeight -from .nbhd_attn import NbhdAttnWeight, NbhdAttnWeightFlashInfer -from .radial_attn import RadialAttnWeight -from .rainfusion_attn import RainfusionAttnWeight -from .ring_attn import RingAttnWeight -from .sage_attn import SageAttn2KInt8VFP8Weight, SageAttn2Weight, SageAttn3Weight, SparseSageAttn2Weight, SparseSageAttn3Weight -from .sol_attn import SolAttnWeight -from .sparge_attn import SpargeAttnWeight -from .sparse_mask_generator import NbhdMaskGenerator, SlaMaskGenerator, SpargeMaskGenerator, SvgMaskGenerator -from .sparse_operator import FlashinferOperator, FlexBlockOperator, MagiOperator, SlaTritonOperator, SparseFlashAttentionV4Operator, SparseSageAttentionV2Operator, SparseSageAttentionV3Operator -from .svg2_attn import Svg2AttnWeight -from .svg_attn import SvgAttnWeight +from lightx2v_platform.base.global_var import AI_DEVICE + from .torch_sdpa import TorchSDPAWeight -from .ulysses_attn import UlyssesAttnWeight + +if str(AI_DEVICE) != "mps": + from .draft_attn import DraftAttnWeight + from .dynamic_sparse_attn import DynamicSparseAttnWeight + from .flash_attn import ( + FlashAttn2Weight, + FlashAttn3Weight, + FlashAttn4Weight, + SparseFlashAttn4Weight, + ) + from .general_sparse_attn import GeneralSparseAttnWeight + from .nbhd_attn import NbhdAttnWeight, NbhdAttnWeightFlashInfer + from .radial_attn import RadialAttnWeight + from .rainfusion_attn import RainfusionAttnWeight + from .ring_attn import RingAttnWeight + from .sage_attn import SageAttn2KInt8VFP8Weight, SageAttn2Weight, SageAttn3Weight, SparseSageAttn2Weight, SparseSageAttn3Weight + from .sol_attn import SolAttnWeight + from .sparge_attn import SpargeAttnWeight + from .sparse_mask_generator import NbhdMaskGenerator, SlaMaskGenerator, SpargeMaskGenerator, SvgMaskGenerator + from .sparse_operator import FlashinferOperator, FlexBlockOperator, MagiOperator, SlaTritonOperator, SparseFlashAttentionV4Operator, SparseSageAttentionV2Operator, SparseSageAttentionV3Operator + from .svg2_attn import Svg2AttnWeight + from .svg_attn import SvgAttnWeight + from .ulysses_attn import UlyssesAttnWeight diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py index 61613dd5a..80313858d 100755 --- a/lightx2v/common/ops/mm/mm_weight.py +++ b/lightx2v/common/ops/mm/mm_weight.py @@ -6,20 +6,32 @@ from loguru import logger from safetensors import safe_open +from lightx2v_platform.base.global_var import AI_DEVICE + try: from magi_compiler import magi_register_custom_op except ImportError: magi_register_custom_op = None from lightx2v.common.ops.mm.sgl_kernel import sgl_fp8_scaled_mm, sgl_fp8_scaled_mm_meta -from lightx2v.common.ops.mm.triton_kernels import ( - fp8_gemm_bias_triton, - fp8_gemm_triton, - fp8_quantize_triton, - int8_gemm_bias_triton, - int8_gemm_triton, - int8_quantize_triton, -) + +if str(AI_DEVICE) == "mps": + fp8_gemm_bias_triton = None + fp8_gemm_triton = None + fp8_quantize_triton = None + int8_gemm_bias_triton = None + int8_gemm_triton = None + int8_quantize_triton = None +else: + from lightx2v.common.ops.mm.triton_kernels import ( + fp8_gemm_bias_triton, + fp8_gemm_triton, + fp8_quantize_triton, + int8_gemm_bias_triton, + int8_gemm_triton, + int8_quantize_triton, + ) + from lightx2v.common.ops.utils import * from lightx2v.utils.envs import * from lightx2v.utils.ggml_tensor import GGMLTensor @@ -27,7 +39,6 @@ from lightx2v.utils.global_paras import CALIB from lightx2v.utils.quant_utils import FloatQuantizer, IntegerQuantizer from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER -from lightx2v_platform.base.global_var import AI_DEVICE try: from lightx2v_kernel.gemm import ( diff --git a/lightx2v/common/ops/norm/layer_norm_weight.py b/lightx2v/common/ops/norm/layer_norm_weight.py index 732d9f273..8307a9abd 100755 --- a/lightx2v/common/ops/norm/layer_norm_weight.py +++ b/lightx2v/common/ops/norm/layer_norm_weight.py @@ -7,8 +7,12 @@ from lightx2v.common.ops.utils import * from lightx2v.utils.envs import * from lightx2v.utils.registry_factory import LN_WEIGHT_REGISTER +from lightx2v_platform.base.global_var import AI_DEVICE -from .triton_ops import norm_infer +if str(AI_DEVICE) == "mps": + norm_infer = None +else: + from .triton_ops import norm_infer try: from magi_compiler import magi_register_custom_op diff --git a/lightx2v/common/ops/norm/rms_norm_weight.py b/lightx2v/common/ops/norm/rms_norm_weight.py index 67503ad1a..d60c7da23 100755 --- a/lightx2v/common/ops/norm/rms_norm_weight.py +++ b/lightx2v/common/ops/norm/rms_norm_weight.py @@ -5,16 +5,24 @@ from loguru import logger from safetensors import safe_open -from lightx2v.common.ops.norm.triton_ops import ( - fused_norm_3drope, - fused_qk_norm_3drope, - fused_qk_rms_norm, - rms_norm_kernel, -) +from lightx2v_platform.base.global_var import AI_DEVICE + +if str(AI_DEVICE) == "mps": + fused_norm_3drope = None + fused_qk_norm_3drope = None + fused_qk_rms_norm = None + rms_norm_kernel = None +else: + from lightx2v.common.ops.norm.triton_ops import ( + fused_norm_3drope, + fused_qk_norm_3drope, + fused_qk_rms_norm, + rms_norm_kernel, + ) + from lightx2v.common.ops.utils import * from lightx2v.utils.envs import * from lightx2v.utils.registry_factory import RMS_WEIGHT_REGISTER -from lightx2v_platform.base.global_var import AI_DEVICE try: import sgl_kernel diff --git a/lightx2v/common/ops/rope/__init__.py b/lightx2v/common/ops/rope/__init__.py index 1cdeff55d..f34d799cf 100644 --- a/lightx2v/common/ops/rope/__init__.py +++ b/lightx2v/common/ops/rope/__init__.py @@ -1,4 +1,8 @@ +from lightx2v_platform.base.global_var import AI_DEVICE + from .chunked_rope import ChunkedRope -from .flashinfer_rope import FlashInferRope from .template import RopeLayout, RopeTemplate from .torch_rope import TorchComplexRope, TorchRealRope + +if str(AI_DEVICE) != "mps": + from .flashinfer_rope import FlashInferRope diff --git a/lightx2v/pipeline.py b/lightx2v/pipeline.py index f25dd674f..215c3c2e3 100755 --- a/lightx2v/pipeline.py +++ b/lightx2v/pipeline.py @@ -12,34 +12,40 @@ from lightx2v.common.ops import * from lightx2v.models.networks.wan.animate2_identity import WAN_ANIMATE2_MODEL_ID -from lightx2v.models.runners.ernie_image.ernie_image_runner import ErnieImageRunner # noqa: F401 -from lightx2v.models.runners.flux2.flux2_runner import Flux2DevRunner, Flux2KleinRunner # noqa: F401 -from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner # noqa: F401 -from lightx2v.models.runners.lingbot_video.lingbot_video_runner import LingBotVideoRunner # noqa: F401 -from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner # noqa: F401 -from lightx2v.models.runners.ltx2.ltx2_runner import LTX2Runner # noqa: F401 -from lightx2v.models.runners.ltx2.ltx25_runner import LTX25Runner # noqa: F401 -from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 -from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 -from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 -from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner # noqa: F401 -from lightx2v.models.runners.swiftvr.swiftvr_runner import SwiftVRRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_animate2_runner import WanAnimate2Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_animate_runner import WanAnimateRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner, WanAudioRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_distill_runner import WanDistillRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_lingbot_fast_runner import LingbotFastRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_lingbot_va_runner import LingbotVARunner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game2_runner import WanSFMtxg2Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game3_runner import WanMatrixGame3Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_sf_runner import WanSFRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_vace_runner import WanVaceRunner # noqa: F401 -from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_ar_runner import WorldPlayARRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_bi_runner import WorldPlayBIRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_distill_runner import WorldPlayDistillRunner # noqa: F401 -from lightx2v.models.runners.z_image.z_image_runner import ZImageRunner # noqa: F401 +from lightx2v_platform.base.global_var import AI_DEVICE + +if str(AI_DEVICE) == "mps": + from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 +else: + from lightx2v.models.runners.ernie_image.ernie_image_runner import ErnieImageRunner # noqa: F401 + from lightx2v.models.runners.flux2.flux2_runner import Flux2DevRunner, Flux2KleinRunner # noqa: F401 + from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner # noqa: F401 + from lightx2v.models.runners.lingbot_video.lingbot_video_runner import LingBotVideoRunner # noqa: F401 + from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner # noqa: F401 + from lightx2v.models.runners.ltx2.ltx2_runner import LTX2Runner # noqa: F401 + from lightx2v.models.runners.ltx2.ltx25_runner import LTX25Runner # noqa: F401 + from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 + from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 + from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 + from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner # noqa: F401 + from lightx2v.models.runners.swiftvr.swiftvr_runner import SwiftVRRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_animate2_runner import WanAnimate2Runner # noqa: F401 + from lightx2v.models.runners.wan.wan_animate_runner import WanAnimateRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner, WanAudioRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_distill_runner import WanDistillRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_lingbot_fast_runner import LingbotFastRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_lingbot_va_runner import LingbotVARunner # noqa: F401 + from lightx2v.models.runners.wan.wan_matrix_game2_runner import WanSFMtxg2Runner # noqa: F401 + from lightx2v.models.runners.wan.wan_matrix_game3_runner import WanMatrixGame3Runner # noqa: F401 + from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_sf_runner import WanSFRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_vace_runner import WanVaceRunner # noqa: F401 + from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner # noqa: F401 + from lightx2v.models.runners.worldplay.worldplay_ar_runner import WorldPlayARRunner # noqa: F401 + from lightx2v.models.runners.worldplay.worldplay_bi_runner import WorldPlayBIRunner # noqa: F401 + from lightx2v.models.runners.worldplay.worldplay_distill_runner import WorldPlayDistillRunner # noqa: F401 + from lightx2v.models.runners.z_image.z_image_runner import ZImageRunner # noqa: F401 + from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict from lightx2v.utils.registry_factory import RUNNER_REGISTER from lightx2v.utils.set_config import set_config, set_parallel_config From 8e000492e041567c0c30328b603b2cc10a47a99c Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Mon, 31 Aug 2026 22:25:14 +0800 Subject: [PATCH 03/31] fix(runtime): support MPS cache management --- lightx2v/models/runners/default_runner.py | 29 +++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/lightx2v/models/runners/default_runner.py b/lightx2v/models/runners/default_runner.py index 2d5ff4c0f..0c9c40919 100755 --- a/lightx2v/models/runners/default_runner.py +++ b/lightx2v/models/runners/default_runner.py @@ -355,6 +355,35 @@ def maybe_empty_cache(self, *, force: bool = False, collect_garbage: bool = Fals min_free_bytes = float(self.config.get("empty_cache_min_free_gib", 4)) * gib min_reclaimable_bytes = float(self.config.get("empty_cache_min_reclaimable_gib", 2)) * gib + if AI_DEVICE == "mps": + driver_allocated_bytes = torch.mps.driver_allocated_memory() + recommended_max_bytes = torch.mps.recommended_max_memory() + headroom_bytes = max(recommended_max_bytes - driver_allocated_bytes, 0) + check_cache = force or headroom_bytes < min_free_bytes + if collect_garbage or check_cache: + gc.collect() + if not check_cache: + return False + + current_allocated_bytes = torch.mps.current_allocated_memory() + driver_allocated_bytes = torch.mps.driver_allocated_memory() + recommended_max_bytes = torch.mps.recommended_max_memory() + headroom_bytes = max(recommended_max_bytes - driver_allocated_bytes, 0) + reclaimable_bytes = max(driver_allocated_bytes - current_allocated_bytes, 0) + + if force or reclaimable_bytes >= min_reclaimable_bytes: + logger.info( + f"[Memory] Emptying MPS cache: headroom={headroom_bytes / gib:.2f} GiB, " + f"current_allocated={current_allocated_bytes / gib:.2f} GiB, " + f"driver_allocated={driver_allocated_bytes / gib:.2f} GiB, " + f"recommended_max={recommended_max_bytes / gib:.2f} GiB, " + f"reclaimable={reclaimable_bytes / gib:.2f} GiB, force={force}" + ) + torch.mps.empty_cache() + return True + + return False + free_bytes, _ = torch_device_module.mem_get_info() check_cache = force or free_bytes < min_free_bytes if collect_garbage or check_cache: From a54fbc248e007813a962e646ebd8b94ea5955089 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Tue, 1 Sep 2026 13:04:40 +0800 Subject: [PATCH 04/31] feat(mps): add MiniMax H3 inference config --- configs/platforms/mps/minimax_h3_t2av.json | 44 ++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 configs/platforms/mps/minimax_h3_t2av.json diff --git a/configs/platforms/mps/minimax_h3_t2av.json b/configs/platforms/mps/minimax_h3_t2av.json new file mode 100644 index 000000000..feaeb60f9 --- /dev/null +++ b/configs/platforms/mps/minimax_h3_t2av.json @@ -0,0 +1,44 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 480, + "target_width": 480, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + + "cpu_offload": true, + "offload_granularity": "block", + "dit_prepost_resident": false, + + "text_encoder_cpu_offload": true, + "text_encoder_offload_granularity": "block", + "text_encoder_host_pinned": false, + "text_encoder_release_block_offload_buffers": true, + + "vae_cpu_offload": true, + "vae_use_compile": false, + "vae_attn_type": "torch_sdpa", + + "lazy_load": false, + "unload_modules": false, + "warmup": false, + + "attn_type": "torch_sdpa", + "rms_type": "torch_native", + "rope_type": "torch_real_rope", + + "dit_quantized": false, + "dit_quant_scheme": "Default", + + "feature_caching": "NoCaching", + "use_compile": false, + + "video_flow_shift": 12.0, + "audio_flow_shift": 3.0, + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true +} From 8534b4e692561b49d9e0f6b45ddead79e3b0bde3 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Tue, 1 Sep 2026 14:46:45 +0800 Subject: [PATCH 05/31] feat(mps): add MiniMax H3 CLI entry and launch script --- lightx2v/infer.py | 86 +++++++++++--------- scripts/platforms/mps/run_minimax_h3_t2av.sh | 47 +++++++++++ 2 files changed, 93 insertions(+), 40 deletions(-) create mode 100755 scripts/platforms/mps/run_minimax_h3_t2av.sh diff --git a/lightx2v/infer.py b/lightx2v/infer.py index 9ca794686..e1f85ebe5 100755 --- a/lightx2v/infer.py +++ b/lightx2v/infer.py @@ -8,46 +8,52 @@ from lightx2v.common.ops import * from lightx2v.models.networks.bagel.sensenova_tasks import OMNI_VISION_SUBTASK_CHOICES from lightx2v.models.networks.wan.animate2_identity import WAN_ANIMATE2_MODEL_ID -from lightx2v.models.runners.bagel.bagel_runner import BagelRunner # noqa: F401 -from lightx2v.models.runners.bagel.sensenova_vision_runner import SenseNovaVisionRunner # noqa: F401 -from lightx2v.models.runners.cosmos3.cosmos3_runner import Cosmos3Runner # noqa: F401 -from lightx2v.models.runners.ernie_image.ernie_image_runner import ErnieImageRunner # noqa: F401 -from lightx2v.models.runners.flux2.flux2_runner import Flux2DevRunner, Flux2KleinRunner # noqa: F401 -from lightx2v.models.runners.hidream_o1_image.hidream_o1_image_runner import HidreamO1ImageRunner # noqa: F401 -from lightx2v.models.runners.hunyuan3d.hunyuan3d_shape_runner import Hunyuan3DShapeRunner # noqa: F401 -from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner # noqa: F401 -from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_distill_runner import HunyuanVideo15DistillRunner # noqa: F401 -from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner # noqa: F401 -from lightx2v.models.runners.lingbot_video.lingbot_video_runner import LingBotVideoRunner # noqa: F401 -from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner # noqa: F401 -from lightx2v.models.runners.ltx2.ltx2_runner import LTX2ARRunner, LTX2Runner # noqa: F401 -from lightx2v.models.runners.ltx2.ltx25_runner import LTX25Runner # noqa: F401 -from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 -from lightx2v.models.runners.motus.motus_runner import MotusRunner # noqa: F401 -from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 -from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 -from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner # noqa: F401 -from lightx2v.models.runners.swiftvr.swiftvr_runner import SwiftVRRunner # noqa: F401 -from lightx2v.models.runners.wan.fastwam_runner import FastWAMRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_animate2_runner import WanAnimate2Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_animate_runner import WanAnimateRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner, WanAudioRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_dancer_runner import WanDancerRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_distill_runner import WanDistillRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_dreamzero_runner import WanDreamZeroRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_infinitetalk_runner import InfiniteTalkRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_lingbot_va_runner import LingbotVARunner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game2_runner import WanSFMtxg2Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_matrix_game3_runner import WanMatrixGame3Runner # noqa: F401 -from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_s2v_runner import WanS2VRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_sf_runner import WanSFRunner # noqa: F401 -from lightx2v.models.runners.wan.wan_vace_runner import Wan22MoeVaceRunner, WanVaceRunner # noqa: F401 -from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_ar_runner import WorldPlayARRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_bi_runner import WorldPlayBIRunner # noqa: F401 -from lightx2v.models.runners.worldplay.worldplay_distill_runner import WorldPlayDistillRunner # noqa: F401 -from lightx2v.models.runners.z_image.z_image_runner import ZImageRunner # noqa: F401 +from lightx2v_platform.base.global_var import AI_DEVICE + +if str(AI_DEVICE) == "mps": + from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 +else: + from lightx2v.models.runners.bagel.bagel_runner import BagelRunner # noqa: F401 + from lightx2v.models.runners.bagel.sensenova_vision_runner import SenseNovaVisionRunner # noqa: F401 + from lightx2v.models.runners.cosmos3.cosmos3_runner import Cosmos3Runner # noqa: F401 + from lightx2v.models.runners.ernie_image.ernie_image_runner import ErnieImageRunner # noqa: F401 + from lightx2v.models.runners.flux2.flux2_runner import Flux2DevRunner, Flux2KleinRunner # noqa: F401 + from lightx2v.models.runners.hidream_o1_image.hidream_o1_image_runner import HidreamO1ImageRunner # noqa: F401 + from lightx2v.models.runners.hunyuan3d.hunyuan3d_shape_runner import Hunyuan3DShapeRunner # noqa: F401 + from lightx2v.models.runners.hunyuan_image3.hunyuan_image3_runner import HunyuanImage3Runner # noqa: F401 + from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_distill_runner import HunyuanVideo15DistillRunner # noqa: F401 + from lightx2v.models.runners.hunyuan_video.hunyuan_video_15_runner import HunyuanVideo15Runner # noqa: F401 + from lightx2v.models.runners.lingbot_video.lingbot_video_runner import LingBotVideoRunner # noqa: F401 + from lightx2v.models.runners.longcat_image.longcat_image_runner import LongCatImageRunner # noqa: F401 + from lightx2v.models.runners.ltx2.ltx2_runner import LTX2ARRunner, LTX2Runner # noqa: F401 + from lightx2v.models.runners.ltx2.ltx25_runner import LTX25Runner # noqa: F401 + from lightx2v.models.runners.minimax_h3.minimax_h3_runner import MiniMaxH3Runner # noqa: F401 + from lightx2v.models.runners.motus.motus_runner import MotusRunner # noqa: F401 + from lightx2v.models.runners.neopp.neopp_runner import NeoppRunner # noqa: F401 + from lightx2v.models.runners.qwen_image.qwen_image_runner import QwenImageRunner # noqa: F401 + from lightx2v.models.runners.seedvr.seedvr_runner import SeedVRRunner # noqa: F401 + from lightx2v.models.runners.swiftvr.swiftvr_runner import SwiftVRRunner # noqa: F401 + from lightx2v.models.runners.wan.fastwam_runner import FastWAMRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_animate2_runner import WanAnimate2Runner # noqa: F401 + from lightx2v.models.runners.wan.wan_animate_runner import WanAnimateRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_audio_runner import Wan22AudioRunner, WanAudioRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_dancer_runner import WanDancerRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_distill_runner import WanDistillRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_dreamzero_runner import WanDreamZeroRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_infinitetalk_runner import InfiniteTalkRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_lingbot_va_runner import LingbotVARunner # noqa: F401 + from lightx2v.models.runners.wan.wan_matrix_game2_runner import WanSFMtxg2Runner # noqa: F401 + from lightx2v.models.runners.wan.wan_matrix_game3_runner import WanMatrixGame3Runner # noqa: F401 + from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_s2v_runner import WanS2VRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_sf_runner import WanSFRunner # noqa: F401 + from lightx2v.models.runners.wan.wan_vace_runner import Wan22MoeVaceRunner, WanVaceRunner # noqa: F401 + from lightx2v.models.runners.worldmirror.worldmirror_runner import WorldMirrorRunner # noqa: F401 + from lightx2v.models.runners.worldplay.worldplay_ar_runner import WorldPlayARRunner # noqa: F401 + from lightx2v.models.runners.worldplay.worldplay_bi_runner import WorldPlayBIRunner # noqa: F401 + from lightx2v.models.runners.worldplay.worldplay_distill_runner import WorldPlayDistillRunner # noqa: F401 + from lightx2v.models.runners.z_image.z_image_runner import ZImageRunner # noqa: F401 + from lightx2v.utils.envs import * from lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict from lightx2v.utils.profiler import * diff --git a/scripts/platforms/mps/run_minimax_h3_t2av.sh b/scripts/platforms/mps/run_minimax_h3_t2av.sh new file mode 100755 index 000000000..3f8ba842b --- /dev/null +++ b/scripts/platforms/mps/run_minimax_h3_t2av.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd -- "${SCRIPT_DIR}/../../.." && pwd) + +lightx2v_path=${LIGHTX2V_PATH:-${REPO_ROOT}} +model_path=${MODEL_PATH:-} +config_json=${CONFIG_JSON:-${lightx2v_path}/configs/platforms/mps/minimax_h3_t2av.json} +output_path=${OUTPUT_PATH:-${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av.mp4} + +export PLATFORM=mps +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 +export TOKENIZERS_PARALLELISM=false +export PYTHONFAULTHANDLER=1 +export PYTHONUNBUFFERED=1 +export PYTHONPATH="${lightx2v_path}:${PYTHONPATH:-}" + +if [[ -z "${model_path}" ]]; then + echo "MODEL_PATH must point to the MiniMax-H3 model directory." + exit 1 +fi + +[[ -d "${model_path}" ]] || { + echo "Model directory not found: ${model_path}" + exit 1 +} + +[[ -f "${config_json}" ]] || { + echo "Config file not found: ${config_json}" + exit 1 +} + +mkdir -p "$(dirname -- "${output_path}")" + +prompt=${PROMPT:-A cinematic fox walking through a snowy forest} +seed=${SEED:-42} + +python -m lightx2v.infer \ + --model_cls minimax_h3 \ + --task t2av \ + --model_path "${model_path}" \ + --config_json "${config_json}" \ + --prompt "${prompt}" \ + --save_result_path "${output_path}" \ + --seed "${seed}" From b6a47aab3241df75a911a10baa53e144889d42e6 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Tue, 1 Sep 2026 19:52:27 +0800 Subject: [PATCH 06/31] feat(mps): stream MiniMax H3 Qwen weights from disk --- .../input_encoders/hf/minimax_h3/qwen3vl.py | 276 ++++++++++++- .../minimax_h3/test_qwen3vl_disk_streaming.py | 370 ++++++++++++++++++ 2 files changed, 641 insertions(+), 5 deletions(-) create mode 100644 tests/models/minimax_h3/test_qwen3vl_disk_streaming.py diff --git a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py index 0d2c03c2f..387839c61 100644 --- a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py +++ b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py @@ -72,6 +72,7 @@ MINIMAX_H3_TEXT_TAG = 1 _CHECKPOINT_PREFIX = "model.language_model" +_QWEN_LAYER_PREFIX = f"{_CHECKPOINT_PREFIX}.layers" _EXPECTED_RELEASE_CONFIG = { "hidden_size": MINIMAX_H3_TEXT_HIDDEN_SIZE, "intermediate_size": 25600, @@ -83,6 +84,46 @@ } +def _resolve_qwen_layer_name(name, layer_index): + layer_prefix = f"{_QWEN_LAYER_PREFIX}." + if not name.startswith(layer_prefix): + return name + parts = name.split(".", 4) + if len(parts) == 5 and parts[3].isdigit(): + return f"{_QWEN_LAYER_PREFIX}.{int(layer_index)}.{parts[4]}" + return name + + +def _iter_base_attrs(module): + if hasattr(module, "base_attrs"): + yield from module.base_attrs + for child in getattr(module, "_modules", {}).values(): + if child is not None: + yield from _iter_base_attrs(child) + + +def _load_selected_checkpoint_tensors(text_encoder_path, weight_map, names): + names = tuple(sorted(dict.fromkeys(names))) + missing = sorted(name for name in names if name not in weight_map) + if missing: + raise KeyError(f"MiniMax-H3 Qwen3-VL checkpoint is missing requested tensors: {missing}") + + by_shard = defaultdict(list) + for name in names: + by_shard[weight_map[name]].append(name) + + root = Path(text_encoder_path) + tensors = {} + for shard_name in sorted(by_shard): + shard_path = root / shard_name + if not shard_path.is_file(): + raise FileNotFoundError(f"Safetensors shard from checkpoint index was not found: {shard_path}") + with safe_open(shard_path, framework="pt", device="cpu") as checkpoint: + for name in sorted(by_shard[shard_name]): + tensors[name] = checkpoint.get_tensor(name) + return tensors + + def _empty_device_cache(): """Release cached accelerator allocations without assuming CUDA.""" with suppress(Exception): @@ -342,16 +383,29 @@ def forward(self, hidden_states, position_embeddings): class _Qwen3VLTextBackboneWeights(WeightModule): """Unbatched native prefix of Qwen3-VL's language backbone.""" - def __init__(self, config, text_config, num_layers=MINIMAX_H3_TEXT_ENCODER_LAYER, attn_type="torch_sdpa", block_offload=False, tp_group=None): + def __init__( + self, + config, + text_config, + num_layers=MINIMAX_H3_TEXT_ENCODER_LAYER, + attn_type="torch_sdpa", + block_offload=False, + tp_group=None, + disk_streaming=False, + ): super().__init__() self.config = config self.text_config = text_config self.num_layers = int(num_layers) self.attn_type = attn_type - self.block_offload = bool(block_offload) + self.disk_streaming = bool(disk_streaming) + self.block_offload = bool(block_offload) and not self.disk_streaming self.offload_manager = None self.offload_cuda_buffers = None self._offload_completion_event = None + self.streaming_layer = None + self._disk_streaming_text_encoder_path = None + self._disk_streaming_weight_map = None self.hidden_size = int(text_config["hidden_size"]) self.head_dim = int(text_config["head_dim"]) self.rope_theta = float(text_config["rope_theta"]) @@ -399,6 +453,17 @@ def _collect_weight_modules(self): modules.update((leaf.weight_name, leaf) for leaf in layer.weight_modules()) return modules + @staticmethod + def _layer_tensor_names(layer, layer_index): + names = [] + seen = set() + for name, _, _ in _iter_base_attrs(layer): + actual_name = _resolve_qwen_layer_name(name, layer_index) + if actual_name not in seen: + seen.add(actual_name) + names.append(actual_name) + return tuple(sorted(names)) + def select_tp_shard(self, name, tensor): if self.tp_size == 1: return tensor @@ -483,6 +548,128 @@ def init_block_offload(self): allocated_bytes / (1024**3), ) + def init_disk_streaming(self, text_encoder_path=None, weight_map=None): + if not self.disk_streaming: + return + if text_encoder_path is not None: + self._disk_streaming_text_encoder_path = Path(text_encoder_path) + if weight_map is not None: + self._disk_streaming_weight_map = dict(weight_map) + if self._disk_streaming_text_encoder_path is None or self._disk_streaming_weight_map is None: + raise RuntimeError("Qwen3-VL disk streaming requires a checkpoint path and weight map") + if self.streaming_layer is not None: + return + + self.streaming_layer = _Qwen3VLDecoderLayerWeights( + 0, + self.config, + self.text_config, + self.attn_type, + tp_group=self.tp_group, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + create_cuda_buffer=True, + ) + layer0_names = self._layer_tensor_names(self.streaming_layer, 0) + layer0_tensors = _load_selected_checkpoint_tensors( + self._disk_streaming_text_encoder_path, + self._disk_streaming_weight_map, + layer0_names, + ) + try: + self.streaming_layer.load(layer0_tensors) + self.streaming_layer.load_state_dict(self._prepare_streaming_state_dict(layer0_tensors, 0), 0) + finally: + del layer0_tensors + gc.collect() + + def _prepare_streaming_state_dict(self, tensors, layer_index): + state_dict = dict(tensors) + for name, _, transpose in _iter_base_attrs(self.streaming_layer): + if transpose: + actual_name = _resolve_qwen_layer_name(name, layer_index) + if actual_name in state_dict: + state_dict[actual_name] = state_dict[actual_name].t() + return state_dict + + def load_streaming_layer(self, layer_index): + if not self.disk_streaming: + raise RuntimeError("Qwen3-VL load_streaming_layer requires text_encoder_disk_streaming=true") + layer_index = int(layer_index) + if layer_index < 0 or layer_index >= self.num_layers: + raise IndexError(f"Qwen3-VL layer index out of range: {layer_index}") + self.init_disk_streaming() + + layer_names = self._layer_tensor_names(self.streaming_layer, layer_index) + tensors = _load_selected_checkpoint_tensors( + self._disk_streaming_text_encoder_path, + self._disk_streaming_weight_map, + layer_names, + ) + try: + self.streaming_layer.load_state_dict(self._prepare_streaming_state_dict(tensors, layer_index), layer_index) + finally: + del tensors + gc.collect() + return self.streaming_layer + + def _forward_streaming_embedding(self, input_ids): + if not self.disk_streaming: + raise RuntimeError("Qwen3-VL streaming embedding requires text_encoder_disk_streaming=true") + embedding_name = self.embed_tokens.weight_name + tensors = _load_selected_checkpoint_tensors( + self._disk_streaming_text_encoder_path, + self._disk_streaming_weight_map, + (embedding_name,), + ) + try: + host_weight = tensors.pop(embedding_name) + device_weight = host_weight.to(AI_DEVICE) + self.embed_tokens.weight = device_weight + try: + hidden_states = self.embed_tokens.apply(input_ids) + torch_device_module.synchronize() + finally: + self.embed_tokens.weight = None + del device_weight + del host_weight + finally: + del tensors + gc.collect() + _empty_device_cache() + if hasattr(self.embed_tokens, "pin_weight"): + self.embed_tokens.pin_weight = None + return hidden_states + + def _forward_with_disk_streaming(self, input_ids, position_ids, vision_mask, vision_embeds, deepstack_embeds): + if vision_mask is not None or vision_embeds is not None or deepstack_embeds is not None: + raise NotImplementedError("MiniMax-H3 Qwen3-VL disk streaming currently supports text-only t2av prompts.") + hidden_states = self._forward_streaming_embedding(input_ids) + position_embeddings = self._position_embeddings(hidden_states, position_ids) + for layer_index in range(self.num_layers): + layer = self.load_streaming_layer(layer_index) + hidden_states = layer.forward(hidden_states, position_embeddings) + torch_device_module.synchronize() + return hidden_states + + def release_disk_streaming_buffer(self): + if self.streaming_layer is None: + return + with suppress(Exception): + torch_device_module.synchronize() + for module in self.streaming_layer.weight_modules(): + storage = unwrap_tp_linear(module) + for _, attr_name, _ in getattr(storage, "base_attrs", ()): + if hasattr(storage, attr_name): + setattr(storage, attr_name, None) + buffer_attr = f"{attr_name}_cuda_buffer" + if hasattr(storage, buffer_attr): + setattr(storage, buffer_attr, None) + self.streaming_layer = None + gc.collect() + _empty_device_cache() + logger.info("MiniMax-H3 Qwen3-VL released its disk-streaming layer buffer") + def release_block_offload_buffers(self): """Release transient device slots while retaining CPU checkpoint views.""" if self.offload_manager is None: @@ -499,12 +686,16 @@ def named_weight_modules(self): @property def device(self): + if self.disk_streaming: + return torch.device(AI_DEVICE) if self.block_offload: return torch.device(AI_DEVICE) return self.embed_tokens.weight.device @property def dtype(self): + if self.disk_streaming: + return GET_DTYPE() return self.embed_tokens.weight.dtype def _position_embeddings(self, hidden_states, position_ids=None): @@ -596,6 +787,8 @@ def _forward_with_block_offload(self, input_ids, position_ids, vision_mask, visi def forward(self, input_ids, position_ids=None, vision_mask=None, vision_embeds=None, deepstack_embeds=None): if input_ids.ndim != 1: raise ValueError(f"MiniMax-H3's native Qwen3-VL backbone expects unbatched token IDs, got {tuple(input_ids.shape)}") + if self.disk_streaming: + return self._forward_with_disk_streaming(input_ids, position_ids, vision_mask, vision_embeds, deepstack_embeds) if self.block_offload: return self._forward_with_block_offload(input_ids, position_ids, vision_mask, vision_embeds, deepstack_embeds) @@ -638,9 +831,23 @@ class MiniMaxH3Qwen3VLTextEncoder: def __init__(self, config): self.config = config + self.disk_streaming = bool(config.get("text_encoder_disk_streaming", False)) if config.get("text_encoder_quantized", False) and not config.get("text_encoder_quantized_ckpt"): raise ValueError("MiniMax-H3 quantized text encoder requires text_encoder_quantized_ckpt") self.tensor_parallel = bool(config.get("text_encoder_tensor_parallel", config.get("tensor_parallel", False))) + if self.disk_streaming: + if torch.device(AI_DEVICE).type != "mps": + raise ValueError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming currently requires AI_DEVICE='mps'.") + if config.get("task") != "t2av": + raise ValueError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming currently supports task='t2av' only.") + if GET_DTYPE() != torch.bfloat16: + raise ValueError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming currently requires BF16.") + if config.get("text_encoder_quantized", False): + raise NotImplementedError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming does not support quantized text encoder weights.") + if self.tensor_parallel: + raise NotImplementedError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming does not support text encoder tensor parallel.") + if config.get("text_encoder_async_prefetch", False) or config.get("text_encoder_double_buffer", False): + raise NotImplementedError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming is synchronous and does not support prefetch or double buffering.") if self.tensor_parallel: if not dist.is_initialized(): raise RuntimeError("MiniMax-H3 text encoder TP requires an initialized distributed process group") @@ -665,7 +872,12 @@ def __init__(self, config): raise ValueError(f"Unsupported text_encoder_offload_granularity={self.offload_granularity!r}; expected 'model' or 'block'") if self.offload_granularity == "block" and not self.cpu_offload: raise ValueError("text_encoder_offload_granularity='block' requires text_encoder_cpu_offload=true") - self.block_offload = self.cpu_offload and self.offload_granularity == "block" + if self.disk_streaming: + if not self.cpu_offload: + raise ValueError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming requires text_encoder_cpu_offload=true.") + if self.offload_granularity != "block": + raise ValueError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming requires text_encoder_offload_granularity='block'.") + self.block_offload = self.cpu_offload and self.offload_granularity == "block" and not self.disk_streaming self.release_block_offload_buffers = bool(config.get("text_encoder_release_block_offload_buffers", False)) self.local_files_only = config.get("local_files_only", True) self.text_encoder = None @@ -917,6 +1129,51 @@ def _load_native_weights(cls, backbone, text_encoder_path, text_config): backbone.to_cpu() return checkpoint_dtypes.pop() + @classmethod + def _preflight_native_checkpoint(cls, backbone, text_encoder_path, text_config): + modules = dict(backbone.named_weight_modules()) + expected_shapes = cls._expected_weight_shapes(text_config) + if modules.keys() != expected_shapes.keys(): + missing_native = sorted(expected_shapes.keys() - modules.keys()) + unexpected_native = sorted(modules.keys() - expected_shapes.keys()) + raise RuntimeError(f"Native Qwen3-VL weight declaration disagrees with its shape schema: missing={missing_native}, unexpected={unexpected_native}") + + root = Path(text_encoder_path) + weight_map = cls._checkpoint_weight_map(root, modules) + missing = sorted(modules.keys() - weight_map.keys()) + if missing: + preview = ", ".join(missing[:8]) + raise KeyError(f"MiniMax-H3 text encoder checkpoint is missing {len(missing)} required tensors: {preview}") + + by_shard = defaultdict(list) + for name in modules: + by_shard[weight_map[name]].append(name) + + checkpoint_dtypes = set() + for shard_name in sorted(by_shard): + shard_path = root / shard_name + if not shard_path.is_file(): + raise FileNotFoundError(f"Safetensors shard from checkpoint index was not found: {shard_path}") + with safe_open(shard_path, framework="pt", device="cpu") as checkpoint: + shard_keys = set(checkpoint.keys()) + for name in by_shard[shard_name]: + if name not in shard_keys: + raise KeyError(f"Checkpoint index maps {name} to {shard_path}, but the tensor is absent") + tensor_slice = checkpoint.get_slice(name) + actual_shape = tuple(tensor_slice.get_shape()) + if actual_shape != expected_shapes[name]: + raise ValueError(f"Unexpected checkpoint shape for {name}: {actual_shape}, expected {expected_shapes[name]}") + checkpoint_dtypes.add(tensor_slice.get_dtype()) + if len(checkpoint_dtypes) != 1: + raise ValueError(f"MiniMax-H3 Qwen3-VL weights must use one floating dtype, got {sorted(checkpoint_dtypes)}") + logger.info( + "Preflighted {} native Qwen3-VL tensors (embedding + layers 0..{}) from {} shards for disk streaming", + len(modules), + MINIMAX_H3_TEXT_ENCODER_LAYER - 1, + len(by_shard), + ) + return weight_map, checkpoint_dtypes.pop() + @staticmethod def _load_quantized_weights(backbone, checkpoint_path): text_encoder_host_pinned = backbone.config.get("text_encoder_host_pinned", True) @@ -1008,8 +1265,12 @@ def load_text_encoder(self): attn_type=attn_type, block_offload=self.block_offload, tp_group=self.tp_group, + disk_streaming=self.disk_streaming, ) - if quantized: + if self.disk_streaming: + weight_map, _ = self._preflight_native_checkpoint(text_encoder, checkpoint_path, text_config) + text_encoder.init_disk_streaming(checkpoint_path, weight_map) + elif quantized: self._load_quantized_weights(text_encoder, checkpoint_path) else: self._load_native_weights(text_encoder, checkpoint_path, text_config) @@ -1207,6 +1468,8 @@ def _encode_vision(self, input_ids, pixel_values, image_grid_thw, pixel_values_v @torch.inference_mode() def infer(self, prompt, image_list=None, references=None): """Return unbatched ``[tokens, 5120]`` conditioning and text tags.""" + if self.disk_streaming and (image_list or references is not None): + raise NotImplementedError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming currently supports text-only t2av prompts.") self._ensure_loaded() try: # Input encoding happens before DefaultRunner enters its main-model @@ -1261,7 +1524,10 @@ def infer(self, prompt, image_list=None, references=None): "text_token_tags": token_tags.to(prompt_embeds.device), } finally: - if self.block_offload: + if self.disk_streaming: + if self.release_block_offload_buffers and self.text_encoder is not None: + self.text_encoder.release_disk_streaming_buffer() + elif self.block_offload: if self.release_block_offload_buffers and self.text_encoder is not None: self.text_encoder.release_block_offload_buffers() elif self.cpu_offload and self.text_encoder is not None: diff --git a/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py b/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py new file mode 100644 index 000000000..7a9f3301a --- /dev/null +++ b/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py @@ -0,0 +1,370 @@ +import importlib.util +import json +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from safetensors.torch import save_file + +REPO_ROOT = Path(__file__).parents[3] + + +class _FakeWeightModule: + def __init__(self): + self._modules = {} + self._parameters = {} + + def add_module(self, name, module): + self._modules[name] = module + setattr(self, name, module) + + def load(self, weight_dict): + for module in self._modules.values(): + if hasattr(module, "load"): + module.load(weight_dict) + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + for module in self._modules.values(): + if hasattr(module, "load_state_dict"): + module.load_state_dict(destination, block_index, adapter_block_index) + + +class _FakeWeightModuleList(_FakeWeightModule): + def __init__(self, modules=None): + super().__init__() + self._list = [] + if modules is not None: + for module in modules: + self.append(module) + + def append(self, module): + self._list.append(module) + self.add_module(str(len(self._list) - 1), module) + + def __getitem__(self, index): + return self._list[index] + + def __len__(self): + return len(self._list) + + def __iter__(self): + return iter(self._list) + + +def _resolve_qwen_layer_name(name, layer_index): + prefix = "model.language_model.layers." + if not name.startswith(prefix): + return name + parts = name.split(".", 4) + if len(parts) == 5 and parts[3].isdigit(): + return f"{prefix}{int(layer_index)}.{parts[4]}" + return name + + +class _FakeLinear: + def __init__(self, weight_name, bias_name=None, create_cuda_buffer=False, **_kwargs): + self.weight_name = weight_name + self.bias_name = bias_name + self.create_cuda_buffer = create_cuda_buffer + self.base_attrs = [(weight_name, "weight", True)] + if bias_name is not None: + self.base_attrs.append((bias_name, "bias", False)) + + def load(self, weight_dict): + for name, attr_name, transpose in self.base_attrs: + tensor = weight_dict[name] + if transpose: + tensor = tensor.t() + if self.create_cuda_buffer: + setattr(self, f"{attr_name}_cuda_buffer", tensor.clone()) + else: + setattr(self, attr_name, tensor.clone()) + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + for name, attr_name, _transpose in self.base_attrs: + actual_name = _resolve_qwen_layer_name(name, block_index) + if actual_name in destination: + buffer = getattr(self, f"{attr_name}_cuda_buffer") + setattr(self, attr_name, buffer.copy_(destination[actual_name])) + + +class _FakeRMS: + def __init__(self, weight_name, create_cuda_buffer=False, **_kwargs): + self.weight_name = weight_name + self.create_cuda_buffer = create_cuda_buffer + self.base_attrs = [(weight_name, "weight", False)] + + def load(self, weight_dict): + tensor = weight_dict[self.weight_name] + if self.create_cuda_buffer: + self.weight_cuda_buffer = tensor.clone() + else: + self.weight = tensor.clone() + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + actual_name = _resolve_qwen_layer_name(self.weight_name, block_index) + if actual_name in destination: + self.weight = self.weight_cuda_buffer.copy_(destination[actual_name]) + + +class _FakeEmbedding: + def __init__(self, weight_name, *_args, **_kwargs): + self.weight_name = weight_name + self.weight = None + + def apply(self, input_indices): + return F.embedding(input_indices, self.weight) + + +class _FakeAttention: + def __init__(self, *_args, **_kwargs): + pass + + def apply(self, q, *_args, **_kwargs): + return q.reshape(q.shape[0], -1) + + +class _FakeLeaf: + def __init__(self, *_args, **_kwargs): + pass + + +class _FakeAttnWeightTemplate: + def __init__(self, *_args, **_kwargs): + pass + + +def _load_qwen_module(monkeypatch): + for package_name in [ + "lightx2v", + "lightx2v.common", + "lightx2v.common.modules", + "lightx2v.common.offload", + "lightx2v.common.ops", + "lightx2v.common.ops.attn", + "lightx2v.common.ops.embedding", + "lightx2v.common.ops.mm", + "lightx2v.common.ops.norm", + "lightx2v.models", + "lightx2v.models.input_encoders", + "lightx2v.models.input_encoders.hf", + "lightx2v.models.input_encoders.hf.minimax_h3", + "lightx2v.models.networks", + "lightx2v.models.networks.minimax_h3", + "lightx2v.models.networks.minimax_h3.weights", + "lightx2v.utils", + "lightx2v_platform", + "lightx2v_platform.base", + ]: + package = types.ModuleType(package_name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, package_name, package) + + weight_module = types.ModuleType("lightx2v.common.modules.weight_module") + weight_module.WeightModule = _FakeWeightModule + weight_module.WeightModuleList = _FakeWeightModuleList + monkeypatch.setitem(sys.modules, "lightx2v.common.modules.weight_module", weight_module) + + event_manager = types.ModuleType("lightx2v.common.offload.event_manager") + + class RaisingEventSlotWeightAsyncStreamManager: + def __init__(self, *_args, **_kwargs): + raise AssertionError("EventSlotWeightAsyncStreamManager must not be instantiated") + + event_manager.EventSlotWeightAsyncStreamManager = RaisingEventSlotWeightAsyncStreamManager + monkeypatch.setitem(sys.modules, "lightx2v.common.offload.event_manager", event_manager) + + attn_template = types.ModuleType("lightx2v.common.ops.attn.template") + attn_template.AttnWeightTemplate = _FakeAttnWeightTemplate + monkeypatch.setitem(sys.modules, "lightx2v.common.ops.attn.template", attn_template) + for module_name, class_name in [ + ("lightx2v.common.ops.attn.torch_sdpa", "TorchSDPAWeight"), + ("lightx2v.common.ops.embedding.embedding_weight", "EmbeddingWeight"), + ("lightx2v.common.ops.mm.mm_weight", "MMWeight"), + ("lightx2v.common.ops.norm.rms_norm_weight", "RMSWeightFP32Qwen"), + ]: + module = types.ModuleType(module_name) + setattr(module, class_name, _FakeLeaf) + monkeypatch.setitem(sys.modules, module_name, module) + + vision = types.ModuleType("lightx2v.models.input_encoders.hf.minimax_h3.qwen3vl_vision") + vision.MiniMaxH3Qwen3VLVisionTower = _FakeLeaf + monkeypatch.setitem(sys.modules, "lightx2v.models.input_encoders.hf.minimax_h3.qwen3vl_vision", vision) + + packing = types.ModuleType("lightx2v.models.networks.minimax_h3.packing") + packing.VIDEO_TAG = 2 + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.packing", packing) + + packing_ref = types.ModuleType("lightx2v.models.networks.minimax_h3.packing_ref2av") + packing_ref.build_ref2av_presentation = lambda *_args, **_kwargs: None + packing_ref.sample_reference_video_frames = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.packing_ref2av", packing_ref) + + tp = types.ModuleType("lightx2v.models.networks.minimax_h3.weights.tensor_parallel") + tp.MiniMaxH3TensorParallelLinear = _FakeLinear + tp.unwrap_tp_linear = lambda obj: obj + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.weights.tensor_parallel", tp) + + envs = types.ModuleType("lightx2v.utils.envs") + envs.GET_DTYPE = lambda: torch.bfloat16 + monkeypatch.setitem(sys.modules, "lightx2v.utils.envs", envs) + + registry = types.ModuleType("lightx2v.utils.registry_factory") + registry.ATTN_WEIGHT_REGISTER = {"torch_sdpa": _FakeAttention} + registry.EMBEDDING_WEIGHT_REGISTER = {"Default": _FakeEmbedding} + registry.MM_WEIGHT_REGISTER = {"Default": _FakeLinear} + registry.RMS_WEIGHT_REGISTER = {"fp32_variance_qwen": _FakeRMS} + monkeypatch.setitem(sys.modules, "lightx2v.utils.registry_factory", registry) + + global_var = types.ModuleType("lightx2v_platform.base.global_var") + global_var.AI_DEVICE = "cpu" + monkeypatch.setitem(sys.modules, "lightx2v_platform.base.global_var", global_var) + + spec = importlib.util.spec_from_file_location( + "qwen3vl_disk_streaming_under_test", + REPO_ROOT / "lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py", + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + module.AI_DEVICE = "cpu" + module.torch_device_module = SimpleNamespace(synchronize=lambda: None) + return module + + +@pytest.fixture() +def qwen_module(monkeypatch): + return _load_qwen_module(monkeypatch) + + +def _tiny_text_config(): + return { + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "vocab_size": 32, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + } + + +def _make_backbone(qwen_module): + return qwen_module._Qwen3VLTextBackboneWeights( + {"text_encoder_disk_streaming": True}, + _tiny_text_config(), + num_layers=2, + attn_type="torch_sdpa", + block_offload=False, + disk_streaming=True, + ) + + +def _tensors_for_names(names, value): + tensors = {} + for name in names: + if name.endswith("embed_tokens.weight"): + tensors[name] = torch.full((32, 8), value, dtype=torch.bfloat16) + elif name.endswith(".weight") and any(proj in name for proj in (".q_proj.", ".k_proj.", ".v_proj.", ".o_proj.", ".gate_proj.", ".up_proj.", ".down_proj.")): + tensors[name] = torch.full((2, 3), value, dtype=torch.bfloat16) + else: + tensors[name] = torch.full((3,), value, dtype=torch.bfloat16) + return tensors + + +def _write_fake_checkpoint(tmp_path, backbone, qwen_module): + template_layer = backbone.layers[0] + embedding_name = backbone.embed_tokens.weight_name + layer0_names = qwen_module._Qwen3VLTextBackboneWeights._layer_tensor_names(template_layer, 0) + layer1_names = qwen_module._Qwen3VLTextBackboneWeights._layer_tensor_names(template_layer, 1) + tensors = {} + tensors.update(_tensors_for_names((embedding_name,), 5)) + tensors.update(_tensors_for_names(layer0_names, 1)) + tensors.update(_tensors_for_names(layer1_names, 2)) + + names = sorted(tensors) + shard_1_names = set(names[::2]) + shard_1 = {name: tensors[name] for name in names if name in shard_1_names} + shard_2 = {name: tensors[name] for name in names if name not in shard_1_names} + save_file(shard_1, tmp_path / "model-00001-of-00002.safetensors") + save_file(shard_2, tmp_path / "model-00002-of-00002.safetensors") + weight_map = { + **{name: "model-00001-of-00002.safetensors" for name in shard_1}, + **{name: "model-00002-of-00002.safetensors" for name in shard_2}, + } + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": 0}, "weight_map": weight_map}), + encoding="utf-8", + ) + return embedding_name, layer0_names, layer1_names, weight_map + + +def test_qwen3vl_disk_streaming_reuses_one_layer_buffer(tmp_path, monkeypatch, qwen_module): + backbone = _make_backbone(qwen_module) + embedding_name, layer0_names, layer1_names, weight_map = _write_fake_checkpoint(tmp_path, backbone, qwen_module) + calls = [] + original_loader = qwen_module._load_selected_checkpoint_tensors + + def spy_loader(text_encoder_path, weight_map_arg, names): + calls.append(tuple(names)) + return original_loader(text_encoder_path, weight_map_arg, names) + + monkeypatch.setattr(qwen_module, "_load_selected_checkpoint_tensors", spy_loader) + + backbone.init_disk_streaming(tmp_path, weight_map) + assert calls == [layer0_names] + assert backbone.offload_manager is None + layer0 = backbone.load_streaming_layer(0) + layer_id = id(layer0) + buffer = layer0.self_attn.q_proj.weight_cuda_buffer + buffer_id = id(buffer) + assert torch.all(layer0.self_attn.q_proj.weight == 1) + + layer1 = backbone.load_streaming_layer(1) + assert id(layer1) == layer_id + assert id(layer1.self_attn.q_proj.weight_cuda_buffer) == buffer_id + assert torch.all(layer1.self_attn.q_proj.weight == 2) + + hidden_states = backbone._forward_streaming_embedding(torch.tensor([0, 1], dtype=torch.long)) + assert hidden_states.dtype == torch.bfloat16 + assert tuple(hidden_states.shape) == (2, 8) + assert getattr(backbone.embed_tokens, "weight", None) is None + assert getattr(backbone.embed_tokens, "pin_weight", None) is None + + assert calls == [layer0_names, layer0_names, layer1_names, (embedding_name,)] + assert all(set(call) in [set(layer0_names), set(layer1_names), {embedding_name}] for call in calls) + + +def test_qwen3vl_disk_streaming_rejects_vision_inputs(tmp_path, qwen_module): + backbone = _make_backbone(qwen_module) + _embedding_name, _layer0_names, _layer1_names, weight_map = _write_fake_checkpoint(tmp_path, backbone, qwen_module) + backbone.init_disk_streaming(tmp_path, weight_map) + + with pytest.raises(NotImplementedError, match="text-only t2av"): + backbone.forward( + torch.tensor([0], dtype=torch.long), + vision_mask=torch.tensor([True]), + vision_embeds=torch.zeros((1, 8), dtype=torch.bfloat16), + ) + + +def test_qwen3vl_release_disk_streaming_buffer_clears_device_refs(tmp_path, qwen_module): + backbone = _make_backbone(qwen_module) + _embedding_name, _layer0_names, _layer1_names, weight_map = _write_fake_checkpoint(tmp_path, backbone, qwen_module) + backbone.init_disk_streaming(tmp_path, weight_map) + old_layer = backbone.streaming_layer + assert old_layer.self_attn.q_proj.weight_cuda_buffer is not None + + backbone.release_disk_streaming_buffer() + + assert backbone.streaming_layer is None + assert old_layer.self_attn.q_proj.weight is None + assert old_layer.self_attn.q_proj.weight_cuda_buffer is None + assert old_layer.input_layernorm.weight is None + assert old_layer.input_layernorm.weight_cuda_buffer is None From a7fb5f5047dced77578131b64774ed092a6b4e04 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Tue, 1 Sep 2026 20:05:45 +0800 Subject: [PATCH 07/31] feat(mps): stream MiniMax H3 DiT weights from disk --- .../models/networks/minimax_h3/checkpoint.py | 104 +++++ .../minimax_h3/infer/transformer_infer.py | 9 + lightx2v/models/networks/minimax_h3/model.py | 70 ++- .../minimax_h3/weights/transformer_weights.py | 88 +++- tests/models/minimax_h3/test_checkpoint.py | 121 ++++++ .../minimax_h3/test_model_disk_streaming.py | 401 ++++++++++++++++++ .../test_transformer_disk_streaming.py | 326 ++++++++++++++ 7 files changed, 1117 insertions(+), 2 deletions(-) create mode 100644 lightx2v/models/networks/minimax_h3/checkpoint.py create mode 100644 tests/models/minimax_h3/test_checkpoint.py create mode 100644 tests/models/minimax_h3/test_model_disk_streaming.py create mode 100644 tests/models/minimax_h3/test_transformer_disk_streaming.py diff --git a/lightx2v/models/networks/minimax_h3/checkpoint.py b/lightx2v/models/networks/minimax_h3/checkpoint.py new file mode 100644 index 000000000..1e75bc1ed --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/checkpoint.py @@ -0,0 +1,104 @@ +import json +import re +from collections import defaultdict +from pathlib import Path + +from safetensors import safe_open + +_H3_BLOCK_KEY_RE = re.compile(r"^transformer_blocks\.(\d+)\.") + + +class MiniMaxH3ShardCheckpoint: + """Synchronous reader for official MiniMax-H3 sharded safetensors.""" + + def __init__(self, checkpoint_dir): + self.checkpoint_dir = Path(checkpoint_dir) + self.index_path = self.checkpoint_dir / "model.safetensors.index.json" + if not self.index_path.is_file(): + raise FileNotFoundError(f"MiniMax-H3 safetensors index not found: {self.index_path}") + + with self.index_path.open("r", encoding="utf-8") as handle: + index = json.load(handle) + + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"MiniMax-H3 safetensors index must contain a dict weight_map: {self.index_path}") + if not weight_map: + raise ValueError(f"MiniMax-H3 safetensors index weight_map is empty: {self.index_path}") + + invalid_shard_names = sorted( + name + for name, shard_name in weight_map.items() + if not isinstance(shard_name, str) or not shard_name + ) + if invalid_shard_names: + raise ValueError( + f"MiniMax-H3 safetensors index contains invalid shard file names for tensors: {invalid_shard_names}" + ) + + self.weight_map = dict(weight_map) + shard_names = set(self.weight_map.values()) + missing_shards = sorted( + shard_name + for shard_name in shard_names + if not (self.checkpoint_dir / shard_name).is_file() + ) + if missing_shards: + raise FileNotFoundError( + f"MiniMax-H3 safetensors index references missing shard files: {missing_shards}" + ) + + @property + def tensor_names(self): + return tuple(sorted(self.weight_map)) + + @property + def block_indices(self): + return tuple( + sorted( + { + int(match.group(1)) + for name in self.weight_map + if (match := _H3_BLOCK_KEY_RE.match(name)) is not None + } + ) + ) + + def tensor_names_for_block(self, block_index): + block_prefix = f"transformer_blocks.{int(block_index)}." + return tuple(sorted(name for name in self.weight_map if name.startswith(block_prefix))) + + def non_block_tensor_names(self): + return tuple(sorted(name for name in self.weight_map if _H3_BLOCK_KEY_RE.match(name) is None)) + + def shard_for_tensor(self, name): + try: + return self.weight_map[name] + except KeyError as error: + raise KeyError(f"MiniMax-H3 checkpoint is missing requested tensor: {name}") from error + + def block_names(self, block_index): + return list(self.tensor_names_for_block(block_index)) + + def non_block_names(self): + return list(self.non_block_tensor_names()) + + def load_tensors(self, names, device="cpu"): + missing = sorted(name for name in names if name not in self.weight_map) + if missing: + raise KeyError(f"MiniMax-H3 checkpoint is missing requested tensors: {missing}") + + by_shard = defaultdict(list) + for name in names: + by_shard[self.weight_map[name]].append(name) + + tensors = {} + for shard_name in sorted(by_shard): + shard_path = self.checkpoint_dir / shard_name + with safe_open(shard_path, framework="pt", device=device) as shard: + for name in sorted(by_shard[shard_name]): + tensors[name] = shard.get_tensor(name) + return tensors + + +__all__ = ["MiniMaxH3ShardCheckpoint"] diff --git a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py index f7366425f..451981af7 100644 --- a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py @@ -169,7 +169,16 @@ def infer_without_offload(self, blocks, hidden_states, pre_infer_out): hidden_states = self.run_block(block_index, block, hidden_states, pre_infer_out) return hidden_states + def infer_with_disk_streaming(self, block_weights, hidden_states, pre_infer_out): + for block_index in block_weights.checkpoint.block_indices: + block = block_weights.load_streaming_block(block_index) + self.block_idx = block_index + hidden_states = self.run_block(block_index, block, hidden_states, pre_infer_out) + return hidden_states + def infer(self, block_weights, pre_infer_out): if self.use_adaln_cache: self._prepare_adaln_cache() + if getattr(block_weights, "disk_streaming", False): + return self.infer_with_disk_streaming(block_weights, pre_infer_out.hidden_states, pre_infer_out) return self.infer_func(block_weights.blocks, pre_infer_out.hidden_states, pre_infer_out) diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 19ca82476..2140dc04a 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -1,3 +1,4 @@ +import gc import glob import math import os @@ -39,6 +40,27 @@ } +def _collect_declared_base_tensor_names(*roots): + names = [] + seen = set() + stack = list(roots) + visited = set() + while stack: + obj = stack.pop() + if obj is None or id(obj) in visited: + continue + visited.add(id(obj)) + for name, _, _ in getattr(unwrap_tp_linear(obj), "base_attrs", ()): + if name.startswith("transformer_blocks."): + raise ValueError(f"MiniMax-H3 pre/post disk-streaming tensor list unexpectedly contains block tensor: {name}") + if name not in seen: + seen.add(name) + names.append(name) + stack.extend(getattr(obj, "_modules", {}).values()) + stack.extend(getattr(obj, "_parameters", {}).values()) + return tuple(sorted(names)) + + class MiniMaxH3Model(BaseTransformerModel): """LightX2V-native MiniMax-H3 joint audio/video transformer.""" @@ -68,6 +90,19 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0 raise ValueError("MiniMax-H3 dit_quant_scheme requires a dit_quantized_ckpt") if config.get("cpu_offload", False) and config.get("offload_granularity", "model") not in {"model", "block"}: raise NotImplementedError("MiniMax-H3 supports model and block CPU offload") + if config.get("dit_disk_streaming", False): + if not config.get("cpu_offload", False): + raise ValueError("MiniMax-H3 dit_disk_streaming requires cpu_offload=true.") + if config.get("offload_granularity", "model") != "block": + raise ValueError("MiniMax-H3 dit_disk_streaming requires offload_granularity='block'.") + if config.get("lazy_load", False): + raise NotImplementedError("MiniMax-H3 dit_disk_streaming cannot be combined with lazy_load.") + if config.get("dit_quantized", False): + raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support quantized DiT checkpoints yet.") + if config.get("tensor_parallel", False): + raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support tensor parallel inference yet.") + if lora_path is not None or config.get("lora_configs"): + raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support LoRA yet.") if config.get("attn_type") == "sol_attn": reorder = str(config.get("sol_attn_setting", {}).get("reorder", "none")).lower() if reorder != "none": @@ -105,6 +140,36 @@ def _apply_weights(self, weight_dict=None): self._h3_weight_shapes = {key: tuple(tensor.shape) for key, tensor in source.items() if isinstance(tensor, torch.Tensor) and tensor.ndim == 2} return super()._apply_weights(weight_dict) + def _init_weights(self, weight_dict=None): + if not self.config.get("dit_disk_streaming", False): + return super()._init_weights(weight_dict) + if weight_dict is not None: + raise ValueError("MiniMax-H3 dit_disk_streaming loads weights directly from the official checkpoint; explicit weight_dict is not supported.") + + self.pre_weight = self.pre_weight_class(self.config) + self.transformer_weights = self.transformer_weight_class(self.config) + self.post_weight = self.post_weight_class(self.config) + + checkpoint = self.transformer_weights.checkpoint + prepost_tensor_names = _collect_declared_base_tensor_names(self.pre_weight, self.post_weight) + missing = sorted(name for name in prepost_tensor_names if name not in checkpoint.weight_map) + if missing: + raise KeyError(f"MiniMax-H3 dit_disk_streaming checkpoint is missing pre/post tensors: {missing}") + + prepost_weights = checkpoint.load_tensors(prepost_tensor_names, device="cpu") + try: + self.pre_weight.load(prepost_weights) + self.post_weight.load(prepost_weights) + if prepost_weights: + raise RuntimeError(f"MiniMax-H3 dit_disk_streaming pre/post tensors were not consumed: {sorted(prepost_weights)}") + finally: + del prepost_weights + gc.collect() + device_module = getattr(torch, torch.device(self.device).type, None) + if device_module is not None and hasattr(device_module, "empty_cache"): + device_module.empty_cache() + return None + @staticmethod def _normalize_dynamic_lora_key(key): for prefix in ("base_model.model.", "model.diffusion_model.", "diffusion_model.", "transformer.", "model."): @@ -469,7 +534,10 @@ def _init_infer_class(self): if self.config.get("feature_caching", "NoCaching") != "NoCaching": raise NotImplementedError("MiniMax-H3 feature caching is not implemented") self.pre_infer_class = MiniMaxH3PreInfer - self.transformer_infer_class = MiniMaxH3OffloadTransformerInfer if self.cpu_offload else MiniMaxH3TransformerInfer + if self.config.get("dit_disk_streaming", False): + self.transformer_infer_class = MiniMaxH3TransformerInfer + else: + self.transformer_infer_class = MiniMaxH3OffloadTransformerInfer if self.cpu_offload else MiniMaxH3TransformerInfer self.post_infer_class = MiniMaxH3PostInfer def _init_infer(self): diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index 9634eb595..01736322f 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -2,10 +2,29 @@ import torch.distributed as dist from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList +from lightx2v.models.networks.minimax_h3.checkpoint import MiniMaxH3ShardCheckpoint from lightx2v.models.networks.minimax_h3.infer.triton_ops import MiniMaxH3TritonRope # noqa: F401 from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER +def _resolve_streaming_block_name(name, block_index): + block_prefix = "transformer_blocks." + if not name.startswith(block_prefix): + return name + parts = name.split(".", 2) + if len(parts) == 3 and parts[1].isdigit(): + return f"{block_prefix}{int(block_index)}.{parts[2]}" + return name + + +def _iter_base_attrs(module): + if hasattr(module, "base_attrs"): + yield from module.base_attrs + for child in getattr(module, "_modules", {}).values(): + if child is not None: + yield from _iter_base_attrs(child) + + def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None): lora_prefix = "transformer_blocks" if config.get("tensor_parallel", False) and tp_split is not None: @@ -129,11 +148,48 @@ def __init__(self, index, config, create_cuda_buffer=False): class MiniMaxH3TransformerWeights(WeightModule): def __init__(self, config, lazy_load_path=None, lora_path=None): super().__init__() + self.num_layers = int(config.get("num_layers", 50)) + self.disk_streaming = bool(config.get("dit_disk_streaming", False)) + if self.disk_streaming: + if config.get("lazy_load", False): + raise NotImplementedError( + "MiniMax-H3 dit_disk_streaming reads the official sharded checkpoint directly and cannot be combined with converted lazy_load block shards." + ) + if config.get("dit_quantized", False): + raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support quantized DiT checkpoints yet.") + if config.get("tensor_parallel", False): + raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support tensor parallel inference yet.") + if lora_path is not None: + raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support LoRA streaming yet.") + + checkpoint_dir = config.get("dit_original_ckpt") + if checkpoint_dir is None: + raise ValueError("MiniMax-H3 dit_disk_streaming requires config['dit_original_ckpt'] to point to the official transformer checkpoint directory.") + self.checkpoint = MiniMaxH3ShardCheckpoint(checkpoint_dir) + expected_block_indices = tuple(range(self.num_layers)) + if self.checkpoint.block_indices != expected_block_indices: + raise ValueError( + "MiniMax-H3 dit_disk_streaming checkpoint block indices mismatch: " + f"expected {expected_block_indices}, found {self.checkpoint.block_indices}" + ) + + self.blocks = WeightModuleList([]) + self.streaming_block = MiniMaxH3TransformerBlockWeights(0, config, create_cuda_buffer=True) + block0_tensors = self.checkpoint.load_tensors(self.checkpoint.tensor_names_for_block(0), device="cpu") + try: + self.streaming_block.load(block0_tensors) + self.streaming_block.load_state_dict(self._prepare_streaming_state_dict(block0_tensors, 0), 0) + finally: + del block0_tensors + self.add_module("blocks", self.blocks) + self.add_module("streaming_block", self.streaming_block) + return + if config.get("lazy_load", False): raise NotImplementedError( "MiniMax-H3 reads the official sharded checkpoint directly; disk lazy_load requires a converted block-sharded checkpoint and is not supported yet. Use lazy_load=false with model or block CPU offload." ) - self.blocks = WeightModuleList([MiniMaxH3TransformerBlockWeights(i, config) for i in range(int(config.get("num_layers", 50)))]) + self.blocks = WeightModuleList([MiniMaxH3TransformerBlockWeights(i, config) for i in range(self.num_layers)]) if config.get("cpu_offload", False) and config.get("offload_granularity", "model") == "block": self.offload_block_cuda_buffers = WeightModuleList([MiniMaxH3TransformerBlockWeights(i, config, create_cuda_buffer=True) for i in range(2)]) # Register device buffers before source blocks: buffer allocation @@ -141,3 +197,33 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): self.add_module("offload_block_cuda_buffers", self.offload_block_cuda_buffers) self.offload_phase_cuda_buffers = None self.add_module("blocks", self.blocks) + + @property + def streaming_block_indices(self): + if not self.disk_streaming: + raise RuntimeError("MiniMax-H3 streaming_block_indices is only available when dit_disk_streaming=true.") + return self.checkpoint.block_indices + + def load_streaming_block(self, block_index): + if not self.disk_streaming: + raise RuntimeError("MiniMax-H3 load_streaming_block requires dit_disk_streaming=true.") + block_index = int(block_index) + if block_index not in self.checkpoint.block_indices: + raise IndexError(f"MiniMax-H3 checkpoint does not contain transformer block {block_index}.") + + tensor_names = self.checkpoint.tensor_names_for_block(block_index) + tensors = self.checkpoint.load_tensors(tensor_names, device="cpu") + try: + self.streaming_block.load_state_dict(self._prepare_streaming_state_dict(tensors, block_index), block_index) + finally: + del tensors + return self.streaming_block + + def _prepare_streaming_state_dict(self, tensors, block_index): + state_dict = dict(tensors) + for name, _, transpose in _iter_base_attrs(self.streaming_block): + if transpose: + actual_name = _resolve_streaming_block_name(name, block_index) + if actual_name in state_dict: + state_dict[actual_name] = state_dict[actual_name].t() + return state_dict diff --git a/tests/models/minimax_h3/test_checkpoint.py b/tests/models/minimax_h3/test_checkpoint.py new file mode 100644 index 000000000..daa27fa17 --- /dev/null +++ b/tests/models/minimax_h3/test_checkpoint.py @@ -0,0 +1,121 @@ +import importlib.util +import json +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + + +def _load_checkpoint_class(): + module_path = Path(__file__).parents[3] / "lightx2v/models/networks/minimax_h3/checkpoint.py" + spec = importlib.util.spec_from_file_location("minimax_h3_checkpoint", module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.MiniMaxH3ShardCheckpoint + + +MiniMaxH3ShardCheckpoint = _load_checkpoint_class() + + +def _write_fake_checkpoint(tmp_path): + shard_1 = { + "proj_in.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3), + "transformer_blocks.0.attn.to_q.weight": torch.ones((2, 2), dtype=torch.bfloat16), + "transformer_blocks.1.attn.to_q.weight": torch.full((2, 2), 3, dtype=torch.bfloat16), + } + shard_2 = { + "transformer_blocks.0.ff.net.2.weight": torch.full((2, 2), 2, dtype=torch.bfloat16), + "norm_out.linear.weight": torch.full((2, 2), 4, dtype=torch.float32), + } + save_file(shard_1, tmp_path / "model-00001-of-00002.safetensors") + save_file(shard_2, tmp_path / "model-00002-of-00002.safetensors") + + weight_map = { + name: "model-00001-of-00002.safetensors" for name in shard_1 + } | { + name: "model-00002-of-00002.safetensors" for name in shard_2 + } + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": 0}, "weight_map": weight_map}), + encoding="utf-8", + ) + return shard_1, shard_2 + + +def test_tensor_block_and_non_block_names_are_deterministic(tmp_path): + _write_fake_checkpoint(tmp_path) + checkpoint = MiniMaxH3ShardCheckpoint(tmp_path) + + assert checkpoint.tensor_names == ( + "norm_out.linear.weight", + "proj_in.weight", + "transformer_blocks.0.attn.to_q.weight", + "transformer_blocks.0.ff.net.2.weight", + "transformer_blocks.1.attn.to_q.weight", + ) + assert checkpoint.block_indices == (0, 1) + assert checkpoint.tensor_names_for_block(0) == ( + "transformer_blocks.0.attn.to_q.weight", + "transformer_blocks.0.ff.net.2.weight", + ) + assert checkpoint.non_block_tensor_names() == ( + "norm_out.linear.weight", + "proj_in.weight", + ) + assert checkpoint.block_names(0) == list(checkpoint.tensor_names_for_block(0)) + assert checkpoint.non_block_names() == list(checkpoint.non_block_tensor_names()) + assert checkpoint.shard_for_tensor("transformer_blocks.0.attn.to_q.weight") == "model-00001-of-00002.safetensors" + assert checkpoint.shard_for_tensor("transformer_blocks.0.ff.net.2.weight") == "model-00002-of-00002.safetensors" + + +def test_load_tensors_reads_requested_block_tensors_across_shards(tmp_path): + shard_1, shard_2 = _write_fake_checkpoint(tmp_path) + checkpoint = MiniMaxH3ShardCheckpoint(tmp_path) + + tensors = checkpoint.load_tensors(checkpoint.tensor_names_for_block(0)) + + assert set(tensors) == set(checkpoint.tensor_names_for_block(0)) + assert torch.equal(tensors["transformer_blocks.0.attn.to_q.weight"], shard_1["transformer_blocks.0.attn.to_q.weight"]) + assert torch.equal(tensors["transformer_blocks.0.ff.net.2.weight"], shard_2["transformer_blocks.0.ff.net.2.weight"]) + + +def test_load_tensors_rejects_unknown_tensor(tmp_path): + _write_fake_checkpoint(tmp_path) + checkpoint = MiniMaxH3ShardCheckpoint(tmp_path) + + with pytest.raises(KeyError, match="missing requested tensors"): + checkpoint.load_tensors(["transformer_blocks.9.attn.to_q.weight"]) + + +def test_shard_for_tensor_rejects_unknown_tensor(tmp_path): + _write_fake_checkpoint(tmp_path) + checkpoint = MiniMaxH3ShardCheckpoint(tmp_path) + + with pytest.raises(KeyError, match="missing requested tensor"): + checkpoint.shard_for_tensor("missing") + + +def test_missing_index_raises_file_not_found(tmp_path): + with pytest.raises(FileNotFoundError, match="safetensors index not found"): + MiniMaxH3ShardCheckpoint(tmp_path) + + +def test_invalid_weight_map_shard_name_raises_value_error(tmp_path): + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"proj_in.weight": ""}}), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="invalid shard file names"): + MiniMaxH3ShardCheckpoint(tmp_path) + + +def test_missing_referenced_shard_raises_file_not_found(tmp_path): + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"proj_in.weight": "missing.safetensors"}}), + encoding="utf-8", + ) + + with pytest.raises(FileNotFoundError, match="missing shard files"): + MiniMaxH3ShardCheckpoint(tmp_path) diff --git a/tests/models/minimax_h3/test_model_disk_streaming.py b/tests/models/minimax_h3/test_model_disk_streaming.py new file mode 100644 index 000000000..22d330b62 --- /dev/null +++ b/tests/models/minimax_h3/test_model_disk_streaming.py @@ -0,0 +1,401 @@ +import importlib.util +import json +import sys +import types +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file + +REPO_ROOT = Path(__file__).parents[3] + + +class _FakeWeightModule: + def __init__(self): + self._modules = {} + self._parameters = {} + + def add_module(self, name, module): + self._modules[name] = module + setattr(self, name, module) + + def load(self, weight_dict): + for module in self._modules.values(): + if hasattr(module, "load"): + module.load(weight_dict) + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + for module in self._modules.values(): + if hasattr(module, "load_state_dict"): + module.load_state_dict(destination, block_index, adapter_block_index) + + def to_cuda(self): + pass + + def to_cpu(self): + pass + + +class _FakeWeightModuleList(_FakeWeightModule): + def __init__(self, modules=None): + super().__init__() + self._list = [] + if modules is not None: + for module in modules: + self.append(module) + + def append(self, module): + self._list.append(module) + self.add_module(str(len(self._list) - 1), module) + + def __getitem__(self, index): + return self._list[index] + + def __len__(self): + return len(self._list) + + def __iter__(self): + return iter(self._list) + + +def _resolve_block_name(name, block_index): + parts = name.split(".", 2) + if len(parts) == 3 and parts[0] == "transformer_blocks" and parts[1].isdigit(): + return f"transformer_blocks.{int(block_index)}.{parts[2]}" + return name + + +class _FakeLinear: + def __init__(self, weight_name, bias_name=None, create_cuda_buffer=False, **_kwargs): + self.weight_name = weight_name + self.bias_name = bias_name + self.create_cuda_buffer = create_cuda_buffer + self.base_attrs = [(weight_name, "weight", True)] + if bias_name is not None: + self.base_attrs.append((bias_name, "bias", False)) + + def load(self, weight_dict): + for name, attr_name, transpose in self.base_attrs: + tensor = weight_dict[name] + if transpose: + tensor = tensor.t() + if self.create_cuda_buffer: + setattr(self, f"{attr_name}_cuda_buffer", tensor.clone()) + else: + setattr(self, attr_name, tensor.clone()) + if tensor.device.type == "cpu": + setattr(self, f"pin_{attr_name}", tensor.clone()) + del weight_dict[name] + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + for name, attr_name, _transpose in self.base_attrs: + actual_name = _resolve_block_name(name, block_index) + if actual_name in destination: + buffer = getattr(self, f"{attr_name}_cuda_buffer") + setattr(self, attr_name, buffer.copy_(destination[actual_name])) + + +class _FakeRMS: + def __init__(self, weight_name, create_cuda_buffer=False, **_kwargs): + self.weight_name = weight_name + self.create_cuda_buffer = create_cuda_buffer + self.base_attrs = [(weight_name, "weight", False)] + + def load(self, weight_dict): + tensor = weight_dict[self.weight_name] + if self.create_cuda_buffer: + self.weight_cuda_buffer = tensor.clone() + else: + self.weight = tensor.clone() + self.pin_weight = tensor.clone() + if tensor.device.type == "cpu": + del weight_dict[self.weight_name] + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + actual_name = _resolve_block_name(self.weight_name, block_index) + if actual_name in destination: + self.weight = self.weight_cuda_buffer.copy_(destination[actual_name]) + + +class _FakeLeaf: + base_attrs = () + + def __init__(self, *_args, **_kwargs): + pass + + def set_config(self, *_args, **_kwargs): + pass + + +class _FakeBaseTransformerModel: + load_ckpt_called = False + init_weights_called = False + + def __init__(self, model_path, config, device, model_type=None, lora_path=None, lora_strength=1.0): + self.device = torch.device(device) + self.model_path = model_path + self.config = config + self.lora_path = lora_path + self.lora_strength = lora_strength + self.model_type = model_type + self.cpu_offload = config.get("cpu_offload", False) + self.offload_granularity = config.get("offload_granularity", "block") + self.lazy_load = config.get("lazy_load", False) + self.dit_quantized = config.get("dit_quantized", False) + self.use_tp = config.get("tensor_parallel", False) + self.tp_size = 1 + self.tp_rank = 0 + self.seq_p_group = None + self.sensitive_layer = {} + + def _init_weights(self, weight_dict=None): + _FakeBaseTransformerModel.init_weights_called = True + if weight_dict is None: + self._load_ckpt(False, {}) + + def _load_ckpt(self, unified_dtype, sensitive_layer): + _FakeBaseTransformerModel.load_ckpt_called = True + raise AssertionError("full checkpoint loading must not run in disk streaming") + + def _apply_weights(self, weight_dict=None): + pass + + def _init_offload_manager(self): + raise AssertionError("WeightAsyncStreamManager/offload manager must not initialize in disk streaming") + + +def _load_module(module_name, relative_path): + module_path = REPO_ROOT / relative_path + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture() +def h3_model_modules(monkeypatch): + for package_name in [ + "lightx2v", + "lightx2v.common", + "lightx2v.common.modules", + "lightx2v.models", + "lightx2v.models.networks", + "lightx2v.models.networks.minimax_h3", + "lightx2v.models.networks.minimax_h3.infer", + "lightx2v.models.networks.minimax_h3.weights", + "lightx2v.utils", + ]: + package = types.ModuleType(package_name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, package_name, package) + + weight_module = types.ModuleType("lightx2v.common.modules.weight_module") + weight_module.WeightModule = _FakeWeightModule + weight_module.WeightModuleList = _FakeWeightModuleList + monkeypatch.setitem(sys.modules, "lightx2v.common.modules.weight_module", weight_module) + + registry = types.ModuleType("lightx2v.utils.registry_factory") + registry.MM_WEIGHT_REGISTER = {"Default": _FakeLinear, "Default-ForceFp32": _FakeLinear} + registry.RMS_WEIGHT_REGISTER = {"torch_native": _FakeRMS} + registry.ROPE_REGISTER = {"torch_real_rope": _FakeLeaf} + registry.ATTN_WEIGHT_REGISTER = {"flash_attn3": _FakeLeaf} + monkeypatch.setitem(sys.modules, "lightx2v.utils.registry_factory", registry) + + envs = types.ModuleType("lightx2v.utils.envs") + envs.GET_DTYPE = lambda: torch.bfloat16 + monkeypatch.setitem(sys.modules, "lightx2v.utils.envs", envs) + + base_model = types.ModuleType("lightx2v.models.networks.base_model") + base_model.BaseTransformerModel = _FakeBaseTransformerModel + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.base_model", base_model) + + triton_ops = types.ModuleType("lightx2v.models.networks.minimax_h3.infer.triton_ops") + triton_ops.MiniMaxH3TritonRope = _FakeLeaf + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.infer.triton_ops", triton_ops) + + infer_module = types.ModuleType("lightx2v.models.networks.minimax_h3.infer.transformer_infer") + + class MiniMaxH3TransformerInfer: + def __init__(self, config): + self.config = config + + infer_module.MiniMaxH3TransformerInfer = MiniMaxH3TransformerInfer + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.infer.transformer_infer", infer_module) + + offload_module = types.ModuleType("lightx2v.models.networks.minimax_h3.infer.offload") + + class MiniMaxH3OffloadTransformerInfer: + pass + + offload_module.MiniMaxH3OffloadTransformerInfer = MiniMaxH3OffloadTransformerInfer + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.infer.offload", offload_module) + + for module_name, class_name in [ + ("module_io", "MiniMaxH3SequenceParallelState"), + ("post_infer", "MiniMaxH3PostInfer"), + ("pre_infer", "MiniMaxH3PreInfer"), + ]: + module = types.ModuleType(f"lightx2v.models.networks.minimax_h3.infer.{module_name}") + setattr(module, class_name, _FakeLeaf) + monkeypatch.setitem(sys.modules, f"lightx2v.models.networks.minimax_h3.infer.{module_name}", module) + + tensor_parallel = types.ModuleType("lightx2v.models.networks.minimax_h3.weights.tensor_parallel") + tensor_parallel.unwrap_tp_linear = lambda obj: obj + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.weights.tensor_parallel", tensor_parallel) + + checkpoint_module = _load_module( + "lightx2v.models.networks.minimax_h3.checkpoint", + "lightx2v/models/networks/minimax_h3/checkpoint.py", + ) + pre_module = _load_module( + "minimax_h3_pre_weights_for_model_test", + "lightx2v/models/networks/minimax_h3/weights/pre_weights.py", + ) + post_module = _load_module( + "minimax_h3_post_weights_for_model_test", + "lightx2v/models/networks/minimax_h3/weights/post_weights.py", + ) + transformer_module = _load_module( + "minimax_h3_transformer_weights_for_model_test", + "lightx2v/models/networks/minimax_h3/weights/transformer_weights.py", + ) + + weights_package = sys.modules["lightx2v.models.networks.minimax_h3.weights"] + weights_package.MiniMaxH3PreWeights = pre_module.MiniMaxH3PreWeights + weights_package.MiniMaxH3PostWeights = post_module.MiniMaxH3PostWeights + weights_package.MiniMaxH3TransformerWeights = transformer_module.MiniMaxH3TransformerWeights + + model_module = _load_module( + "minimax_h3_model_under_test", + "lightx2v/models/networks/minimax_h3/model.py", + ) + return checkpoint_module, pre_module, post_module, transformer_module, infer_module, offload_module, model_module + + +def _iter_base_attrs(module): + if hasattr(module, "base_attrs"): + yield from module.base_attrs + for child in getattr(module, "_modules", {}).values(): + yield from _iter_base_attrs(child) + + +def _tensors_from_roots(roots, block_index=None, value=1): + tensors = {} + for root in roots: + for name, _attr_name, transpose in _iter_base_attrs(root): + actual_name = _resolve_block_name(name, block_index) if block_index is not None else name + if transpose: + tensor = torch.full((2, 3), value, dtype=torch.bfloat16) + elif actual_name.endswith(".bias"): + tensor = torch.full((2,), value, dtype=torch.bfloat16) + else: + tensor = torch.full((3,), value, dtype=torch.bfloat16) + tensors[actual_name] = tensor + return tensors + + +def _write_fake_checkpoint(tmp_path, pre_module, post_module, transformer_module, num_layers=2): + config = {"num_layers": num_layers, "num_refiner_layers": 1} + tensors = {} + tensors.update(_tensors_from_roots([pre_module.MiniMaxH3PreWeights(config)], value=3)) + tensors.update(_tensors_from_roots([post_module.MiniMaxH3PostWeights(config)], value=4)) + block_template = transformer_module.MiniMaxH3TransformerBlockWeights(0, config) + for block_index in range(num_layers): + tensors.update(_tensors_from_roots([block_template], block_index=block_index, value=block_index + 1)) + + names = sorted(tensors) + shard_1_names = set(names[::2]) + shard_1 = {name: tensors[name] for name in names if name in shard_1_names} + shard_2 = {name: tensors[name] for name in names if name not in shard_1_names} + save_file(shard_1, tmp_path / "model-00001-of-00002.safetensors") + save_file(shard_2, tmp_path / "model-00002-of-00002.safetensors") + weight_map = { + **{name: "model-00001-of-00002.safetensors" for name in shard_1}, + **{name: "model-00002-of-00002.safetensors" for name in shard_2}, + } + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": 0}, "weight_map": weight_map}), + encoding="utf-8", + ) + return tensors + + +def _config(tmp_path, **overrides): + config = { + "seq_parallel": False, + "cfg_parallel": False, + "enable_cfg": False, + "cpu_offload": True, + "offload_granularity": "block", + "dit_disk_streaming": True, + "dit_original_ckpt": str(tmp_path), + "lazy_load": False, + "dit_quantized": False, + "dit_quant_scheme": "Default", + "tensor_parallel": False, + "num_layers": 2, + "num_refiner_layers": 1, + } + config.update(overrides) + return config + + +def test_non_disk_streaming_uses_base_init_weights(tmp_path, h3_model_modules): + _checkpoint_module, _pre_module, _post_module, _transformer_module, _infer_module, _offload_module, model_module = h3_model_modules + _FakeBaseTransformerModel.init_weights_called = False + _FakeBaseTransformerModel.load_ckpt_called = False + + with pytest.raises(AssertionError, match="full checkpoint loading"): + model_module.MiniMaxH3Model(str(tmp_path), _config(tmp_path, dit_disk_streaming=False), torch.device("cpu")) + + assert _FakeBaseTransformerModel.init_weights_called is True + assert _FakeBaseTransformerModel.load_ckpt_called is True + + +def test_disk_streaming_model_init_skips_full_checkpoint_load(tmp_path, monkeypatch, h3_model_modules): + checkpoint_module, pre_module, post_module, transformer_module, infer_module, offload_module, model_module = h3_model_modules + _write_fake_checkpoint(tmp_path, pre_module, post_module, transformer_module, num_layers=2) + calls = [] + + class SpyCheckpoint(checkpoint_module.MiniMaxH3ShardCheckpoint): + def load_tensors(self, names, device="cpu"): + calls.append(tuple(names)) + return super().load_tensors(names, device=device) + + monkeypatch.setattr(transformer_module, "MiniMaxH3ShardCheckpoint", SpyCheckpoint) + _FakeBaseTransformerModel.init_weights_called = False + _FakeBaseTransformerModel.load_ckpt_called = False + + model = model_module.MiniMaxH3Model(str(tmp_path), _config(tmp_path), torch.device("cpu")) + + assert _FakeBaseTransformerModel.load_ckpt_called is False + assert model.transformer_weights.disk_streaming is True + assert len(model.transformer_weights.blocks) == 0 + assert model.transformer_weights.streaming_block is not None + assert model.pre_weight.proj_in.pin_weight is not None + assert model.post_weight.proj_out.pin_weight is not None + assert model.transformer_infer_class is infer_module.MiniMaxH3TransformerInfer + assert model.transformer_infer_class is not offload_module.MiniMaxH3OffloadTransformerInfer + assert not hasattr(model.transformer_infer, "offload_manager") + + block0_names = model.transformer_weights.checkpoint.tensor_names_for_block(0) + block1_names = model.transformer_weights.checkpoint.tensor_names_for_block(1) + prepost_names = model_module._collect_declared_base_tensor_names(model.pre_weight, model.post_weight) + assert calls == [block0_names, prepost_names] + assert not any(set(call) == set(block1_names) for call in calls) + + +def test_disk_streaming_rejects_cpu_offload_false(tmp_path, h3_model_modules): + _checkpoint_module, _pre_module, _post_module, _transformer_module, _infer_module, _offload_module, model_module = h3_model_modules + + with pytest.raises(ValueError, match="requires cpu_offload=true"): + model_module.MiniMaxH3Model(str(tmp_path), _config(tmp_path, cpu_offload=False), torch.device("cpu")) + + +def test_disk_streaming_rejects_non_block_offload(tmp_path, h3_model_modules): + _checkpoint_module, _pre_module, _post_module, _transformer_module, _infer_module, _offload_module, model_module = h3_model_modules + + with pytest.raises(ValueError, match="requires offload_granularity='block'"): + model_module.MiniMaxH3Model(str(tmp_path), _config(tmp_path, offload_granularity="model"), torch.device("cpu")) diff --git a/tests/models/minimax_h3/test_transformer_disk_streaming.py b/tests/models/minimax_h3/test_transformer_disk_streaming.py new file mode 100644 index 000000000..d8d39d7a5 --- /dev/null +++ b/tests/models/minimax_h3/test_transformer_disk_streaming.py @@ -0,0 +1,326 @@ +import importlib.util +import json +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +REPO_ROOT = Path(__file__).parents[3] + + +class _FakeWeightModule: + def __init__(self): + self._modules = {} + self._parameters = {} + + def add_module(self, name, module): + self._modules[name] = module + setattr(self, name, module) + + def load(self, weight_dict): + for module in self._modules.values(): + if hasattr(module, "load"): + module.load(weight_dict) + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + for module in self._modules.values(): + if hasattr(module, "load_state_dict"): + module.load_state_dict(destination, block_index, adapter_block_index) + + +class _FakeWeightModuleList(_FakeWeightModule): + def __init__(self, modules=None): + super().__init__() + self._list = [] + if modules is not None: + for module in modules: + self.append(module) + + def append(self, module): + self._list.append(module) + self.add_module(str(len(self._list) - 1), module) + + def __getitem__(self, index): + return self._list[index] + + def __len__(self): + return len(self._list) + + def __iter__(self): + return iter(self._list) + + +def _resolve_block_name(name, block_index): + parts = name.split(".", 2) + if len(parts) == 3 and parts[0] == "transformer_blocks" and parts[1].isdigit(): + return f"transformer_blocks.{int(block_index)}.{parts[2]}" + return name + + +class _FakeLinear: + def __init__(self, weight_name, bias_name=None, create_cuda_buffer=False, **_kwargs): + self.weight_name = weight_name + self.bias_name = bias_name + self.create_cuda_buffer = create_cuda_buffer + self.base_attrs = [(weight_name, "weight", True)] + if bias_name is not None: + self.base_attrs.append((bias_name, "bias", False)) + + def load(self, weight_dict): + for name, attr_name, transpose in self.base_attrs: + tensor = weight_dict[name] + if transpose: + tensor = tensor.t() + if self.create_cuda_buffer: + setattr(self, f"{attr_name}_cuda_buffer", tensor.clone()) + else: + setattr(self, attr_name, tensor.clone()) + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + for name, attr_name, _transpose in self.base_attrs: + actual_name = _resolve_block_name(name, block_index) + if actual_name in destination: + buffer = getattr(self, f"{attr_name}_cuda_buffer") + setattr(self, attr_name, buffer.copy_(destination[actual_name])) + + +class _FakeRMS: + def __init__(self, weight_name, create_cuda_buffer=False, **_kwargs): + self.weight_name = weight_name + self.create_cuda_buffer = create_cuda_buffer + self.base_attrs = [(weight_name, "weight", False)] + + def load(self, weight_dict): + tensor = weight_dict[self.weight_name] + if self.create_cuda_buffer: + self.weight_cuda_buffer = tensor.clone() + else: + self.weight = tensor.clone() + + def load_state_dict(self, destination, block_index, adapter_block_index=None): + actual_name = _resolve_block_name(self.weight_name, block_index) + if actual_name in destination: + self.weight = self.weight_cuda_buffer.copy_(destination[actual_name]) + + +class _FakeLeaf: + base_attrs = () + + def __init__(self, *_args, **_kwargs): + pass + + def set_config(self, *_args, **_kwargs): + pass + + +def _load_module(module_name, relative_path): + module_path = REPO_ROOT / relative_path + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +@pytest.fixture() +def h3_modules(monkeypatch): + for package_name in [ + "lightx2v", + "lightx2v.common", + "lightx2v.common.modules", + "lightx2v.common.transformer_infer", + "lightx2v.models", + "lightx2v.models.networks", + "lightx2v.models.networks.minimax_h3", + "lightx2v.models.networks.minimax_h3.infer", + "lightx2v.utils", + ]: + package = types.ModuleType(package_name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, package_name, package) + + weight_module = types.ModuleType("lightx2v.common.modules.weight_module") + weight_module.WeightModule = _FakeWeightModule + weight_module.WeightModuleList = _FakeWeightModuleList + monkeypatch.setitem(sys.modules, "lightx2v.common.modules.weight_module", weight_module) + + registry = types.ModuleType("lightx2v.utils.registry_factory") + registry.MM_WEIGHT_REGISTER = {"Default": _FakeLinear} + registry.RMS_WEIGHT_REGISTER = {"torch_native": _FakeRMS} + registry.ROPE_REGISTER = {"torch_real_rope": _FakeLeaf} + registry.ATTN_WEIGHT_REGISTER = {"flash_attn3": _FakeLeaf} + monkeypatch.setitem(sys.modules, "lightx2v.utils.registry_factory", registry) + + triton_ops = types.ModuleType("lightx2v.models.networks.minimax_h3.infer.triton_ops") + triton_ops.MiniMaxH3TritonRope = _FakeLeaf + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.infer.triton_ops", triton_ops) + + checkpoint_module = _load_module( + "lightx2v.models.networks.minimax_h3.checkpoint", + "lightx2v/models/networks/minimax_h3/checkpoint.py", + ) + weights_module = _load_module( + "minimax_h3_transformer_weights_under_test", + "lightx2v/models/networks/minimax_h3/weights/transformer_weights.py", + ) + + base_infer = types.ModuleType("lightx2v.common.transformer_infer.transformer_infer") + + class _BaseTransformerInfer: + def init_compile(self, config): + self.use_compile = config.get("use_compile", False) + + def run_block(self, block_idx, block, *args): + return self.infer_block(block, *args) + + base_infer.BaseTransformerInfer = _BaseTransformerInfer + monkeypatch.setitem(sys.modules, "lightx2v.common.transformer_infer.transformer_infer", base_infer) + + envs = types.ModuleType("lightx2v.utils.envs") + envs.GET_DTYPE = lambda: torch.float32 + monkeypatch.setitem(sys.modules, "lightx2v.utils.envs", envs) + infer_module = _load_module( + "minimax_h3_transformer_infer_under_test", + "lightx2v/models/networks/minimax_h3/infer/transformer_infer.py", + ) + + return checkpoint_module, weights_module, infer_module + + +def _iter_base_attrs(module): + if hasattr(module, "base_attrs"): + yield from module.base_attrs + for child in getattr(module, "_modules", {}).values(): + yield from _iter_base_attrs(child) + + +def _block_tensors_from_template(block, block_index, value): + tensors = {} + for name, _attr_name, transpose in _iter_base_attrs(block): + actual_name = _resolve_block_name(name, block_index) + if transpose: + tensor = torch.full((2, 3), value, dtype=torch.float32) + elif actual_name.endswith(".bias"): + tensor = torch.full((2,), value, dtype=torch.float32) + else: + tensor = torch.full((3,), value, dtype=torch.float32) + tensors[actual_name] = tensor + return tensors + + +def _write_fake_official_checkpoint(tmp_path, weights_module, num_layers=2): + template = weights_module.MiniMaxH3TransformerBlockWeights(0, {"num_layers": num_layers}) + block_tensors = {} + for block_index in range(num_layers): + block_tensors.update(_block_tensors_from_template(template, block_index, block_index + 1)) + + names = sorted(block_tensors) + shard_1_names = set(names[::2]) + shard_1 = {name: block_tensors[name] for name in names if name in shard_1_names} + shard_2 = {name: block_tensors[name] for name in names if name not in shard_1_names} + save_file(shard_1, tmp_path / "model-00001-of-00002.safetensors") + save_file(shard_2, tmp_path / "model-00002-of-00002.safetensors") + + weight_map = { + **{name: "model-00001-of-00002.safetensors" for name in shard_1}, + **{name: "model-00002-of-00002.safetensors" for name in shard_2}, + } + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": 0}, "weight_map": weight_map}), + encoding="utf-8", + ) + return block_tensors + + +def test_transformer_weights_stream_official_shards_one_block_at_a_time(tmp_path, monkeypatch, h3_modules): + checkpoint_module, weights_module, _infer_module = h3_modules + _write_fake_official_checkpoint(tmp_path, weights_module, num_layers=2) + requested_names = [] + + class SpyCheckpoint(checkpoint_module.MiniMaxH3ShardCheckpoint): + def load_tensors(self, names, device="cpu"): + requested_names.append(tuple(names)) + return super().load_tensors(names, device=device) + + monkeypatch.setattr(weights_module, "MiniMaxH3ShardCheckpoint", SpyCheckpoint) + + weights = weights_module.MiniMaxH3TransformerWeights( + { + "dit_disk_streaming": True, + "dit_original_ckpt": str(tmp_path), + "num_layers": 2, + "dit_quantized": False, + "tensor_parallel": False, + } + ) + + assert weights.disk_streaming is True + assert len(weights.blocks) == 0 + assert weights.checkpoint.block_indices == (0, 1) + assert weights.streaming_block_indices == (0, 1) + assert weights.streaming_block.attn.to_q.weight_cuda_buffer.device.type == "cpu" + + block0_names = weights.checkpoint.tensor_names_for_block(0) + block1_names = weights.checkpoint.tensor_names_for_block(1) + assert requested_names == [block0_names] + + block0 = weights.load_streaming_block(0) + block_id = id(block0) + buffer = block0.attn.to_q.weight_cuda_buffer + buffer_id = id(buffer) + assert torch.all(block0.attn.to_q.weight == 1) + assert requested_names[-1] == block0_names + + block1 = weights.load_streaming_block(1) + assert id(block1) == block_id + assert id(block1.attn.to_q.weight_cuda_buffer) == buffer_id + assert torch.all(block1.attn.to_q.weight == 2) + assert block1.attn.to_q.weight.shape == (3, 2) + assert requested_names[-1] == block1_names + assert all(set(names) in [set(block0_names), set(block1_names)] for names in requested_names) + + +def test_transformer_disk_streaming_rejects_missing_block(tmp_path, h3_modules): + _checkpoint_module, weights_module, _infer_module = h3_modules + _write_fake_official_checkpoint(tmp_path, weights_module, num_layers=1) + + with pytest.raises(ValueError, match="checkpoint block indices mismatch"): + weights_module.MiniMaxH3TransformerWeights( + { + "dit_disk_streaming": True, + "dit_original_ckpt": str(tmp_path), + "num_layers": 2, + "dit_quantized": False, + "tensor_parallel": False, + } + ) + + +def test_transformer_infer_dispatches_to_disk_streaming(h3_modules): + _checkpoint_module, _weights_module, infer_module = h3_modules + infer = infer_module.MiniMaxH3TransformerInfer({"num_attention_heads": 1, "use_adaln_cache": False}) + loaded = [] + ran = [] + + class FakeBlockWeights: + disk_streaming = True + checkpoint = SimpleNamespace(block_indices=(0, 1)) + + def load_streaming_block(self, block_index): + loaded.append(block_index) + return f"block-{block_index}" + + def run_block(block_index, block, hidden_states, pre_infer_out): + ran.append((block_index, block, hidden_states)) + return hidden_states + block_index + 1 + + infer.run_block = run_block + pre_infer_out = SimpleNamespace(hidden_states=0) + + assert infer.infer(FakeBlockWeights(), pre_infer_out) == 3 + assert loaded == [0, 1] + assert ran == [(0, "block-0", 0), (1, "block-1", 1)] From 87dbd18bbf8d813b868604b7932979bd3feecd97 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Tue, 1 Sep 2026 20:13:45 +0800 Subject: [PATCH 08/31] fix(runtime): make cache cleanup device-aware --- lightx2v/models/networks/base_model.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py index 8bab34a75..5f645cd02 100644 --- a/lightx2v/models/networks/base_model.py +++ b/lightx2v/models/networks/base_model.py @@ -336,7 +336,9 @@ def _apply_weights(self, weight_dict=None): self._register_lora(self.lora_path, self.lora_strength) del self.original_weight_dict - torch.cuda.empty_cache() + device_module = getattr(torch, AI_DEVICE, None) + if device_module is not None and hasattr(device_module, "empty_cache"): + device_module.empty_cache() gc.collect() def _load_lora_file(self, file_path): From 0fb815cad2d47cc3cd42f15da957a093d742205a Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Tue, 1 Sep 2026 20:13:48 +0800 Subject: [PATCH 09/31] fix(mps): support MPS random seeding --- lightx2v/utils/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lightx2v/utils/utils.py b/lightx2v/utils/utils.py index 8e11aec69..36150c894 100755 --- a/lightx2v/utils/utils.py +++ b/lightx2v/utils/utils.py @@ -32,7 +32,8 @@ def seed_all(seed): np.random.seed(seed) torch.manual_seed(seed) torch_device_module.manual_seed(seed) - torch_device_module.manual_seed_all(seed) + if hasattr(torch_device_module, "manual_seed_all"): + torch_device_module.manual_seed_all(seed) torch.backends.cudnn.benchmark = False torch.backends.cudnn.deterministic = True From c2a3d0164c79580f34e07ccabf74cf7a794f729f Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Tue, 1 Sep 2026 22:01:59 +0800 Subject: [PATCH 10/31] feat(mps): add MiniMax H3 low-memory lifecycle --- .../minimax_h3/weights/transformer_weights.py | 73 +++- .../runners/minimax_h3/minimax_h3_runner.py | 42 ++- .../minimax_h3/test_runner_mps_low_memory.py | 318 ++++++++++++++++++ .../test_transformer_disk_streaming.py | 42 +++ 4 files changed, 463 insertions(+), 12 deletions(-) create mode 100644 tests/models/minimax_h3/test_runner_mps_low_memory.py diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index 01736322f..ea4200e90 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -1,3 +1,6 @@ +import gc +from contextlib import suppress + import torch import torch.distributed as dist @@ -5,6 +8,7 @@ from lightx2v.models.networks.minimax_h3.checkpoint import MiniMaxH3ShardCheckpoint from lightx2v.models.networks.minimax_h3.infer.triton_ops import MiniMaxH3TritonRope # noqa: F401 from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, ROPE_REGISTER +from lightx2v_platform.base.global_var import AI_DEVICE def _resolve_streaming_block_name(name, block_index): @@ -25,6 +29,14 @@ def _iter_base_attrs(module): yield from _iter_base_attrs(child) +def _empty_device_cache(): + if not isinstance(AI_DEVICE, str): + return + device_module = getattr(torch, AI_DEVICE, None) + if device_module is not None and hasattr(device_module, "empty_cache"): + device_module.empty_cache() + + def _linear(config, name, bias=False, create_cuda_buffer=False, tp_split=None): lora_prefix = "transformer_blocks" if config.get("tensor_parallel", False) and tp_split is not None: @@ -148,6 +160,7 @@ def __init__(self, index, config, create_cuda_buffer=False): class MiniMaxH3TransformerWeights(WeightModule): def __init__(self, config, lazy_load_path=None, lora_path=None): super().__init__() + self.config = config self.num_layers = int(config.get("num_layers", 50)) self.disk_streaming = bool(config.get("dit_disk_streaming", False)) if self.disk_streaming: @@ -174,15 +187,9 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): ) self.blocks = WeightModuleList([]) - self.streaming_block = MiniMaxH3TransformerBlockWeights(0, config, create_cuda_buffer=True) - block0_tensors = self.checkpoint.load_tensors(self.checkpoint.tensor_names_for_block(0), device="cpu") - try: - self.streaming_block.load(block0_tensors) - self.streaming_block.load_state_dict(self._prepare_streaming_state_dict(block0_tensors, 0), 0) - finally: - del block0_tensors + self.streaming_block = None self.add_module("blocks", self.blocks) - self.add_module("streaming_block", self.streaming_block) + self._ensure_streaming_block() return if config.get("lazy_load", False): @@ -211,6 +218,7 @@ def load_streaming_block(self, block_index): if block_index not in self.checkpoint.block_indices: raise IndexError(f"MiniMax-H3 checkpoint does not contain transformer block {block_index}.") + self._ensure_streaming_block() tensor_names = self.checkpoint.tensor_names_for_block(block_index) tensors = self.checkpoint.load_tensors(tensor_names, device="cpu") try: @@ -219,6 +227,55 @@ def load_streaming_block(self, block_index): del tensors return self.streaming_block + def _ensure_streaming_block(self): + if self.streaming_block is not None: + return + self.streaming_block = MiniMaxH3TransformerBlockWeights(0, self.config, create_cuda_buffer=True) + self.add_module("streaming_block", self.streaming_block) + block0_tensors = self.checkpoint.load_tensors(self.checkpoint.tensor_names_for_block(0), device="cpu") + try: + self.streaming_block.load(block0_tensors) + self.streaming_block.load_state_dict(self._prepare_streaming_state_dict(block0_tensors, 0), 0) + finally: + del block0_tensors + gc.collect() + _empty_device_cache() + + def release_disk_streaming_buffer(self): + if not self.disk_streaming: + return + block = self.streaming_block + if block is None: + return + with suppress(Exception): + device_module = getattr(torch, AI_DEVICE, None) + if device_module is not None and hasattr(device_module, "synchronize"): + device_module.synchronize() + + stack = [block] + visited = set() + while stack: + module = stack.pop() + if module is None or id(module) in visited: + continue + visited.add(id(module)) + for _, attr_name, _ in getattr(module, "base_attrs", ()): + if hasattr(module, attr_name): + setattr(module, attr_name, None) + buffer_attr = f"{attr_name}_cuda_buffer" + if hasattr(module, buffer_attr): + setattr(module, buffer_attr, None) + for attr_name in tuple(vars(module)): + if attr_name.endswith("_cuda_buffer"): + setattr(module, attr_name, None) + stack.extend(getattr(module, "_modules", {}).values()) + stack.extend(getattr(module, "_parameters", {}).values()) + + self.streaming_block = None + self._modules["streaming_block"] = None + gc.collect() + _empty_device_cache() + def _prepare_streaming_state_dict(self, tensors, block_index): state_dict = dict(tensors) for name, _, transpose in _iter_base_attrs(self.streaming_block): diff --git a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py index 7cd2698ea..1f6c37559 100644 --- a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py +++ b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py @@ -1,3 +1,4 @@ +import gc import os from contextlib import suppress @@ -194,10 +195,30 @@ def init_scheduler(self): @ProfilingContext4DebugL2("Load models") def load_model(self): + if self._is_mps_low_memory_streaming(): + self._validate_mps_low_memory_streaming_config() self.model = self.load_transformer() self.text_encoders = self.load_text_encoder() + if self._is_mps_low_memory_streaming(): + self.video_vae = None + self.audio_vae = None + return self.video_vae, self.audio_vae = self.load_vae() + def _is_mps_low_memory_streaming(self): + return ( + AI_DEVICE == "mps" + and self.config.get("task") == "t2av" + and self.config.get("dit_disk_streaming", False) + and self.config.get("text_encoder_disk_streaming", False) + ) + + def _validate_mps_low_memory_streaming_config(self): + if not self.config.get("text_encoder_release_block_offload_buffers", False): + raise ValueError("MiniMax-H3 MPS low-memory streaming requires text_encoder_release_block_offload_buffers=true.") + if self.config.get("warmup", False): + raise ValueError("MiniMax-H3 MPS low-memory streaming requires warmup=false in the first implementation.") + def load_transformer(self): model_kwargs = { "model_path": self.config["model_path"], @@ -219,14 +240,14 @@ def load_text_encoder(self): @staticmethod def _validate_vae_decode_tile_shapes(tile_shapes, video_vae): if not isinstance(tile_shapes, dict): - raise ValueError("vae_decode_tile_shape must map 'HEIGHTxWIDTH' to [tile_height, tile_width]") + raise TypeError("vae_decode_tile_shape must map 'HEIGHTxWIDTH' to [tile_height, tile_width]") ratio = video_vae.spatial_compression_ratio overlap_height = video_vae.tile_sample_min_overlap_height overlap_width = video_vae.tile_sample_min_overlap_width for resolution, tile_shape in tile_shapes.items(): if not isinstance(resolution, str): - raise ValueError(f"invalid VAE tile resolution: {resolution!r}") + raise TypeError(f"invalid VAE tile resolution: {resolution!r}") dimensions = resolution.split("x") if len(dimensions) != 2: raise ValueError(f"invalid VAE tile resolution: {resolution!r}") @@ -594,7 +615,13 @@ def _offload_transformer(self): if not self.config.get("cpu_offload", False): return if self.model.block_offload: - if not self.model.prepost_resident: + if self._is_mps_low_memory_streaming() and self.config.get("dit_disk_streaming", False): + logger.info("Offloading MiniMax-H3 pre/post weights and releasing the disk-streaming DiT device block buffer") + if not self.model.prepost_resident: + self.model.pre_weight.to_cpu() + self.model.post_weight.to_cpu() + self.model.transformer_weights.release_disk_streaming_buffer() + elif not self.model.prepost_resident: logger.info("Offloading MiniMax-H3 pre/post weights; retaining the two block-offload device buffers") self.model.pre_weight.to_cpu() self.model.post_weight.to_cpu() @@ -611,6 +638,8 @@ def _offload_transformer(self): metrics_labels=["MiniMaxH3Runner"], ) def run_vae_decoder(self, video_rows, audio_rows): + if self._is_mps_low_memory_streaming() and (self.video_vae is None or self.audio_vae is None): + self.video_vae, self.audio_vae = self.load_vae() video_rows = video_rows[self.scheduler.num_condition_video_rows :] audio_rows = audio_rows[self.scheduler.num_condition_audio_rows :] video_latents = unpatchify_video_tokens( @@ -706,8 +735,13 @@ def run_main(self): with suppress(Exception): self._offload_transformer() try: - self.end_run() + if self._is_mps_low_memory_streaming() and (self.video_vae is not None or self.audio_vae is not None): + self.video_vae = None + self.audio_vae = None + gc.collect() + self.maybe_empty_cache(force=True, collect_garbage=True) finally: + self.end_run() # Decoded FP32 video is large (roughly 1.5 GiB at the default # shape). Returned tensors keep their own references/copies; # the runner should not retain another request-sized result. diff --git a/tests/models/minimax_h3/test_runner_mps_low_memory.py b/tests/models/minimax_h3/test_runner_mps_low_memory.py new file mode 100644 index 000000000..78c61d9a1 --- /dev/null +++ b/tests/models/minimax_h3/test_runner_mps_low_memory.py @@ -0,0 +1,318 @@ +import importlib.util +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +REPO_ROOT = Path(__file__).parents[3] + + +class _Profiler: + def __init__(self, *_args, **_kwargs): + pass + + def __call__(self, func): + return func + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + +def _install_module(monkeypatch, name, **attrs): + module = types.ModuleType(name) + for attr_name, value in attrs.items(): + setattr(module, attr_name, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_runner_module(monkeypatch): + for package_name in [ + "lightx2v", + "lightx2v.models", + "lightx2v.models.audio_encoders", + "lightx2v.models.audio_encoders.hf", + "lightx2v.models.input_encoders", + "lightx2v.models.input_encoders.hf", + "lightx2v.models.networks", + "lightx2v.models.networks.minimax_h3", + "lightx2v.models.runners", + "lightx2v.models.runners.default_runner", + "lightx2v.models.schedulers", + "lightx2v.models.video_encoders", + "lightx2v.models.video_encoders.hf", + "lightx2v.models.video_encoders.hf.ltx2", + "lightx2v.models.video_encoders.hf.ltx2.audio_vae", + "lightx2v.server", + "lightx2v.utils", + "lightx2v_platform", + "lightx2v_platform.base", + ]: + package = types.ModuleType(package_name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, package_name, package) + + class DefaultRunner: + def __init__(self, config): + self.config = config + + def maybe_empty_cache(self, **_kwargs): + return False + + def end_run(self): + pass + + _install_module(monkeypatch, "lightx2v.models.runners.default_runner", DefaultRunner=DefaultRunner) + _install_module(monkeypatch, "lightx2v.models.audio_encoders.hf.minimax_h3", MiniMaxH3AudioVAE=object) + _install_module(monkeypatch, "lightx2v.models.input_encoders.hf.minimax_h3", MiniMaxH3Qwen3VLTextEncoder=object) + _install_module(monkeypatch, "lightx2v.models.networks.minimax_h3.lora", MiniMaxH3LoraAdapter=object) + _install_module(monkeypatch, "lightx2v.models.networks.minimax_h3.model", MiniMaxH3Model=object) + _install_module(monkeypatch, "lightx2v.models.schedulers.minimax_h3", MiniMaxH3Scheduler=object) + _install_module(monkeypatch, "lightx2v.models.video_encoders.hf.minimax_h3", MiniMaxH3VideoVAE=object) + + class Audio: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + _install_module(monkeypatch, "lightx2v.models.video_encoders.hf.ltx2.audio_vae.ops", Audio=Audio) + class _Metrics: + def __getattr__(self, _name): + return None + + _install_module(monkeypatch, "lightx2v.server.metrics", monitor_cli=_Metrics()) + _install_module(monkeypatch, "lightx2v.utils.envs", DTYPE_MAP={"fp32": torch.float32}, GET_RECORDER_MODE=lambda: None) + _install_module( + monkeypatch, + "lightx2v.utils.input_info", + FL2AVInputInfo=object, + I2AVInputInfo=object, + L2AVInputInfo=object, + Ref2AVInputInfo=object, + T2AVInputInfo=object, + ) + _install_module(monkeypatch, "lightx2v.utils.ltx2_media_io", encode_video=lambda **_kwargs: None) + _install_module(monkeypatch, "lightx2v.utils.profiler", ProfilingContext4DebugL1=_Profiler, ProfilingContext4DebugL2=_Profiler) + _install_module(monkeypatch, "lightx2v.utils.registry_factory", RUNNER_REGISTER=lambda _name: lambda cls: cls) + _install_module(monkeypatch, "lightx2v_platform.base.global_var", AI_DEVICE="mps") + + packing_names = { + "TEXT_TAG": 1, + "align_num_frames": lambda value: value, + "prepare_keyframe_image": lambda image, *_args, **_kwargs: image, + "resolve_canvas_size": lambda width, height: (height, width), + "unpack_audio_tokens": lambda rows, *_args, **_kwargs: rows, + "unpatchify_video_tokens": lambda rows, *_args, **_kwargs: rows, + "validate_t2av_geometry": lambda *_args, **_kwargs: None, + } + _install_module(monkeypatch, "lightx2v.models.networks.minimax_h3.packing", **packing_names) + _install_module( + monkeypatch, + "lightx2v.models.networks.minimax_h3.packing_ref2av", + DEFAULT_REFERENCE_IMAGE_RESIZE_MODE="contain", + MAX_REFERENCES=12, + MAX_REFERENCE_AUDIOS=3, + MAX_REFERENCE_IMAGES=9, + MAX_REFERENCE_VIDEOS=3, + REFERENCE_IMAGE_RESIZE_MODES=("contain",), + MiniMaxH3PreparedReference=object, + decode_reference_audio=lambda *_args, **_kwargs: None, + decode_reference_video=lambda *_args, **_kwargs: None, + prepare_reference_frames=lambda frames, *_args, **_kwargs: frames, + prepare_reference_image=lambda image, *_args, **_kwargs: image, + prepare_reference_waveform=lambda waveform, *_args, **_kwargs: waveform, + resample_reference_frames=lambda frames, *_args, **_kwargs: frames, + resolve_reference_image_size=lambda width, height, **_kwargs: (height, width), + trim_reference_num_frames=lambda value: value, + ) + + module_path = REPO_ROOT / "lightx2v/models/runners/minimax_h3/minimax_h3_runner.py" + spec = importlib.util.spec_from_file_location("minimax_h3_runner_under_test", module_path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + module.torch_device_module = SimpleNamespace(synchronize=lambda: None) + return module + + +@pytest.fixture() +def runner_module(monkeypatch): + return _load_runner_module(monkeypatch) + + +def _low_memory_config(**overrides): + config = { + "task": "t2av", + "dit_disk_streaming": True, + "text_encoder_disk_streaming": True, + "text_encoder_release_block_offload_buffers": True, + "warmup": False, + "cpu_offload": True, + } + config.update(overrides) + return config + + +def _make_runner(runner_module, config): + runner = object.__new__(runner_module.MiniMaxH3Runner) + runner.config = config + return runner + + +def test_low_memory_load_model_defers_vae(runner_module): + runner = _make_runner(runner_module, _low_memory_config()) + calls = [] + runner.load_transformer = lambda: calls.append("transformer") or object() + runner.load_text_encoder = lambda: calls.append("text_encoder") or [object()] + runner.load_vae = lambda: calls.append("vae") or (object(), object()) + + runner.load_model() + + assert calls == ["transformer", "text_encoder"] + assert runner.video_vae is None + assert runner.audio_vae is None + + +def test_non_low_memory_load_model_keeps_eager_vae(runner_module): + runner = _make_runner(runner_module, _low_memory_config(text_encoder_disk_streaming=False)) + video_vae = object() + audio_vae = object() + calls = [] + runner.load_transformer = lambda: calls.append("transformer") or object() + runner.load_text_encoder = lambda: calls.append("text_encoder") or [object()] + runner.load_vae = lambda: calls.append("vae") or (video_vae, audio_vae) + + runner.load_model() + + assert calls == ["transformer", "text_encoder", "vae"] + assert runner.video_vae is video_vae + assert runner.audio_vae is audio_vae + + +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"text_encoder_release_block_offload_buffers": False}, "text_encoder_release_block_offload_buffers=true"), + ({"warmup": True}, "warmup=false"), + ], +) +def test_low_memory_load_model_rejects_unsupported_first_version_configs(runner_module, override, message): + runner = _make_runner(runner_module, _low_memory_config(**override)) + runner.load_transformer = lambda: object() + runner.load_text_encoder = lambda: [object()] + runner.load_vae = lambda: (object(), object()) + + with pytest.raises(ValueError, match=message): + runner.load_model() + + +def test_offload_transformer_releases_disk_streaming_buffer(runner_module): + runner = _make_runner(runner_module, _low_memory_config()) + calls = [] + runner.maybe_empty_cache = lambda **kwargs: calls.append(("empty", kwargs)) + runner.model = SimpleNamespace( + block_offload=True, + prepost_resident=False, + pre_weight=SimpleNamespace(to_cpu=lambda: calls.append("pre_cpu")), + post_weight=SimpleNamespace(to_cpu=lambda: calls.append("post_cpu")), + transformer_weights=SimpleNamespace(release_disk_streaming_buffer=lambda: calls.append("release_dit")), + ) + + runner._offload_transformer() + + assert calls == ["pre_cpu", "post_cpu", "release_dit", ("empty", {"force": True, "collect_garbage": True})] + + +def test_offload_transformer_preserves_regular_block_offload_behavior(runner_module): + runner = _make_runner(runner_module, _low_memory_config(dit_disk_streaming=False)) + calls = [] + runner.maybe_empty_cache = lambda **kwargs: calls.append(("empty", kwargs)) + runner.model = SimpleNamespace( + block_offload=True, + prepost_resident=False, + pre_weight=SimpleNamespace(to_cpu=lambda: calls.append("pre_cpu")), + post_weight=SimpleNamespace(to_cpu=lambda: calls.append("post_cpu")), + transformer_weights=SimpleNamespace(release_disk_streaming_buffer=lambda: calls.append("release_dit")), + ) + + runner._offload_transformer() + + assert calls == ["pre_cpu", "post_cpu", ("empty", {"force": True, "collect_garbage": True})] + + +def test_run_vae_decoder_lazy_loads_once(runner_module): + runner = _make_runner(runner_module, _low_memory_config()) + calls = [] + video_vae = SimpleNamespace( + decode_parallel=False, + decode=lambda latents: ("video", latents), + ) + audio_vae = SimpleNamespace(decode=lambda latents: ("audio", latents)) + runner.load_vae = lambda: calls.append("load_vae") or (video_vae, audio_vae) + runner.video_vae = None + runner.audio_vae = None + runner._vae_decode_tile_shapes = {} + runner.scheduler = SimpleNamespace( + num_condition_video_rows=0, + num_condition_audio_rows=0, + num_latent_frames=1, + latent_height=1, + latent_width=1, + num_audio_latents=1, + ) + + first = runner.run_vae_decoder(torch.tensor([1]), torch.tensor([2])) + second = runner.run_vae_decoder(torch.tensor([3]), torch.tensor([4])) + + assert calls == ["load_vae"] + assert first[0][0] == "video" + assert first[1][0] == "audio" + assert torch.equal(first[0][1], torch.tensor([1])) + assert torch.equal(first[1][1], torch.tensor([2])) + assert second[0][0] == "video" + assert second[1][0] == "audio" + assert torch.equal(second[0][1], torch.tensor([3])) + assert torch.equal(second[1][1], torch.tensor([4])) + + +def test_run_main_releases_vae_after_processing_result(runner_module): + runner = _make_runner(runner_module, _low_memory_config()) + calls = [] + runner.maybe_empty_cache = lambda **kwargs: calls.append(("empty", kwargs)) + runner.init_run = lambda: calls.append("init") + runner.run_segment = lambda _segment: calls.append("dit") or ("video_rows", "audio_rows") + runner._offload_transformer = lambda: calls.append("offload_dit") + runner.run_vae_decoder = lambda *_args: calls.append("decode") or ("decoded_video", "decoded_audio") + + def process(): + calls.append(("process", runner.video_vae, runner.audio_vae)) + return "result" + + runner.process_images_after_vae_decoder = process + runner.end_run = lambda: calls.append("end_run") + runner.video_vae = object() + runner.audio_vae = object() + + assert runner.run_main() == "result" + + assert [call if isinstance(call, str) else call[0] for call in calls] == [ + "init", + "dit", + "offload_dit", + "decode", + "process", + "empty", + "end_run", + ] + process_call = calls[4] + assert process_call[1] is not None + assert process_call[2] is not None + assert runner.video_vae is None + assert runner.audio_vae is None + assert runner.gen_video is None + assert runner.gen_audio is None diff --git a/tests/models/minimax_h3/test_transformer_disk_streaming.py b/tests/models/minimax_h3/test_transformer_disk_streaming.py index d8d39d7a5..5feb63701 100644 --- a/tests/models/minimax_h3/test_transformer_disk_streaming.py +++ b/tests/models/minimax_h3/test_transformer_disk_streaming.py @@ -284,6 +284,48 @@ def load_tensors(self, names, device="cpu"): assert all(set(names) in [set(block0_names), set(block1_names)] for names in requested_names) +def test_transformer_weights_release_and_reinitialize_streaming_block(tmp_path, h3_modules): + _checkpoint_module, weights_module, _infer_module = h3_modules + _write_fake_official_checkpoint(tmp_path, weights_module, num_layers=2) + + weights = weights_module.MiniMaxH3TransformerWeights( + { + "dit_disk_streaming": True, + "dit_original_ckpt": str(tmp_path), + "num_layers": 2, + "dit_quantized": False, + "tensor_parallel": False, + } + ) + old_block = weights.streaming_block + assert old_block.attn.to_q.weight_cuda_buffer is not None + + weights.release_disk_streaming_buffer() + + assert weights.streaming_block is None + assert old_block.attn.to_q.weight is None + assert old_block.attn.to_q.weight_cuda_buffer is None + assert old_block.adaln.bias is None + assert old_block.adaln.bias_cuda_buffer is None + assert old_block.norm1.weight is None + assert old_block.norm1.weight_cuda_buffer is None + + block0 = weights.load_streaming_block(0) + assert block0 is weights.streaming_block + assert block0 is not old_block + assert torch.all(block0.attn.to_q.weight == 1) + assert block0.attn.to_q.weight_cuda_buffer is not None + + weights.release_disk_streaming_buffer() + assert weights.streaming_block is None + + block1 = weights.load_streaming_block(1) + assert block1 is weights.streaming_block + assert block1 is not block0 + assert torch.all(block1.attn.to_q.weight == 2) + assert block1.attn.to_q.weight_cuda_buffer is not None + + def test_transformer_disk_streaming_rejects_missing_block(tmp_path, h3_modules): _checkpoint_module, weights_module, _infer_module = h3_modules _write_fake_official_checkpoint(tmp_path, weights_module, num_layers=1) From f6d1845516ff49a1b82672074dee4164c80cd5d4 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Wed, 2 Sep 2026 10:19:10 +0800 Subject: [PATCH 11/31] feat(mps): enable MiniMax H3 low-memory t2av --- configs/platforms/mps/minimax_h3_t2av.json | 5 + scripts/platforms/mps/run_minimax_h3_t2av.sh | 3 + .../minimax_h3/test_mps_low_memory_config.py | 97 +++++++++++++++++++ 3 files changed, 105 insertions(+) create mode 100644 tests/models/minimax_h3/test_mps_low_memory_config.py diff --git a/configs/platforms/mps/minimax_h3_t2av.json b/configs/platforms/mps/minimax_h3_t2av.json index feaeb60f9..2fb44c0ad 100644 --- a/configs/platforms/mps/minimax_h3_t2av.json +++ b/configs/platforms/mps/minimax_h3_t2av.json @@ -10,15 +10,19 @@ "cpu_offload": true, "offload_granularity": "block", "dit_prepost_resident": false, + "dit_disk_streaming": true, "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", + "text_encoder_disk_streaming": true, "text_encoder_host_pinned": false, "text_encoder_release_block_offload_buffers": true, + "text_encoder_quantized": false, "vae_cpu_offload": true, "vae_use_compile": false, "vae_attn_type": "torch_sdpa", + "video_vae_quantized": false, "lazy_load": false, "unload_modules": false, @@ -28,6 +32,7 @@ "rms_type": "torch_native", "rope_type": "torch_real_rope", + "tensor_parallel": false, "dit_quantized": false, "dit_quant_scheme": "Default", diff --git a/scripts/platforms/mps/run_minimax_h3_t2av.sh b/scripts/platforms/mps/run_minimax_h3_t2av.sh index 3f8ba842b..a7a976ba9 100755 --- a/scripts/platforms/mps/run_minimax_h3_t2av.sh +++ b/scripts/platforms/mps/run_minimax_h3_t2av.sh @@ -37,6 +37,9 @@ mkdir -p "$(dirname -- "${output_path}")" prompt=${PROMPT:-A cinematic fox walking through a snowy forest} seed=${SEED:-42} +echo "Starting MiniMax-H3 t2av on platform=${PLATFORM}, dtype=${DTYPE}" +echo "Config: dit_disk_streaming=true, text_encoder_disk_streaming=true, VAE lazy lifecycle active" + python -m lightx2v.infer \ --model_cls minimax_h3 \ --task t2av \ diff --git a/tests/models/minimax_h3/test_mps_low_memory_config.py b/tests/models/minimax_h3/test_mps_low_memory_config.py new file mode 100644 index 000000000..1aaff4d72 --- /dev/null +++ b/tests/models/minimax_h3/test_mps_low_memory_config.py @@ -0,0 +1,97 @@ +import importlib.util +import json +import os +import subprocess +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[3] +CONFIG_PATH = REPO_ROOT / "configs/platforms/mps/minimax_h3_t2av.json" +LAUNCHER_PATH = REPO_ROOT / "scripts/platforms/mps/run_minimax_h3_t2av.sh" + + +def _load_config(): + return json.loads(CONFIG_PATH.read_text(encoding="utf-8")) + + +@pytest.fixture() +def runner_module(monkeypatch): + helper_path = REPO_ROOT / "tests/models/minimax_h3/test_runner_mps_low_memory.py" + spec = importlib.util.spec_from_file_location("minimax_h3_runner_low_memory_test_helpers", helper_path) + helper = importlib.util.module_from_spec(spec) + spec.loader.exec_module(helper) + return helper._load_runner_module(monkeypatch) + + +def test_mps_minimax_h3_config_enables_low_memory_streaming(): + config = _load_config() + + assert config["cpu_offload"] is True + assert config["offload_granularity"] == "block" + assert config["dit_disk_streaming"] is True + assert config["text_encoder_cpu_offload"] is True + assert config["text_encoder_offload_granularity"] == "block" + assert config["text_encoder_disk_streaming"] is True + assert config["text_encoder_host_pinned"] is False + assert config["text_encoder_release_block_offload_buffers"] is True + assert config["vae_cpu_offload"] is True + assert config["lazy_load"] is False + assert config["unload_modules"] is False + assert config["warmup"] is False + assert config["attn_type"] == "torch_sdpa" + assert config["rms_type"] == "torch_native" + assert config["rope_type"] == "torch_real_rope" + assert config["vae_attn_type"] == "torch_sdpa" + assert config["dit_quantized"] is False + assert config["dit_quant_scheme"] == "Default" + assert config["text_encoder_quantized"] is False + assert config["video_vae_quantized"] is False + assert config["tensor_parallel"] is False + assert config["use_compile"] is False + assert config["vae_use_compile"] is False + assert "dit_original_ckpt" not in config + + +def test_real_mps_config_triggers_runner_low_memory_load_model(runner_module): + config = _load_config() + config.update({"task": "t2av", "model_path": "/tmp/minimax-h3"}) + runner = object.__new__(runner_module.MiniMaxH3Runner) + runner.config = config + calls = [] + runner.load_transformer = lambda: calls.append("transformer") or object() + runner.load_text_encoder = lambda: calls.append("text_encoder") or [object()] + runner.load_vae = lambda: calls.append("vae") or (object(), object()) + + assert runner._is_mps_low_memory_streaming() is True + runner.load_model() + + assert calls == ["transformer", "text_encoder"] + assert runner.video_vae is None + assert runner.audio_vae is None + + +def test_mps_launcher_has_safe_static_defaults(): + text = LAUNCHER_PATH.read_text(encoding="utf-8") + + assert "export PLATFORM=mps" in text + assert "export DTYPE=BF16" in text + assert "PYTORCH_ENABLE_MPS_FALLBACK" not in text + assert "CUDA" not in text + assert "PYTORCH_CUDA_ALLOC_CONF" not in text + + +def test_mps_launcher_fails_fast_without_model_path(): + env = os.environ.copy() + env.pop("MODEL_PATH", None) + result = subprocess.run( + [str(LAUNCHER_PATH)], + cwd=REPO_ROOT, + env=env, + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 1 + assert "MODEL_PATH must point to the MiniMax-H3 model directory." in result.stdout From 7fe6a7d94213d6edb785f5634520b91896bedf18 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Wed, 2 Sep 2026 11:49:30 +0800 Subject: [PATCH 12/31] fix(minimax_h3): support official video VAE checkpoint layout --- .../video_encoders/hf/minimax_h3/video_vae.py | 102 +++++++++- .../minimax_h3/test_video_vae_loader.py | 186 ++++++++++++++++++ 2 files changed, 284 insertions(+), 4 deletions(-) create mode 100644 tests/models/minimax_h3/test_video_vae_loader.py diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py index 318348b2d..84bdabd72 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py @@ -81,6 +81,102 @@ def _component_dir(model_path: str | Path, component: str) -> Path: raise FileNotFoundError(f"Cannot find MiniMax-H3 {component!r} below {model_path}") +def _resolve_video_vae_dir(model_path: str | Path) -> Path: + """Resolve official ``video_vae`` and legacy ``vae`` component layouts.""" + model_path = Path(model_path) + candidates = [ + model_path / "video_vae", + model_path / "vae", + ] + if model_path.name in {"video_vae", "vae"}: + candidates.append(model_path) + + tried = [] + for candidate in candidates: + if candidate in tried: + continue + tried.append(candidate) + if candidate.is_dir(): + return candidate + + formatted = ", ".join(str(path) for path in tried) + raise FileNotFoundError(f"Cannot find MiniMax-H3 video VAE directory. Tried: {formatted}") + + +def _read_json_file(path: Path) -> dict: + with path.open("r", encoding="utf-8") as handle: + return json.load(handle) + + +def _normalize_official_video_vae_config(wrapper_config: dict, source_config: dict | None = None) -> dict: + """Convert the official FL2VA wrapper/source config pair to native keys.""" + if source_config is None: + return dict(wrapper_config) + + config = dict(wrapper_config) + if "in_channels" in source_config: + config["in_channels"] = source_config["in_channels"] + if "out_ch" in source_config: + config["out_channels"] = source_config["out_ch"] + if "z_channels" in source_config and "latent_channels" not in config: + config["latent_channels"] = source_config["z_channels"] + if "ch" in source_config and "ch_mult" in source_config: + config["block_out_channels"] = [int(source_config["ch"]) * int(value) for value in source_config["ch_mult"]] + if "num_res_blocks" in source_config: + config["layers_per_block"] = source_config["num_res_blocks"] + if "space_down" in source_config: + config["spatial_downsample_factors"] = source_config["space_down"] + if "time_down" in source_config: + config["temporal_downsample_factors"] = source_config["time_down"] + if "padding_mode" in source_config: + config["spatial_padding_mode"] = source_config["padding_mode"] + + vit_config = source_config.get("vit_decoder_kwargs") + if isinstance(vit_config, dict): + vit_mappings = { + "num_layers": "decoder_num_layers", + "heads": "decoder_num_attention_heads", + "dim_head": "decoder_attention_head_dim", + "rope_theta": "decoder_rope_theta", + "rope_dim_ratio": "decoder_rope_dim_ratio", + } + for source_key, target_key in vit_mappings.items(): + if source_key in vit_config: + config[target_key] = vit_config[source_key] + + wrapper_mappings = { + "vae_clip_length": "clip_length", + "vae_token_drop": "token_drop", + } + for source_key, target_key in wrapper_mappings.items(): + if source_key in wrapper_config: + config[target_key] = wrapper_config[source_key] + for key in ("latent_channels", "latents_mean", "latents_std"): + if key in wrapper_config: + config[key] = wrapper_config[key] + return config + + +def _load_video_vae_config_and_weight_path( + vae_dir: Path, + checkpoint_path: str | Path | None, +) -> tuple[dict, Path | str]: + wrapper_config = _read_json_file(vae_dir / "config.json") + source_config = None + default_weight_path: Path | str = vae_dir + + source_path = wrapper_config.get("source_path") + source_safetensors_path = wrapper_config.get("source_safetensors_path") + source_dir = vae_dir / source_path if isinstance(source_path, str) else None + if source_dir is not None and (source_dir / "config.json").is_file(): + source_config = _read_json_file(source_dir / "config.json") + if source_dir is not None and isinstance(source_safetensors_path, str): + default_weight_path = source_dir / source_safetensors_path + + weight_path = checkpoint_path if checkpoint_path is not None else default_weight_path + return _normalize_official_video_vae_config(wrapper_config, source_config), weight_path + + class _SwiGLU(nn.Module): """Checkpoint-compatible SwiGLU used by the ViT decoder.""" @@ -699,12 +795,10 @@ def from_pretrained( use_compile: bool = False, attn_type: str = "torch_sdpa", ) -> "MiniMaxH3VideoVAE": - vae_dir = _component_dir(model_path, "vae") + vae_dir = _resolve_video_vae_dir(model_path) if (checkpoint_path is None) != (quant_scheme is None): raise ValueError("MiniMax-H3 video VAE checkpoint_path and quant_scheme must be configured together") - weight_path = checkpoint_path if checkpoint_path is not None else vae_dir - with (vae_dir / "config.json").open("r", encoding="utf-8") as handle: - config = json.load(handle) + config, weight_path = _load_video_vae_config_and_weight_path(vae_dir, checkpoint_path) # The released decoder is several GiB. Constructing it on meta avoids # allocating and then immediately overwriting random initialized weights. diff --git a/tests/models/minimax_h3/test_video_vae_loader.py b/tests/models/minimax_h3/test_video_vae_loader.py new file mode 100644 index 000000000..701cc72cc --- /dev/null +++ b/tests/models/minimax_h3/test_video_vae_loader.py @@ -0,0 +1,186 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).parents[3] + + +def _install_module(monkeypatch, name, **attrs): + module = types.ModuleType(name) + for attr_name, value in attrs.items(): + setattr(module, attr_name, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +@pytest.fixture() +def video_vae_module(monkeypatch): + for package_name in [ + "lightx2v", + "lightx2v.models", + "lightx2v.models.video_encoders", + "lightx2v.models.video_encoders.hf", + "lightx2v.models.video_encoders.hf.minimax_h3", + "lightx2v.utils", + "lightx2v_platform", + "lightx2v_platform.base", + ]: + package = types.ModuleType(package_name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, package_name, package) + + _install_module( + monkeypatch, + "lightx2v.models.video_encoders.hf.minimax_h3.weights", + SafetensorsSubsetReport=object, + load_safetensors_subset=lambda *_args, **_kwargs: None, + ) + _install_module(monkeypatch, "lightx2v.utils.registry_factory", ATTN_WEIGHT_REGISTER={}) + _install_module(monkeypatch, "lightx2v_platform.base.global_var", AI_DEVICE="cpu") + + module_path = REPO_ROOT / "lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py" + spec = importlib.util.spec_from_file_location("minimax_h3_video_vae_under_test", module_path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _write_json(path, data): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(__import__("json").dumps(data), encoding="utf-8") + + +def _wrapper_config(): + return { + "source_path": "source", + "source_safetensors_path": "model.safetensors", + "latent_channels": 24, + "latents_mean": [0.1] * 24, + "latents_std": [1.1] * 24, + "vae_clip_length": 17, + "vae_token_drop": 3, + } + + +def _source_config(): + return { + "in_channels": 3, + "out_ch": 3, + "z_channels": 24, + "ch": 8, + "ch_mult": [1, 2, 4], + "num_res_blocks": 2, + "space_down": [2, 2, 1], + "time_down": [1, 2, 1], + "padding_mode": "reflect", + "vit_decoder_kwargs": { + "num_layers": 12, + "heads": 4, + "dim_head": 16, + "rope_theta": 100.0, + "rope_dim_ratio": 0.75, + }, + } + + +def test_resolves_official_video_vae_layout_before_legacy(video_vae_module, tmp_path): + official = tmp_path / "video_vae" + legacy = tmp_path / "vae" + official.mkdir() + legacy.mkdir() + + assert video_vae_module._resolve_video_vae_dir(tmp_path) == official + + +def test_resolves_legacy_vae_layout(video_vae_module, tmp_path): + legacy = tmp_path / "vae" + legacy.mkdir() + + assert video_vae_module._resolve_video_vae_dir(tmp_path) == legacy + + +@pytest.mark.parametrize("name", ["video_vae", "vae"]) +def test_resolves_direct_component_path(video_vae_module, tmp_path, name): + component = tmp_path / name + component.mkdir() + + assert video_vae_module._resolve_video_vae_dir(component) == component + + +def test_missing_video_vae_and_vae_reports_tried_paths(video_vae_module, tmp_path): + with pytest.raises(FileNotFoundError, match="video_vae.*vae"): + video_vae_module._resolve_video_vae_dir(tmp_path) + + +def test_official_layout_uses_source_safetensors_by_default(video_vae_module, tmp_path): + vae_dir = tmp_path / "video_vae" + _write_json(vae_dir / "config.json", _wrapper_config()) + _write_json(vae_dir / "source/config.json", _source_config()) + (vae_dir / "source/model.safetensors").write_bytes(b"") + + config, weight_path = video_vae_module._load_video_vae_config_and_weight_path(vae_dir, None) + + assert weight_path == vae_dir / "source/model.safetensors" + assert config["block_out_channels"] == [8, 16, 32] + + +def test_explicit_checkpoint_path_wins_for_official_layout(video_vae_module, tmp_path): + vae_dir = tmp_path / "video_vae" + explicit = tmp_path / "quantized.safetensors" + _write_json(vae_dir / "config.json", _wrapper_config()) + _write_json(vae_dir / "source/config.json", _source_config()) + (vae_dir / "source/model.safetensors").write_bytes(b"") + + _config, weight_path = video_vae_module._load_video_vae_config_and_weight_path(vae_dir, explicit) + + assert weight_path == explicit + + +def test_official_config_normalization_maps_wrapper_and_source_fields(video_vae_module): + config = video_vae_module._normalize_official_video_vae_config(_wrapper_config(), _source_config()) + + assert config["in_channels"] == 3 + assert config["out_channels"] == 3 + assert config["latent_channels"] == 24 + assert config["block_out_channels"] == [8, 16, 32] + assert config["layers_per_block"] == 2 + assert config["spatial_downsample_factors"] == [2, 2, 1] + assert config["temporal_downsample_factors"] == [1, 2, 1] + assert config["spatial_padding_mode"] == "reflect" + assert config["decoder_num_layers"] == 12 + assert config["decoder_num_attention_heads"] == 4 + assert config["decoder_attention_head_dim"] == 16 + assert config["decoder_rope_theta"] == 100.0 + assert config["decoder_rope_dim_ratio"] == 0.75 + assert config["clip_length"] == 17 + assert config["token_drop"] == 3 + assert config["latents_mean"] == [0.1] * 24 + assert config["latents_std"] == [1.1] * 24 + + +def test_legacy_config_normalization_keeps_existing_fields(video_vae_module): + legacy = { + "latent_channels": 8, + "block_out_channels": [4, 8], + "clip_length": 9, + "token_drop": 1, + "custom": "kept", + } + + assert video_vae_module._normalize_official_video_vae_config(legacy, None) == legacy + + +def test_legacy_layout_uses_component_dir_as_weight_path(video_vae_module, tmp_path): + vae_dir = tmp_path / "vae" + legacy = {"latent_channels": 8, "block_out_channels": [4, 8]} + _write_json(vae_dir / "config.json", legacy) + (vae_dir / "model.safetensors").write_bytes(b"") + + config, weight_path = video_vae_module._load_video_vae_config_and_weight_path(vae_dir, None) + + assert config == legacy + assert weight_path == vae_dir From a02eae8d5806889091fc90d40090788c169deffd Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Fri, 4 Sep 2026 19:04:09 +0800 Subject: [PATCH 13/31] fix(minimax_h3): load official video VAE checkpoints --- .../video_encoders/hf/minimax_h3/video_vae.py | 7 +- .../video_encoders/hf/minimax_h3/weights.py | 164 ++++++++++++++++++ .../test_video_vae_checkpoint_adapter.py | 147 ++++++++++++++++ .../minimax_h3/test_video_vae_loader.py | 58 +++++++ 4 files changed, 375 insertions(+), 1 deletion(-) create mode 100644 tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py index 84bdabd72..1c68a0ab5 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py @@ -45,6 +45,8 @@ from lightx2v.models.video_encoders.hf.minimax_h3.weights import ( SafetensorsSubsetReport, + _is_official_video_vae_checkpoint, + load_minimax_h3_video_vae_checkpoint, load_safetensors_subset, ) from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER @@ -813,7 +815,10 @@ def from_pretrained( attn_type=attn_type, ) model._reset_runtime_buffers() - model.load_report = load_safetensors_subset(model, weight_path) + if quant_scheme is None and _is_official_video_vae_checkpoint(weight_path): + model.load_report = load_minimax_h3_video_vae_checkpoint(model, weight_path) + else: + model.load_report = load_safetensors_subset(model, weight_path) if quant_scheme is not None: # Pack only after loading the checkpoint's original Q/K/V keys. model._pack_decoder_fp8_qkv() diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/weights.py b/lightx2v/models/video_encoders/hf/minimax_h3/weights.py index ceb1452e1..8c24596c4 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/weights.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/weights.py @@ -23,6 +23,7 @@ from __future__ import annotations +import re from dataclasses import dataclass from pathlib import Path @@ -159,3 +160,166 @@ def load_safetensors_subset(module: nn.Module, component_dir: str | Path) -> Saf if missing: raise RuntimeError(f"Failed to load MiniMax-H3 tensors: {missing[:20]}") return SafetensorsSubsetReport(report.component_dir, report.files, tuple(sorted(loaded)), report.ignored_keys) + + +_OFFICIAL_VIDEO_VAE_SIGNATURE = { + "decoder.x_embedder.weight", + "decoder.transformer_blocks.0.attn.to_qkv.weight", + "decoder.transformer_blocks.0.ff.w1.weight", +} + + +def _is_official_video_vae_checkpoint(component_dir: str | Path) -> bool: + """Detect the released Video VAE schema from safetensors keys.""" + keys: set[str] = set() + for filename in _component_files(component_dir): + with safe_open(filename, framework="pt", device="cpu") as checkpoint: + keys.update(checkpoint.keys()) + return _OFFICIAL_VIDEO_VAE_SIGNATURE <= keys and "decoder.proj_in.weight" not in keys + + +def _official_video_vae_targets(source_key: str) -> tuple[str, ...] | None: + if source_key == "decoder.mask_token": + return () + + qkv = re.fullmatch(r"(decoder\.transformer_blocks\.\d+\.attn)\.to_qkv\.(weight|bias)", source_key) + if qkv: + prefix, suffix = qkv.groups() + return tuple(f"{prefix}.to_{name}.{suffix}" for name in ("q", "k", "v")) + + w1 = re.fullmatch(r"(decoder\.transformer_blocks\.\d+\.ff)\.w1\.(weight|bias)", source_key) + if w1: + prefix, suffix = w1.groups() + return (f"{prefix}.net.0.proj.{suffix}",) + + target = source_key + down_block = re.fullmatch(r"encoder\.down\.(\d+)\.block\.(\d+)\.(.+)", target) + if down_block: + stage, block, suffix = down_block.groups() + suffix = suffix.replace("nin_shortcut", "conv_shortcut") + target = f"encoder.down_blocks.{stage}.resnets.{block}.{suffix}" + else: + downsample = re.fullmatch(r"encoder\.down\.(\d+)\.downsample\.conv\.(weight|bias)", target) + if downsample: + stage, suffix = downsample.groups() + target = f"encoder.down_blocks.{stage}.downsamplers.0.conv.{suffix}" + + target = target.replace("decoder.x_embedder.", "decoder.proj_in.") + target = re.sub(r"(decoder\.transformer_blocks\.\d+\.attn)\.to_out\.", r"\1.to_out.0.", target) + target = re.sub(r"(decoder\.transformer_blocks\.\d+\.ff)\.w2\.", r"\1.net.2.", target) + return (target,) + + +def _official_validation_error( + *, unknown: list[str], missing: list[str], duplicates: list[str], shape_mismatches: list[str], dtype_mismatches: list[str] +) -> RuntimeError: + details = [] + for label, values in ( + ("unknown", unknown), + ("missing", missing), + ("duplicate", duplicates), + ("shape_mismatch", shape_mismatches), + ("dtype_mismatch", dtype_mismatches), + ): + if values: + details.append(f"{label}={values[:20]}{' ...' if len(values) > 20 else ''}") + return RuntimeError("Official MiniMax-H3 Video VAE checkpoint validation failed: " + ", ".join(details)) + + +def validate_minimax_h3_video_vae_checkpoint(module: nn.Module, component_dir: str | Path) -> SafetensorsSubsetReport: + """Validate official-to-native coverage using safetensors metadata only.""" + files = _component_files(component_dir) + expected = _expected_specs(module) + assigned: set[str] = set() + unknown: list[str] = [] + duplicates: list[str] = [] + shape_mismatches: list[str] = [] + dtype_mismatches: list[str] = [] + ignored = 0 + + for filename in files: + with safe_open(filename, framework="pt", device="cpu") as checkpoint: + for source_key in checkpoint.keys(): + tensor_slice = checkpoint.get_slice(source_key) + source_shape = tuple(tensor_slice.get_shape()) + source_dtype = str(tensor_slice.get_dtype()) + targets = _official_video_vae_targets(source_key) + if targets == (): + decoder_dim = expected.get("decoder.proj_in.weight", ((0,), torch.float32))[0][0] + if source_shape != (1, 1, decoder_dim): + shape_mismatches.append(f"{source_key}:{source_shape}!={(1, 1, decoder_dim)}") + elif source_dtype != _SAFETENSORS_DTYPES.get(expected["decoder.proj_in.weight"][1]): + dtype_mismatches.append(source_key) + ignored += 1 + continue + if targets is None or any(target not in expected for target in targets): + unknown.append(source_key) + continue + + if len(targets) == 3: + if not source_shape or source_shape[0] % 3: + shape_mismatches.append(source_key) + continue + mapped_shape = (source_shape[0] // 3, *source_shape[1:]) + else: + mapped_shape = source_shape + for target in targets: + expected_shape, expected_dtype = expected[target] + if mapped_shape != expected_shape: + shape_mismatches.append(f"{source_key}->{target}:{mapped_shape}!={expected_shape}") + if source_dtype != _SAFETENSORS_DTYPES.get(expected_dtype): + dtype_mismatches.append(f"{source_key}->{target}") + if target in assigned: + duplicates.append(target) + assigned.add(target) + + missing = sorted(set(expected) - assigned) + if unknown or missing or duplicates or shape_mismatches or dtype_mismatches or ignored != 1: + if ignored != 1: + unknown.append(f"ignored_count:{ignored} (expected decoder.mask_token exactly once)") + raise _official_validation_error( + unknown=unknown, + missing=missing, + duplicates=duplicates, + shape_mismatches=shape_mismatches, + dtype_mismatches=dtype_mismatches, + ) + return SafetensorsSubsetReport(Path(component_dir), files, tuple(sorted(assigned)), ignored) + + +def load_minimax_h3_video_vae_checkpoint(module: nn.Module, component_dir: str | Path) -> SafetensorsSubsetReport: + """Stream the released official Video VAE schema into the native module.""" + report = validate_minimax_h3_video_vae_checkpoint(module, component_dir) + loaded: set[str] = set() + for filename in report.files: + with safe_open(filename, framework="pt", device="cpu") as checkpoint: + for source_key in checkpoint.keys(): + targets = _official_video_vae_targets(source_key) + if targets == (): + continue + tensor_slice = checkpoint.get_slice(source_key) + if len(targets) == 3: + chunk_size = tensor_slice.get_shape()[0] // 3 + for index, target in enumerate(targets): + _assign_tensor(module, target, tensor_slice[index * chunk_size : (index + 1) * chunk_size]) + loaded.add(target) + elif ".ff.w1." in source_key: + target = targets[0] + rows = tensor_slice.get_shape()[0] + half = rows // 2 + gate = tensor_slice[:half] + reordered = torch.empty(tuple(tensor_slice.get_shape()), dtype=gate.dtype) + reordered[half:].copy_(gate) + del gate + value = tensor_slice[half:] + reordered[:half].copy_(value) + del value + _assign_tensor(module, target, reordered) + loaded.add(target) + else: + target = targets[0] + _assign_tensor(module, target, checkpoint.get_tensor(source_key)) + loaded.add(target) + if loaded != set(report.loaded_keys): + raise RuntimeError("Official MiniMax-H3 Video VAE load did not reproduce validated target coverage") + return SafetensorsSubsetReport(report.component_dir, report.files, tuple(sorted(loaded)), report.ignored_keys) diff --git a/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py b/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py new file mode 100644 index 000000000..ffd3f6046 --- /dev/null +++ b/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py @@ -0,0 +1,147 @@ +import importlib.util +import sys +from pathlib import Path + +import pytest +import torch +import torch.nn as nn +from safetensors.torch import save_file + +REPO_ROOT = Path(__file__).parents[3] +_SPEC = importlib.util.spec_from_file_location("minimax_h3_weights_under_test", REPO_ROOT / "lightx2v/models/video_encoders/hf/minimax_h3/weights.py") +_WEIGHTS = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _WEIGHTS +_SPEC.loader.exec_module(_WEIGHTS) +_is_official_video_vae_checkpoint = _WEIGHTS._is_official_video_vae_checkpoint +load_minimax_h3_video_vae_checkpoint = _WEIGHTS.load_minimax_h3_video_vae_checkpoint +load_safetensors_subset = _WEIGHTS.load_safetensors_subset +validate_minimax_h3_video_vae_checkpoint = _WEIGHTS.validate_minimax_h3_video_vae_checkpoint + + +def _parameter_module(specs): + class _IndexedModule(nn.Module): + def __getitem__(self, index): + return self._modules[str(index)] + + root = _IndexedModule() + for key, tensor in specs.items(): + parent = root + parts = key.split(".") + for part in parts[:-1]: + if not hasattr(parent, part): + parent.add_module(part, _IndexedModule()) + parent = getattr(parent, part) + parent.register_parameter(parts[-1], nn.Parameter(torch.empty_like(tensor, device="meta"))) + return root + + +def _official_tensors(): + qkv_weight = torch.arange(48, dtype=torch.float32).reshape(12, 4) + qkv_bias = torch.arange(12, dtype=torch.float32) + w1_weight = torch.cat((torch.ones(4, 4), torch.full((4, 4), 2.0))) + w1_bias = torch.cat((torch.ones(4), torch.full((4,), 2.0))) + return { + "encoder.conv_in.weight": torch.full((2, 3, 1, 1, 1), 3.0), + "encoder.down.0.block.0.nin_shortcut.weight": torch.full((2, 2, 1, 1, 1), 4.0), + "encoder.down.0.downsample.conv.bias": torch.full((2,), 5.0), + "decoder.x_embedder.weight": torch.full((4, 2), 6.0), + "decoder.transformer_blocks.0.attn.to_out.weight": torch.full((4, 4), 7.0), + "decoder.transformer_blocks.0.attn.to_qkv.weight": qkv_weight, + "decoder.transformer_blocks.0.attn.to_qkv.bias": qkv_bias, + "decoder.transformer_blocks.0.ff.w1.weight": w1_weight, + "decoder.transformer_blocks.0.ff.w1.bias": w1_bias, + "decoder.transformer_blocks.0.ff.w2.weight": torch.full((4, 4), 8.0), + "decoder.mask_token": torch.zeros(1, 1, 4), + } + + +def _native_specs(): + tensors = _official_tensors() + return { + "encoder.conv_in.weight": tensors["encoder.conv_in.weight"], + "encoder.down_blocks.0.resnets.0.conv_shortcut.weight": tensors["encoder.down.0.block.0.nin_shortcut.weight"], + "encoder.down_blocks.0.downsamplers.0.conv.bias": tensors["encoder.down.0.downsample.conv.bias"], + "decoder.proj_in.weight": tensors["decoder.x_embedder.weight"], + "decoder.transformer_blocks.0.attn.to_out.0.weight": tensors["decoder.transformer_blocks.0.attn.to_out.weight"], + "decoder.transformer_blocks.0.attn.to_q.weight": torch.empty(4, 4), + "decoder.transformer_blocks.0.attn.to_k.weight": torch.empty(4, 4), + "decoder.transformer_blocks.0.attn.to_v.weight": torch.empty(4, 4), + "decoder.transformer_blocks.0.attn.to_q.bias": torch.empty(4), + "decoder.transformer_blocks.0.attn.to_k.bias": torch.empty(4), + "decoder.transformer_blocks.0.attn.to_v.bias": torch.empty(4), + "decoder.transformer_blocks.0.ff.net.0.proj.weight": tensors["decoder.transformer_blocks.0.ff.w1.weight"], + "decoder.transformer_blocks.0.ff.net.0.proj.bias": tensors["decoder.transformer_blocks.0.ff.w1.bias"], + "decoder.transformer_blocks.0.ff.net.2.weight": tensors["decoder.transformer_blocks.0.ff.w2.weight"], + } + + +def _write(path: Path, tensors=None): + save_file(tensors or _official_tensors(), path) + return path + + +def test_official_detection_uses_strict_key_signature(tmp_path): + official = _write(tmp_path / "official.safetensors") + legacy = tmp_path / "legacy.safetensors" + save_file({"decoder.proj_in.weight": torch.zeros(4, 2)}, legacy) + assert _is_official_video_vae_checkpoint(official) + assert not _is_official_video_vae_checkpoint(legacy) + + +def test_official_mapping_qkv_ffn_and_mask_token(tmp_path): + tensors = _official_tensors() + module = _parameter_module(_native_specs()) + report = load_minimax_h3_video_vae_checkpoint(module, _write(tmp_path / "model.safetensors", tensors)) + state = module.state_dict() + + assert len(report.loaded_keys) == len(_native_specs()) + assert report.ignored_keys == 1 + assert torch.equal(state["encoder.down_blocks.0.resnets.0.conv_shortcut.weight"], tensors["encoder.down.0.block.0.nin_shortcut.weight"]) + assert torch.equal(state["decoder.proj_in.weight"], tensors["decoder.x_embedder.weight"]) + assert torch.equal(state["decoder.transformer_blocks.0.attn.to_out.0.weight"], tensors["decoder.transformer_blocks.0.attn.to_out.weight"]) + assert torch.equal(state["decoder.transformer_blocks.0.ff.net.2.weight"], tensors["decoder.transformer_blocks.0.ff.w2.weight"]) + for index, name in enumerate(("q", "k", "v")): + assert torch.equal(state[f"decoder.transformer_blocks.0.attn.to_{name}.weight"], tensors["decoder.transformer_blocks.0.attn.to_qkv.weight"][index * 4 : (index + 1) * 4]) + assert torch.equal(state[f"decoder.transformer_blocks.0.attn.to_{name}.bias"], tensors["decoder.transformer_blocks.0.attn.to_qkv.bias"][index * 4 : (index + 1) * 4]) + assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.weight"][:4] == 2) + assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.weight"][4:] == 1) + assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.bias"][:4] == 2) + assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.bias"][4:] == 1) + + +def test_mask_token_shape_is_validated(tmp_path): + tensors = _official_tensors() + tensors["decoder.mask_token"] = torch.zeros(1, 2, 4) + with pytest.raises(RuntimeError, match="shape_mismatch.*decoder.mask_token"): + validate_minimax_h3_video_vae_checkpoint(_parameter_module(_native_specs()), _write(tmp_path / "bad.safetensors", tensors)) + + +def test_unexpected_official_key_fails(tmp_path): + tensors = _official_tensors() + tensors["decoder.surprise"] = torch.zeros(1) + with pytest.raises(RuntimeError, match="unknown.*decoder.surprise"): + validate_minimax_h3_video_vae_checkpoint(_parameter_module(_native_specs()), _write(tmp_path / "bad.safetensors", tensors)) + + +def test_missing_native_target_fails(tmp_path): + specs = _native_specs() + specs["decoder.proj_out.weight"] = torch.empty(4, 4) + with pytest.raises(RuntimeError, match="missing.*decoder.proj_out.weight"): + validate_minimax_h3_video_vae_checkpoint(_parameter_module(specs), _write(tmp_path / "bad.safetensors")) + + +def test_duplicate_target_fails(tmp_path): + tensors = _official_tensors() + tensors["decoder.proj_in.weight"] = tensors["decoder.x_embedder.weight"] + with pytest.raises(RuntimeError, match="duplicate.*decoder.proj_in.weight"): + validate_minimax_h3_video_vae_checkpoint(_parameter_module(_native_specs()), _write(tmp_path / "bad.safetensors", tensors)) + + +def test_legacy_subset_loader_regression(tmp_path): + expected = {"layer.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3)} + module = _parameter_module(expected) + path = tmp_path / "native.safetensors" + save_file(expected, path) + report = load_safetensors_subset(module, path) + assert report.loaded_keys == ("layer.weight",) + assert torch.equal(module.state_dict()["layer.weight"], expected["layer.weight"]) diff --git a/tests/models/minimax_h3/test_video_vae_loader.py b/tests/models/minimax_h3/test_video_vae_loader.py index 701cc72cc..c4fca54c5 100644 --- a/tests/models/minimax_h3/test_video_vae_loader.py +++ b/tests/models/minimax_h3/test_video_vae_loader.py @@ -36,6 +36,8 @@ def video_vae_module(monkeypatch): monkeypatch, "lightx2v.models.video_encoders.hf.minimax_h3.weights", SafetensorsSubsetReport=object, + _is_official_video_vae_checkpoint=lambda *_args, **_kwargs: False, + load_minimax_h3_video_vae_checkpoint=lambda *_args, **_kwargs: None, load_safetensors_subset=lambda *_args, **_kwargs: None, ) _install_module(monkeypatch, "lightx2v.utils.registry_factory", ATTN_WEIGHT_REGISTER={}) @@ -184,3 +186,59 @@ def test_legacy_layout_uses_component_dir_as_weight_path(video_vae_module, tmp_p assert config == legacy assert weight_path == vae_dir + + +def test_from_pretrained_dispatches_official_and_legacy(video_vae_module, monkeypatch, tmp_path): + calls = [] + + class TinyVAE(video_vae_module.MiniMaxH3VideoVAE): + def __init__(self, _config, **_kwargs): + video_vae_module.nn.Module.__init__(self) + self.weight = video_vae_module.nn.Parameter(video_vae_module.torch.empty(1)) + self.execution_device = video_vae_module.torch.device("cpu") + + def _reset_runtime_buffers(self): + pass + + def _prepare_inference_dtypes(self): + pass + + monkeypatch.setattr(video_vae_module, "_resolve_video_vae_dir", lambda _path: tmp_path) + monkeypatch.setattr(video_vae_module, "_load_video_vae_config_and_weight_path", lambda *_args: ({}, tmp_path / "model.safetensors")) + monkeypatch.setattr(video_vae_module, "load_minimax_h3_video_vae_checkpoint", lambda *_args: calls.append("official") or "official-report") + monkeypatch.setattr(video_vae_module, "load_safetensors_subset", lambda *_args: calls.append("legacy") or "legacy-report") + + monkeypatch.setattr(video_vae_module, "_is_official_video_vae_checkpoint", lambda _path: True) + assert TinyVAE.from_pretrained(tmp_path, cpu_offload=True).load_report == "official-report" + monkeypatch.setattr(video_vae_module, "_is_official_video_vae_checkpoint", lambda _path: False) + assert TinyVAE.from_pretrained(tmp_path, cpu_offload=True).load_report == "legacy-report" + assert calls == ["official", "legacy"] + + +def test_from_pretrained_quantized_path_bypasses_official_adapter(video_vae_module, monkeypatch, tmp_path): + calls = [] + + class TinyQuantizedVAE(video_vae_module.MiniMaxH3VideoVAE): + def __init__(self, _config, **_kwargs): + video_vae_module.nn.Module.__init__(self) + self.weight = video_vae_module.nn.Parameter(video_vae_module.torch.empty(1)) + self.execution_device = video_vae_module.torch.device("cpu") + + def _reset_runtime_buffers(self): + pass + + def _pack_decoder_fp8_qkv(self): + calls.append("pack") + + def _prepare_inference_dtypes(self): + pass + + monkeypatch.setattr(video_vae_module, "_resolve_video_vae_dir", lambda _path: tmp_path) + monkeypatch.setattr(video_vae_module, "_load_video_vae_config_and_weight_path", lambda *_args: ({}, tmp_path / "quant.safetensors")) + monkeypatch.setattr(video_vae_module, "_is_official_video_vae_checkpoint", lambda _path: pytest.fail("official detection must not run for quantized checkpoints")) + monkeypatch.setattr(video_vae_module, "load_minimax_h3_video_vae_checkpoint", lambda *_args: pytest.fail("official adapter must not run for quantized checkpoints")) + monkeypatch.setattr(video_vae_module, "load_safetensors_subset", lambda *_args: calls.append("legacy") or "quant-report") + + model = TinyQuantizedVAE.from_pretrained(tmp_path, checkpoint_path=tmp_path / "quant.safetensors", quant_scheme="fp8-sgl", cpu_offload=True) + assert model.load_report == "quant-report" + assert calls == ["legacy", "pack"] From 8bbf600e4a7094a4069d7a64714623b8a0beb696 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Sat, 5 Sep 2026 09:22:28 +0800 Subject: [PATCH 14/31] fix(minimax_h3): fix Qwen disk streaming dispatch --- .../input_encoders/hf/minimax_h3/qwen3vl.py | 2 +- .../minimax_h3/test_qwen3vl_disk_streaming.py | 88 +++++++++++++++++++ 2 files changed, 89 insertions(+), 1 deletion(-) diff --git a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py index 387839c61..bc5006258 100644 --- a/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py +++ b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py @@ -1499,7 +1499,7 @@ def infer(self, prompt, image_list=None, references=None): video_grid_thw, ) vision_mask, vision_embeds, deepstack = self._encode_vision(input_ids, pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw) - if self.cpu_offload and not self.block_offload: + if self.cpu_offload and not self.block_offload and not self.disk_streaming: self.text_encoder.to_cuda() elif self.block_offload: # Recreate transient device slots if the previous request was diff --git a/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py b/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py index 7a9f3301a..760924246 100644 --- a/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py +++ b/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py @@ -368,3 +368,91 @@ def test_qwen3vl_release_disk_streaming_buffer_clears_device_refs(tmp_path, qwen assert old_layer.self_attn.q_proj.weight_cuda_buffer is None assert old_layer.input_layernorm.weight is None assert old_layer.input_layernorm.weight_cuda_buffer is None + + +@pytest.mark.parametrize("mode", ["disk", "resident", "block"]) +@pytest.mark.parametrize("release_buffers", [True, False]) +@pytest.mark.parametrize("fail_forward", [False, True]) +def test_public_infer_offload_lifecycle(tmp_path, monkeypatch, qwen_module, mode, release_buffers, fail_forward): + from unittest.mock import Mock + + # Exercise the real constructor's mode flags without needing an accelerator. + monkeypatch.setattr(qwen_module, "AI_DEVICE", "mps") + encoder = qwen_module.MiniMaxH3Qwen3VLTextEncoder( + { + "task": "t2av", + "text_encoder_cpu_offload": True, + "text_encoder_offload_granularity": "model" if mode == "resident" else "block", + "text_encoder_disk_streaming": mode == "disk", + "text_encoder_release_block_offload_buffers": release_buffers, + "text_encoder_load_on_init": False, + } + ) + monkeypatch.setattr(qwen_module, "AI_DEVICE", "cpu") + monkeypatch.setattr(qwen_module, "MINIMAX_H3_TEXT_HIDDEN_SIZE", 8) + backbone = qwen_module._Qwen3VLTextBackboneWeights( + encoder.config, _tiny_text_config(), num_layers=2, + block_offload=encoder.block_offload, disk_streaming=encoder.disk_streaming, + ) + encoder.text_encoder = backbone + encoder.tokenizer = Mock(return_value={"input_ids": [0, 1, 2]}) + assert (encoder.cpu_offload, encoder.block_offload, encoder.disk_streaming) == (True, mode == "block", mode == "disk") + events = [] + + def layer_forward(self, hidden_states, position_embeddings): + events.append("layer") + if fail_forward and events.count("layer") == 2: + raise RuntimeError("injected forward failure") + return hidden_states + 1 + + monkeypatch.setattr(qwen_module._Qwen3VLDecoderLayerWeights, "forward", layer_forward) + if mode == "disk": + _, _, _, weight_map = _write_fake_checkpoint(tmp_path, backbone, qwen_module) + backbone.init_disk_streaming(tmp_path, weight_map) + else: + backbone.embed_tokens.weight = torch.full((32, 8), 5, dtype=torch.bfloat16) + + def block_forward(input_ids, *args): + events.append("block") + if fail_forward: + raise RuntimeError("injected forward failure") + return torch.full((input_ids.numel(), 8), 7, dtype=torch.bfloat16) + + monkeypatch.setattr(backbone, "_forward_with_block_offload", Mock(side_effect=block_forward)) + for name in ("to_cuda", "to_cpu", "init_block_offload", "release_block_offload_buffers"): + monkeypatch.setattr(backbone, name, Mock()) + for name in ("_forward_streaming_embedding", "load_streaming_layer", "release_disk_streaming_buffer"): + monkeypatch.setattr(backbone, name, Mock(wraps=getattr(backbone, name))) + + prompt = "A cat walking on the grass." + if fail_forward: + with pytest.raises(RuntimeError, match="injected forward failure"): + encoder.infer(prompt) + else: + result = encoder.infer(prompt) + assert set(result) == {"prompt_embeds", "text_token_tags"} + embeds, tags = result["prompt_embeds"], result["text_token_tags"] + assert isinstance(embeds, torch.Tensor) and isinstance(tags, torch.Tensor) + assert embeds.shape == (3, 8) and embeds.dtype == torch.bfloat16 + assert embeds.device.type == "cpu" and embeds.is_contiguous() + assert torch.isfinite(embeds).all() and torch.all(embeds == 7) + assert tags.shape == (3,) and tags.dtype == torch.long + assert tags.device == embeds.device + assert torch.all(tags == qwen_module.MINIMAX_H3_TEXT_TAG) + + encoder.tokenizer.assert_called_once_with(prompt, add_special_tokens=False) + assert backbone.to_cuda.call_count == int(mode == "resident") + assert backbone.to_cpu.call_count == int(mode == "resident") + assert backbone.init_block_offload.call_count == int(mode == "block") + assert backbone.release_block_offload_buffers.call_count == int(mode == "block" and release_buffers) + assert backbone.release_disk_streaming_buffer.call_count == int(mode == "disk" and release_buffers) + assert backbone._forward_with_block_offload.call_count == int(mode == "block") + if mode == "disk": + backbone._forward_streaming_embedding.assert_called_once() + assert [call.args[0] for call in backbone.load_streaming_layer.call_args_list] == [0, 1] + assert (backbone.streaming_layer is None) == release_buffers + assert backbone.embed_tokens.weight is None + assert not hasattr(backbone.embed_tokens, "pin_weight") + else: + backbone._forward_streaming_embedding.assert_not_called() + backbone.load_streaming_layer.assert_not_called() From fd0a91ff40c7799c79c6278d4affbcc984735ca7 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Sat, 5 Sep 2026 11:14:13 +0800 Subject: [PATCH 15/31] feat(minimax_h3): load official DiT checkpoints --- .../models/networks/minimax_h3/checkpoint.py | 317 ++++++++++++++++- lightx2v/models/networks/minimax_h3/model.py | 5 +- .../minimax_h3/weights/transformer_weights.py | 14 +- .../minimax_h3/test_checkpoint_adapter.py | 334 ++++++++++++++++++ 4 files changed, 651 insertions(+), 19 deletions(-) create mode 100644 tests/models/minimax_h3/test_checkpoint_adapter.py diff --git a/lightx2v/models/networks/minimax_h3/checkpoint.py b/lightx2v/models/networks/minimax_h3/checkpoint.py index 1e75bc1ed..00e3de4a7 100644 --- a/lightx2v/models/networks/minimax_h3/checkpoint.py +++ b/lightx2v/models/networks/minimax_h3/checkpoint.py @@ -1,17 +1,124 @@ import json import re from collections import defaultdict +from contextlib import contextmanager from pathlib import Path +from typing import NamedTuple +import torch from safetensors import safe_open _H3_BLOCK_KEY_RE = re.compile(r"^transformer_blocks\.(\d+)\.") -class MiniMaxH3ShardCheckpoint: - """Synchronous reader for official MiniMax-H3 sharded safetensors.""" +class TargetSpec(NamedTuple): + name: str + shape: tuple + dtype: str - def __init__(self, checkpoint_dir): + +class SourcePlan(NamedTuple): + source_name: str + source_shard: str + transform: str + shape: tuple + dtype: str + targets: tuple + + +def validate_mapping(entries, source_names, target_names): + """Require an exact source classification and exactly one producer per target.""" + sources, targets = [], [] + for entry in entries: + sources.append(entry.source_name) + targets.extend(target.name for target in entry.targets) + for label, actual, expected in (("source", sources, set(source_names)), ("target", targets, set(target_names))): + if len(actual) != len(set(actual)): + raise ValueError(f"MiniMax-H3 duplicate {label} coverage") + missing, unknown = sorted(expected - set(actual)), sorted(set(actual) - expected) + if missing or unknown: + raise ValueError(f"MiniMax-H3 {label} mapping mismatch: missing={missing}, unknown={unknown}") + + +def _official_schema(config): + """Logical checkpoint shapes only; runtime MM transposes belong to the writer.""" + h = config["hidden_size"] + if config.get("adaln_out_features", 18 * h) != 18 * h or config.get("final_adaln_out_features", 2 * h) != 2 * h: + raise ValueError("MiniMax-H3 incompatible AdaLN output dimensions") + inner = config["num_attention_heads"] * config["attention_head_dim"] + ffn, time = config["ffn_dim"], config["time_embed_dim"] + patch = config["in_channels"] + for axis in config["patch_size"]: + patch *= axis + schema = {} + + def add(source, target, shape, dtype="BF16", transform="rename", target_shape=None): + names = (target,) if isinstance(target, str) else target + schema[source] = (transform, tuple(shape), dtype, tuple(TargetSpec(n, tuple(target_shape or shape), dtype) for n in names)) + + for source, target, shape, dtype in ( + ("video_patch_proj", "proj_in", (h, patch), "F32"), + ("audio_patch_proj", "audio_proj_in", (h, config["audio_in_channels"]), "F32"), + ("condition_proj", "context_embedder", (h, config["text_dim"]), "BF16"), + ("time_embedder.proj_in", "time_embedder.linear_1", (config["time_embed_hidden_dim"], config["freq_dim"]), "F32"), + ("time_embedder.proj_out", "time_embedder.linear_2", (time, config["time_embed_hidden_dim"]), "F32"), + ("final_layer.adaln_proj.linear", "norm_out.linear", (2 * h, time), "BF16"), + ("final_layer.video_out", "proj_out", (patch, h), "F32"), + ("final_layer.audio_out", "audio_proj_out", (config["audio_in_channels"], h), "F32"), + ): + add(source + ".weight", target + ".weight", shape, dtype) + add(source + ".bias", target + ".bias", shape[:1], dtype) + add("token_refiner.final_norm.weight", "token_refiner.final_norm.weight", (h,), transform="identity") + add("final_layer.norm.weight", "norm_out.norm.weight", (h,)) + add("rope.inv_freq", (), (config["rope_freq_dim"],), "F32", "validate_rope") + for source_prefix, target_prefix, layers, adaln in ( + ("blocks", "transformer_blocks", config["num_layers"], True), + ("token_refiner.blocks", "token_refiner.refiner_blocks", config["num_refiner_layers"], False), + ): + for i in range(layers): + source, target = f"{source_prefix}.{i}", f"{target_prefix}.{i}" + add(source + ".attn.qkv_proj.weight", tuple(target + f".attn.to_{q}.weight" for q in "qkv"), (3 * inner, h), transform="qkv_head_interleaved", target_shape=(inner, h)) + for q in "qk": + add(source + f".attn.{q}_norm.weight", target + f".attn.norm_{q}.weight", (config["attention_head_dim"],)) + add(source + ".attn.out_proj.weight", target + ".attn.to_out.0.weight", (h, inner)) + add(source + ".mlp.fc1.weight", target + ".ff.net.0.proj.weight", (2 * ffn, h), transform="swap_gate_value") + add(source + ".mlp.fc2.weight", target + ".ff.net.2.weight", (h, ffn)) + for n in (1, 2): + add(source + f".norm{n}.weight", target + f".norm{n}.weight", (h,)) + if adaln: + add(source + ".adaln_proj.linear.weight", target + ".adaln_proj.linear.weight", (18 * h, time)) + add(source + ".adaln_proj.linear.bias", target + ".adaln_proj.linear.bias", (18 * h,)) + return schema + + +_CONFIG_ALIASES = { + "token_refiner_num_layers": "num_refiner_layers", + "ffn_hidden_size": "ffn_dim", + "latents_dim": "in_channels", + "audio_latents_dim": "audio_in_channels", + "timestep_input_dim": "freq_dim", + "time_embed_hidden_size": "time_embed_hidden_dim", + "rope_inv_freq_len": "rope_freq_dim", +} + + +def _native_config(config): + normalized = dict(config) + if "patch_size" in normalized: + normalized["patch_size"] = tuple(normalized["patch_size"]) + for source, target in _CONFIG_ALIASES.items(): + if source in config: + if target in config and config[source] != config[target]: + raise ValueError(f"MiniMax-H3 conflicting config fields: {source}, {target}") + normalized[target] = config[source] + normalized.setdefault("rope_theta", 10000.0) + return normalized + + +class MiniMaxH3CheckpointPlan: + """Index/config-only planning. No shard existence checks or tensor reads.""" + + def __init__(self, checkpoint_dir, config=None): self.checkpoint_dir = Path(checkpoint_dir) self.index_path = self.checkpoint_dir / "model.safetensors.index.json" if not self.index_path.is_file(): @@ -37,16 +144,39 @@ def __init__(self, checkpoint_dir): ) self.weight_map = dict(weight_map) - shard_names = set(self.weight_map.values()) - missing_shards = sorted( - shard_name - for shard_name in shard_names - if not (self.checkpoint_dir / shard_name).is_file() - ) - if missing_shards: - raise FileNotFoundError( - f"MiniMax-H3 safetensors index references missing shard files: {missing_shards}" - ) + raw = any(n.startswith(("blocks.", "video_patch_proj.", "audio_patch_proj.", "condition_proj.", "final_layer.", "token_refiner.blocks.")) for n in weight_map) + native = any(n.startswith(("transformer_blocks.", "proj_in.", "audio_proj_in.", "context_embedder.", "norm_out.", "token_refiner.refiner_blocks.")) for n in weight_map) + if raw and native: + raise ValueError("MiniMax-H3 mixed official raw and native checkpoint keys") + self.format = "official_raw" if raw else "native" + self.entries = {} + self.targets = {} + if raw: + signature = {"blocks.0.attn.qkv_proj.weight", "blocks.0.mlp.fc1.weight", "video_patch_proj.weight", "final_layer.video_out.weight"} + if not signature.issubset(weight_map): + raise ValueError(f"MiniMax-H3 incomplete official signature: missing {sorted(signature - weight_map.keys())}") + config_path = self.checkpoint_dir / "config.json" + self.config = _native_config(json.loads(config_path.read_text()) if config_path.is_file() else (config or {})) + schema = _official_schema(self.config) + if config is not None: + runtime = _native_config(config) + keys = ( + "hidden_size", "num_layers", "num_attention_heads", "attention_head_dim", "ffn_dim", "time_embed_dim", + "num_refiner_layers", "freq_dim", "rope_freq_dim", "rope_theta", "time_embed_hidden_dim", "in_channels", + "audio_in_channels", "text_dim", "patch_size", + ) + for key in keys: + if key == "rope_theta" and key not in config: + continue + if key in runtime and runtime[key] != self.config[key]: + raise ValueError(f"MiniMax-H3 checkpoint/runtime config mismatch: {key}") + missing, unknown = sorted(schema.keys() - weight_map.keys()), sorted(weight_map.keys() - schema.keys()) + if missing or unknown: + raise ValueError(f"MiniMax-H3 source/target mapping mismatch: missing={missing}, unknown={unknown}") + self.entries = {name: SourcePlan(name, weight_map[name], *schema[name]) for name in sorted(schema)} + expected = [target.name for _, _, _, targets in schema.values() for target in targets] + validate_mapping(self.entries.values(), weight_map, expected) + self.targets = {target.name: (entry, target) for entry in self.entries.values() for target in entry.targets} @property def tensor_names(self): @@ -54,22 +184,25 @@ def tensor_names(self): @property def block_indices(self): + pattern = re.compile(r"^blocks\.(\d+)\.") if self.format == "official_raw" else _H3_BLOCK_KEY_RE return tuple( sorted( { int(match.group(1)) for name in self.weight_map - if (match := _H3_BLOCK_KEY_RE.match(name)) is not None + if (match := pattern.match(name)) is not None } ) ) def tensor_names_for_block(self, block_index): - block_prefix = f"transformer_blocks.{int(block_index)}." + prefix = "blocks" if self.format == "official_raw" else "transformer_blocks" + block_prefix = f"{prefix}.{int(block_index)}." return tuple(sorted(name for name in self.weight_map if name.startswith(block_prefix))) def non_block_tensor_names(self): - return tuple(sorted(name for name in self.weight_map if _H3_BLOCK_KEY_RE.match(name) is None)) + pattern = re.compile(r"^blocks\.(\d+)\.") if self.format == "official_raw" else _H3_BLOCK_KEY_RE + return tuple(sorted(name for name in self.weight_map if pattern.match(name) is None)) def shard_for_tensor(self, name): try: @@ -83,7 +216,157 @@ def block_names(self, block_index): def non_block_names(self): return list(self.non_block_tensor_names()) + def shards_for_sources(self, names): + return tuple(sorted({self.shard_for_tensor(name) for name in names})) + + +class MiniMaxH3SelectedSourceReader: + """Bounded CPU slice staging, with no large mmap views retained across sources.""" + + def __init__(self, plan, row_chunk_size=128): + if plan.format != "official_raw": + raise ValueError("MiniMax-H3 selected adapter requires official raw format") + if row_chunk_size < 1: + raise ValueError("row_chunk_size must be positive") + self.plan = plan + self.row_chunk_size = row_chunk_size + + @contextmanager + def _source(self, name): + entry = self.plan.entries[name] + path = self.plan.checkpoint_dir / entry.source_shard + if not path.is_file(): + raise FileNotFoundError(f"MiniMax-H3 requested shard missing for {name}: {path}") + with safe_open(path, framework="pt", device="cpu") as reader: + source = reader.get_slice(name) + if tuple(source.get_shape()) != entry.shape: + raise ValueError(f"MiniMax-H3 shape mismatch for {name}: {source.get_shape()} != {entry.shape}") + if source.get_dtype() != entry.dtype: + raise ValueError(f"MiniMax-H3 dtype mismatch for {name}: {source.get_dtype()} != {entry.dtype}") + yield source + + def read_source_slice(self, name, row_start, row_end): + entry = self.plan.entries[name] + if not 0 <= row_start < row_end <= entry.shape[0]: + raise ValueError(f"MiniMax-H3 invalid row slice for {name}: {row_start}:{row_end}") + with self._source(name) as source: + # get_slice's PyTorch result can retain the whole source mmap storage. + # Copy only the requested rows so callers cannot retain that mapping. + return source[row_start:row_end].clone() + + def validate_sources(self, names): + for name in names: + with self._source(name) as source: + if name == "rope.inv_freq": + n = self.plan.config["rope_freq_dim"] + expected = 1.0 / (self.plan.config["rope_theta"] ** (torch.arange(0, 2 * n, 2, dtype=torch.float32, device="cpu") / (2 * n))) + actual = source[:].clone() + if not torch.equal(actual.view(torch.int32), expected.view(torch.int32)): + raise ValueError("MiniMax-H3 rope.inv_freq differs from native reconstruction") + + def _ranges(self, entry, requested): + if entry.transform == "qkv_head_interleaved": + dim = self.plan.config["attention_head_dim"] + for head in range(self.plan.config["num_attention_heads"]): + for component, target in enumerate(entry.targets): + if target.name in requested: + yield target.name, (head * 3 + component) * dim, head * dim, dim + elif entry.transform == "swap_gate_value": + half = entry.shape[0] // 2 + for source_start, target_start in ((half, 0), (0, half)): + for row in range(0, half, self.row_chunk_size): + yield entry.targets[0].name, source_start + row, target_start + row, min(self.row_chunk_size, half - row) + else: + size = entry.shape[0] + step = size if len(entry.shape) == 1 else self.row_chunk_size + for row in range(0, size, step): + yield entry.targets[0].name, row, row, min(step, size - row) + + def write_targets(self, destinations): + """Write {logical_name: (runtime_tensor, transpose)} in place, never reallocating.""" + requested = set(destinations) + unknown = requested - self.plan.targets.keys() + if unknown: + raise KeyError(f"MiniMax-H3 missing target plan: {sorted(unknown)}") + sources = sorted({self.plan.targets[name][0].source_name for name in requested}) + self.validate_sources(sources) + for name, (tensor, transpose) in destinations.items(): + spec = self.plan.targets[name][1] + shape = tuple(reversed(spec.shape)) if transpose else spec.shape + dtype = torch.float32 if spec.dtype == "F32" else torch.bfloat16 + if tuple(tensor.shape) != shape or tensor.dtype != dtype: + raise ValueError(f"MiniMax-H3 destination shape/dtype mismatch for {name}: expected {shape}, {dtype}") + for name in sources: + entry = self.plan.entries[name] + with self._source(name) as source: + for target, start, out_start, rows in self._ranges(entry, requested): + tile = source[start : start + rows].clone() + destination, transpose = destinations[target] + if transpose: + destination[:, out_start : out_start + rows].copy_(tile.t()) + else: + destination[out_start : out_start + rows].copy_(tile) + del tile + + def load_modules(self, roots, device="cpu", block_index=None, reusable=False): + """Bind official slices to native base_attrs; transpose exactly at this boundary. + + Reusable leaves retain both their *_cuda_buffer and active tensor identity. + Pre/post CPU leaves use pin_* attributes, matching native offload semantics. + No attention, RoPE, or padding computation is changed here. + """ + destinations, bindings, visited = {}, [], set() + stack = list(roots) + while stack: + module = stack.pop() + if id(module) in visited: + continue + visited.add(id(module)) + for name, attr, transpose in getattr(module, "base_attrs", ()): + if block_index is not None: + name = _H3_BLOCK_KEY_RE.sub(f"transformer_blocks.{int(block_index)}.", name) + if name in destinations: + raise ValueError(f"MiniMax-H3 duplicate target binding: {name}") + if name not in self.plan.targets: + raise KeyError(f"MiniMax-H3 missing target plan: {name}") + spec = self.plan.targets[name][1] + dtype = torch.float32 if spec.dtype == "F32" else torch.bfloat16 + storage_attr = f"{attr}_cuda_buffer" if reusable else f"pin_{attr}" if torch.device(device).type == "cpu" else attr + tensor = getattr(module, storage_attr, None) + if tensor is None: + tensor = torch.empty(spec.shape, dtype=dtype, device=device) + if transpose: + tensor = tensor.t() + setattr(module, storage_attr, tensor) + elif tensor.device.type != torch.device(device).type or (torch.device(device).index is not None and tensor.device.index != torch.device(device).index): + raise ValueError(f"MiniMax-H3 destination device mismatch: {name}") + destinations[name] = (tensor, transpose) + bindings.append((module, attr, tensor)) + stack.extend(child for child in getattr(module, "_modules", {}).values() if child is not None) + stack.extend(child for child in getattr(module, "_parameters", {}).values() if child is not None) + self.write_targets(destinations) + for module, attr, tensor in bindings: + setattr(module, attr, tensor if reusable or tensor.device.type != "cpu" else None) + if hasattr(module, "bias_name") and module.bias_name is None: + module.bias = None + module.pin_bias = None + + +class MiniMaxH3ShardCheckpoint(MiniMaxH3CheckpointPlan): + """Production entry: retain fail-fast validation of the complete shard set.""" + + def __init__(self, checkpoint_dir, config=None): + super().__init__(checkpoint_dir, config=config) + missing = sorted(name for name in set(self.weight_map.values()) if not (self.checkpoint_dir / name).is_file()) + if missing: + raise FileNotFoundError(f"MiniMax-H3 safetensors index references missing shard files: {missing}") + self.selected_reader = MiniMaxH3SelectedSourceReader(self) if self.format == "official_raw" else None + if self.selected_reader is not None: + self.selected_reader.validate_sources(["rope.inv_freq"]) + def load_tensors(self, names, device="cpu"): + if self.selected_reader is not None: + raise ValueError("MiniMax-H3 official raw tensors must use the selected slice adapter") missing = sorted(name for name in names if name not in self.weight_map) if missing: raise KeyError(f"MiniMax-H3 checkpoint is missing requested tensors: {missing}") @@ -101,4 +384,4 @@ def load_tensors(self, names, device="cpu"): return tensors -__all__ = ["MiniMaxH3ShardCheckpoint"] +__all__ = ["MiniMaxH3CheckpointPlan", "MiniMaxH3SelectedSourceReader", "MiniMaxH3ShardCheckpoint"] diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 2140dc04a..5b0288f33 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -146,11 +146,14 @@ def _init_weights(self, weight_dict=None): if weight_dict is not None: raise ValueError("MiniMax-H3 dit_disk_streaming loads weights directly from the official checkpoint; explicit weight_dict is not supported.") - self.pre_weight = self.pre_weight_class(self.config) self.transformer_weights = self.transformer_weight_class(self.config) + self.pre_weight = self.pre_weight_class(self.config) self.post_weight = self.post_weight_class(self.config) checkpoint = self.transformer_weights.checkpoint + if checkpoint.selected_reader is not None: + checkpoint.selected_reader.load_modules([self.pre_weight, self.post_weight], device="cpu") + return None prepost_tensor_names = _collect_declared_base_tensor_names(self.pre_weight, self.post_weight) missing = sorted(name for name in prepost_tensor_names if name not in checkpoint.weight_map) if missing: diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index ea4200e90..c70a79297 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -178,7 +178,13 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): checkpoint_dir = config.get("dit_original_ckpt") if checkpoint_dir is None: raise ValueError("MiniMax-H3 dit_disk_streaming requires config['dit_original_ckpt'] to point to the official transformer checkpoint directory.") - self.checkpoint = MiniMaxH3ShardCheckpoint(checkpoint_dir) + self.checkpoint = MiniMaxH3ShardCheckpoint(checkpoint_dir, config=config) + if self.checkpoint.selected_reader is not None: + # Raw config aliases (e.g. token_refiner_num_layers) must reach + # the native pre/post constructors as well as the mapping plan. + for name, value in self.checkpoint.config.items(): + config.setdefault(name, value) + self.num_layers = int(config["num_layers"]) expected_block_indices = tuple(range(self.num_layers)) if self.checkpoint.block_indices != expected_block_indices: raise ValueError( @@ -219,6 +225,9 @@ def load_streaming_block(self, block_index): raise IndexError(f"MiniMax-H3 checkpoint does not contain transformer block {block_index}.") self._ensure_streaming_block() + if self.checkpoint.selected_reader is not None: + self.checkpoint.selected_reader.load_modules([self.streaming_block], device=AI_DEVICE, block_index=block_index, reusable=True) + return self.streaming_block tensor_names = self.checkpoint.tensor_names_for_block(block_index) tensors = self.checkpoint.load_tensors(tensor_names, device="cpu") try: @@ -232,6 +241,9 @@ def _ensure_streaming_block(self): return self.streaming_block = MiniMaxH3TransformerBlockWeights(0, self.config, create_cuda_buffer=True) self.add_module("streaming_block", self.streaming_block) + if self.checkpoint.selected_reader is not None: + self.checkpoint.selected_reader.load_modules([self.streaming_block], device=AI_DEVICE, block_index=0, reusable=True) + return block0_tensors = self.checkpoint.load_tensors(self.checkpoint.tensor_names_for_block(0), device="cpu") try: self.streaming_block.load(block0_tensors) diff --git a/tests/models/minimax_h3/test_checkpoint_adapter.py b/tests/models/minimax_h3/test_checkpoint_adapter.py new file mode 100644 index 000000000..f20f41027 --- /dev/null +++ b/tests/models/minimax_h3/test_checkpoint_adapter.py @@ -0,0 +1,334 @@ +import importlib.util +import json +import sys +from pathlib import Path + +import pytest +import torch +from safetensors import safe_open +from safetensors.torch import save_file +from test_model_disk_streaming import _config, h3_model_modules # noqa: F401 +from test_transformer_disk_streaming import h3_modules # noqa: F401 + +ROOT = Path(__file__).parents[3] +spec = importlib.util.spec_from_file_location("h3_checkpoint_adapter_under_test", ROOT / "lightx2v/models/networks/minimax_h3/checkpoint.py") +C = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = C +spec.loader.exec_module(C) + + +def write_raw(path): + config = { + "hidden_size": 3, + "num_layers": 2, + "token_refiner_num_layers": 1, + "num_attention_heads": 2, + "attention_head_dim": 2, + "ffn_hidden_size": 4, + "latents_dim": 1, + "audio_latents_dim": 2, + "patch_size": [1, 1, 2], + "text_dim": 5, + "timestep_input_dim": 4, + "time_embed_hidden_size": 3, + "time_embed_dim": 2, + "rope_inv_freq_len": 2, + } + schema = C._official_schema(C._native_config(config)) + tensors = {} + for i, (name, (_, shape, dtype, _)) in enumerate(schema.items()): + tensor = torch.arange(torch.Size(shape).numel(), dtype=torch.float32).reshape(shape) + i + tensors[name] = tensor.to(torch.float32 if dtype == "F32" else torch.bfloat16) + tensors["rope.inv_freq"] = 1.0 / (10000.0 ** (torch.arange(0, 4, 2, dtype=torch.float32) / 4)) + weight_map = {name: f"shard-{i % 2}.safetensors" for i, name in enumerate(sorted(tensors))} + for shard in set(weight_map.values()): + save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, path / shard) + (path / "config.json").write_text(json.dumps(config)) + (path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + return config, tensors, weight_map + + +@pytest.fixture +def raw(tmp_path): + config, tensors, weight_map = write_raw(tmp_path) + plan = C.MiniMaxH3CheckpointPlan(tmp_path) + return config, tensors, weight_map, plan, C.MiniMaxH3SelectedSourceReader(plan, row_chunk_size=2) + + +def destinations(plan, names, transpose=False): + return { + name: ( + torch.empty(tuple(reversed(plan.targets[name][1].shape)) if transpose else plan.targets[name][1].shape, dtype=torch.float32 if plan.targets[name][1].dtype == "F32" else torch.bfloat16), + transpose, + ) + for name in names + } + + +def test_official_detection_and_complete_config_driven_plan(raw): + _, tensors, _, plan, _ = raw + assert plan.format == "official_raw" + assert plan.block_indices == (0, 1) + assert set(plan.entries) == set(tensors) + C.validate_mapping(plan.entries.values(), tensors, plan.targets) + assert len(plan.entries) == 47 # Tiny config, not the release's 535/638 counts. + assert len(plan.targets) == 52 + assert plan.entries["rope.inv_freq"].targets == () + + +@pytest.mark.parametrize("name", ["transformer_blocks.0.attn.to_q.weight", "proj_in.weight"]) +def test_native_format_not_misclassified(tmp_path, name): + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": {name: "absent.safetensors"}})) + assert C.MiniMaxH3CheckpointPlan(tmp_path).format == "native" + + +def test_mixed_format_rejected(raw, tmp_path): + mapping = dict(raw[2], **{"transformer_blocks.0.attn.to_q.weight": "shard-0.safetensors"}) + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping})) + with pytest.raises(ValueError, match="mixed"): + C.MiniMaxH3CheckpointPlan(tmp_path) + + +@pytest.mark.parametrize("prefix,target", [("blocks.0", "transformer_blocks.0"), ("token_refiner.blocks.0", "token_refiner.refiner_blocks.0")]) +def test_qkv_head_interleave_and_transpose(raw, prefix, target): + _, tensors, _, plan, reader = raw + names = [f"{target}.attn.to_{q}.weight" for q in "qkv"] + raw_qkv = tensors[f"{prefix}.attn.qkv_proj.weight"] + for transpose in (False, True): + dest = destinations(plan, names, transpose) + reader.write_targets(dest) + for component, name in enumerate(names): + expected = torch.vstack([raw_qkv[component * 2 : component * 2 + 2], raw_qkv[6 + component * 2 : 8 + component * 2]]) + actual = dest[name][0].t() if transpose else dest[name][0] + assert actual.shape == (4, 3) + assert torch.equal(actual, expected) + assert not torch.equal(actual, raw_qkv.chunk(3, dim=0)[component]) + + +def test_slice_only_reading_and_owned_slice_storage(raw, monkeypatch): + _, _, _, plan, reader = raw + reads = [] + + class Slice: + def __init__(self, inner, name): + self.inner, self.name = inner, name + + def get_shape(self): + return self.inner.get_shape() + + def get_dtype(self): + return self.inner.get_dtype() + + def __getitem__(self, index): + reads.append((self.name, index.start, index.stop)) + return self.inner[index] + + class Open: + def __init__(self, *args, **kwargs): + self.inner = safe_open(*args, **kwargs) + + def __enter__(self): + self.inner.__enter__() + return self + + def __exit__(self, *args): + return self.inner.__exit__(*args) + + def get_slice(self, name): + return Slice(self.inner.get_slice(name), name) + + def get_tensor(self, name): + pytest.fail("adapter must not get_tensor") + + monkeypatch.setattr(C, "safe_open", Open) + name = "blocks.0.attn.qkv_proj.weight" + sl = reader.read_source_slice(name, 2, 4) + assert sl.untyped_storage().nbytes() == 2 * 3 * 2 + reads.clear() + dest = destinations(plan, [t.name for t in plan.entries[name].targets]) + reader.write_targets(dest) + assert reads == [(name, i, i + 2) for i in range(0, 12, 2)] + + +@pytest.mark.parametrize("prefix,target", [("blocks.0", "transformer_blocks.0"), ("token_refiner.blocks.0", "token_refiner.refiner_blocks.0")]) +def test_fc1_swap_without_full_cat(raw, monkeypatch, prefix, target): + _, tensors, _, plan, reader = raw + name = target + ".ff.net.0.proj.weight" + source = tensors[prefix + ".mlp.fc1.weight"] + dest = destinations(plan, [name], transpose=True) + monkeypatch.setattr(torch, "cat", lambda *a, **k: pytest.fail("no full fused cat")) + reader.write_targets(dest) + actual = dest[name][0].t() + assert torch.equal(actual[:4], source[4:]) + assert torch.equal(actual[4:], source[:4]) + + +def test_norm_refiner_and_nonblock_mapping_and_fp32(raw): + _, tensors, _, plan, reader = raw + names = plan.non_block_tensor_names() + reader.validate_sources(names) + targets = [t.name for name in names for t in plan.entries[name].targets] + dest = destinations(plan, targets) + reader.write_targets(dest) + for source, target in [ + ("video_patch_proj", "proj_in"), + ("audio_patch_proj", "audio_proj_in"), + ("condition_proj", "context_embedder"), + ("time_embedder.proj_in", "time_embedder.linear_1"), + ("time_embedder.proj_out", "time_embedder.linear_2"), + ("final_layer.adaln_proj.linear", "norm_out.linear"), + ("final_layer.video_out", "proj_out"), + ("final_layer.audio_out", "audio_proj_out"), + ]: + for suffix in ("weight", "bias"): + assert torch.equal(dest[f"{target}.{suffix}"][0], tensors[f"{source}.{suffix}"]) + assert dest[f"{target}.{suffix}"][0].dtype == tensors[f"{source}.{suffix}"].dtype + for q in "qk": + assert torch.equal(dest[f"token_refiner.refiner_blocks.0.attn.norm_{q}.weight"][0], tensors[f"token_refiner.blocks.0.attn.{q}_norm.weight"]) + assert torch.equal(dest["norm_out.norm.weight"][0], tensors["final_layer.norm.weight"]) + assert torch.equal(dest["token_refiner.final_norm.weight"][0], tensors["token_refiner.final_norm.weight"]) + + +@pytest.mark.parametrize("failure", ["shape", "dtype", "rope_value"]) +def test_invalid_real_header_or_rope_fails(raw, tmp_path, failure): + _, tensors, weight_map, _, reader = raw + name = "rope.inv_freq" if failure == "rope_value" else "blocks.0.attn.qkv_proj.weight" + tensors[name] = tensors[name][:-1] if failure == "shape" else tensors[name].float() if failure == "dtype" else tensors[name] + 1 + shard = weight_map[name] + save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, tmp_path / shard) + with pytest.raises(ValueError, match="rope.inv_freq" if failure == "rope_value" else failure + " mismatch"): + reader.validate_sources([name]) + + +def test_fp32_special_parameter_must_not_be_bf16(raw, tmp_path): + _, tensors, weight_map, _, reader = raw + name = "video_patch_proj.weight" + tensors[name] = tensors[name].bfloat16() + shard = weight_map[name] + save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, tmp_path / shard) + with pytest.raises(ValueError, match="dtype mismatch"): + reader.validate_sources([name]) + + +def test_wrong_destination_orientation_or_dtype_fails_before_copy(raw): + reader = raw[4] + name = "transformer_blocks.0.attn.to_q.weight" + with pytest.raises(ValueError, match="destination shape/dtype"): + reader.write_targets({name: (torch.zeros(4, 3, dtype=torch.bfloat16), True)}) + with pytest.raises(ValueError, match="destination shape/dtype"): + reader.write_targets({name: (torch.zeros(4, 3, dtype=torch.float32), False)}) + with pytest.raises(KeyError, match="missing target plan"): + reader.write_targets({"unknown": (torch.zeros(1), False)}) + + +@pytest.mark.parametrize("failure", ["unknown", "missing"]) +def test_source_coverage_rejected(raw, tmp_path, failure): + mapping = dict(raw[2]) + if failure == "unknown": + mapping["rope.not_allowed"] = "shard-0.safetensors" + else: + del mapping["blocks.0.norm1.weight"] + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping})) + with pytest.raises(ValueError, match=failure): + C.MiniMaxH3CheckpointPlan(tmp_path) + + +def test_duplicate_and_missing_target_coverage_rejected(raw): + plan = raw[3] + entries = list(plan.entries.values()) + first = next(i for i, e in enumerate(entries) if e.targets) + entry = entries[first] + entries[first] = entry._replace(targets=entry.targets + entry.targets) + with pytest.raises(ValueError, match="duplicate target"): + C.validate_mapping(entries, plan.entries, plan.targets) + entries[first] = entry._replace(targets=()) + with pytest.raises(ValueError, match="target mapping mismatch: missing"): + C.validate_mapping(entries, plan.entries, plan.targets) + + +def test_cross_shard_block_load_and_missing_shard_fail(raw, tmp_path): + _, _, _, plan, reader = raw + sources = plan.tensor_names_for_block(0) + assert len(sources) == 10 + assert len(plan.shards_for_sources(sources)) == 2 + dest = destinations(plan, [t.name for n in sources for t in plan.entries[n].targets]) + reader.write_targets(dest) + (tmp_path / "shard-1.safetensors").unlink() + assert C.MiniMaxH3CheckpointPlan(tmp_path).block_indices == (0, 1) + with pytest.raises(FileNotFoundError, match="requested shard missing"): + reader.write_targets(dest) + with pytest.raises(FileNotFoundError, match="missing shard files"): + C.MiniMaxH3ShardCheckpoint(tmp_path) + + +def test_official_streaming_reuses_storage_and_recreates(tmp_path, monkeypatch, h3_modules): + _, weights_module, _ = h3_modules + config, tensors, _ = write_raw(tmp_path) + config.update(C._native_config(config), dit_disk_streaming=True, dit_original_ckpt=str(tmp_path)) + monkeypatch.setattr(weights_module, "AI_DEVICE", "cpu") + weights = weights_module.MiniMaxH3TransformerWeights(config) + block = weights.streaming_block + records = {} + for name, attr, _ in weights_module._iter_base_attrs(block): + records[name] = attr + + def buffers(): + stack, result = [weights.streaming_block], {} + while stack: + module = stack.pop() + for name, attr, _ in getattr(module, "base_attrs", ()): + tensor = getattr(module, attr) + assert tensor is getattr(module, attr + "_cuda_buffer") + result[name] = (id(tensor), tensor.data_ptr()) + stack.extend(getattr(module, "_modules", {}).values()) + return result + + before = buffers() + assert weights.load_streaming_block(1) is block + assert buffers() == before + assert torch.equal(block.ff.in_proj.weight[:, :4].t(), tensors["blocks.1.mlp.fc1.weight"][4:]) + weights.release_disk_streaming_buffer() + assert weights.streaming_block is None and block.attn.to_q.weight is None + new = weights.load_streaming_block(0) + assert new is not block + before = buffers() + assert weights.load_streaming_block(1) is new and buffers() == before + + +def test_raw_config_drives_native_dimensions_without_runtime_defaults(tmp_path, monkeypatch, h3_modules): + _, weights_module, _ = h3_modules + write_raw(tmp_path) + monkeypatch.setattr(weights_module, "AI_DEVICE", "cpu") + weights = weights_module.MiniMaxH3TransformerWeights({"dit_disk_streaming": True, "dit_original_ckpt": str(tmp_path)}) + assert weights.num_layers == 2 + assert weights.config["hidden_size"] == 3 + assert weights.config["num_refiner_layers"] == 1 + + +def test_official_model_prepost_initialization(tmp_path, monkeypatch, h3_model_modules): + _, _, _, transformer, _, _, model_module = h3_model_modules + config, tensors, _ = write_raw(tmp_path) + monkeypatch.setattr(transformer, "AI_DEVICE", "cpu") + config = _config(tmp_path, **config) + del config["num_refiner_layers"] + model = model_module.MiniMaxH3Model(str(tmp_path), config, torch.device("cpu")) + assert model.config["num_refiner_layers"] == 1 + assert model.config["freq_dim"] == 4 + assert model.pre_weight.proj_in.pin_weight.dtype == torch.float32 + assert model.post_weight.proj_out.pin_weight.dtype == torch.float32 + assert torch.equal(model.pre_weight.context_embedder.pin_weight.t(), tensors["condition_proj.weight"]) + assert torch.equal(model.post_weight.proj_out.pin_weight.t(), tensors["final_layer.video_out.weight"]) + + +def test_quantized_streaming_remains_rejected(tmp_path, h3_modules): + _, weights_module, _ = h3_modules + with pytest.raises(NotImplementedError, match="quantized"): + weights_module.MiniMaxH3TransformerWeights({"dit_disk_streaming": True, "dit_quantized": True}) + + +def test_quantized_nonstreaming_keeps_base_loader(tmp_path, h3_model_modules): + model_module = h3_model_modules[-1] + with pytest.raises(AssertionError, match="full checkpoint loading"): + model_module.MiniMaxH3Model( + str(tmp_path), _config(tmp_path, dit_disk_streaming=False, dit_quantized=True, dit_quant_scheme="int8-torchao", dit_quantized_ckpt=str(tmp_path)), torch.device("cpu") + ) From 44581ec91b437369ec89cada7cef01a5895a7aa5 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Sat, 5 Sep 2026 16:07:43 +0800 Subject: [PATCH 16/31] fix(minimax_h3): cast packed positions before device transfer --- .../models/schedulers/minimax_h3/scheduler.py | 2 +- .../minimax_h3/test_scheduler_layout.py | 69 +++++++++++++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) create mode 100644 tests/models/minimax_h3/test_scheduler_layout.py diff --git a/lightx2v/models/schedulers/minimax_h3/scheduler.py b/lightx2v/models/schedulers/minimax_h3/scheduler.py index fe00f905a..f3935fee4 100644 --- a/lightx2v/models/schedulers/minimax_h3/scheduler.py +++ b/lightx2v/models/schedulers/minimax_h3/scheduler.py @@ -30,7 +30,7 @@ def _make_schedule(num_grid_points: int, shift: float, device) -> tuple[torch.Te def _layout_to_device(layout: MiniMaxH3PackedSequence, device) -> MiniMaxH3PackedSequence: return replace( layout, - position_ids=layout.position_ids.to(device), + position_ids=layout.position_ids.to(device=device, dtype=torch.float32), token_tags=layout.token_tags.to(device), video_indices=layout.video_indices.to(device), audio_indices=layout.audio_indices.to(device), diff --git a/tests/models/minimax_h3/test_scheduler_layout.py b/tests/models/minimax_h3/test_scheduler_layout.py new file mode 100644 index 000000000..e917987e3 --- /dev/null +++ b/tests/models/minimax_h3/test_scheduler_layout.py @@ -0,0 +1,69 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch + +REPO_ROOT = Path(__file__).parents[3] + + +@pytest.fixture +def modules(monkeypatch): + # Keep geometry tests independent of accelerator initialization and media dependencies. + for name in ( + "lightx2v", + "lightx2v.models", + "lightx2v.models.networks", + "lightx2v.models.networks.minimax_h3", + "lightx2v.models.schedulers", + "lightx2v.models.schedulers.minimax_h3", + "lightx2v_platform", + "lightx2v_platform.base", + ): + module = types.ModuleType(name) + module.__path__ = [] + monkeypatch.setitem(sys.modules, name, module) + for name, attrs in ( + ("lightx2v.models.schedulers.scheduler", {"BaseScheduler": object}), + ("lightx2v_platform.base.global_var", {"AI_DEVICE": "cpu"}), + ("lightx2v.models.networks.minimax_h3.packing_ref2av", {"build_ref2av_packed_sequence": None}), + ): + module = types.ModuleType(name) + module.__dict__.update(attrs) + monkeypatch.setitem(sys.modules, name, module) + + def load(name): + spec = importlib.util.spec_from_file_location(name, REPO_ROOT / (name.replace(".", "/") + ".py")) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + return load("lightx2v.models.networks.minimax_h3.packing"), load("lightx2v.models.schedulers.minimax_h3.scheduler") + + +@pytest.mark.parametrize("anchors", [(), ("first", "last")]) +@pytest.mark.parametrize("device", ["cpu", pytest.param("mps", marks=pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS unavailable"))]) +def test_layout_device_boundary_preserves_geometry_and_metadata(modules, anchors, device): + packing, scheduler = modules + layout = packing.build_packed_sequence(torch.ones(1, dtype=torch.long), 37, 2, 2, 207, keyframe_anchors=anchors) + assert layout.position_ids.device.type == "cpu" + assert layout.position_ids.dtype == torch.float64 + expected = layout.position_ids.to(torch.float32) + + actual = scheduler._layout_to_device(layout, device) + + assert actual.position_ids.device.type == device + assert actual.position_ids.dtype == torch.float32 + assert torch.isfinite(actual.position_ids).all() + assert torch.equal(actual.position_ids.cpu(), expected) + assert layout.position_ids.dtype == torch.float64 + for name in ("token_tags", "video_indices", "audio_indices", "text_indices"): + value = getattr(actual, name) + assert value.device.type == device + assert value.dtype == torch.long + assert torch.equal(value.cpu(), getattr(layout, name)) + for name in ("sequence_length", "num_condition_video_rows", "num_condition_audio_rows"): + assert getattr(actual, name) == getattr(layout, name) From 117322193b1c3ecc4aee73734f68b134bea67f18 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Sat, 5 Sep 2026 17:56:34 +0800 Subject: [PATCH 17/31] fix(minimax_h3): release Audio VAE weight norm cache --- .../audio_encoders/hf/minimax_h3/audio_vae.py | 14 +++ .../minimax_h3/test_audio_vae_offload.py | 101 ++++++++++++++++++ 2 files changed, 115 insertions(+) create mode 100644 tests/models/minimax_h3/test_audio_vae_offload.py diff --git a/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py b/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py index 7b82dafc2..81a27fc07 100644 --- a/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py +++ b/lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py @@ -33,6 +33,7 @@ import torch.nn as nn import torch.nn.functional as F from torch.nn.utils import weight_norm +from torch.nn.utils.weight_norm import WeightNorm from lightx2v.models.video_encoders.hf.minimax_h3.weights import ( SafetensorsSubsetReport, @@ -497,9 +498,22 @@ def _activate(self) -> torch.device: def offload(self) -> None: self.to("cpu") + self._clear_weight_norm_derived_tensors() _empty_device_cache(self.execution_device) gc.collect() + def _clear_weight_norm_derived_tensors(self) -> None: + # Module.to() only moves registered parameters/buffers. Legacy + # WeightNorm leaves its computed weight as an ordinary tensor attribute. + # Its pre-hook recreates that attribute from weight_g/weight_v before + # every forward, so dropping it preserves both math and checkpoint keys. + for module in self.modules(): + for hook in module._forward_pre_hooks.values(): + if isinstance(hook, WeightNorm): + name = hook.name + if name not in module._parameters and name not in module._buffers and isinstance(vars(module).get(name), torch.Tensor): + delattr(module, name) + def _prepare_stereo_latents( self, latents: torch.Tensor, diff --git a/tests/models/minimax_h3/test_audio_vae_offload.py b/tests/models/minimax_h3/test_audio_vae_offload.py new file mode 100644 index 000000000..7cdde4a73 --- /dev/null +++ b/tests/models/minimax_h3/test_audio_vae_offload.py @@ -0,0 +1,101 @@ +import importlib.util +import json +import sys +import types +import weakref +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file +from torch.nn.utils.weight_norm import WeightNorm + +REPO_ROOT = Path(__file__).parents[3] + + +@pytest.fixture() +def audio_module(monkeypatch): + def load(name, relative): + spec = importlib.util.spec_from_file_location(name, REPO_ROOT / relative) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + platform = types.ModuleType("lightx2v_platform.base.global_var") + platform.AI_DEVICE = "cpu" + monkeypatch.setitem(sys.modules, platform.__name__, platform) + load("lightx2v.models.video_encoders.hf.minimax_h3.weights", "lightx2v/models/video_encoders/hf/minimax_h3/weights.py") + return load("h3_audio_offload_under_test", "lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py") + + +def tiny_model(audio_module): + model = audio_module.MiniMaxH3AudioVAE( + { + "encoder_dim": 4, + "latent_dim": 8, + "latent_channels": 2, + "num_attention_heads": 2, + "decoder_dim": 8, + "encoder_rates": [2], + "decoder_rates": [2], + "decoder_kernel_sizes": [4], + "resblock_kernel_sizes": [3], + "resblock_dilation_sizes": [[1]], + }, + device="cpu", + cpu_offload=True, + ) + return model.eval().requires_grad_(False) + + +def test_repeated_public_decode_and_checkpoint_contract(audio_module, tmp_path): + model = tiny_model(audio_module) + latent = torch.randn(2, 2, 16) + state = {name: value.clone() for name, value in model.state_dict().items()} + hooks = [(module, key, hook) for module in model.modules() for key, hook in module._forward_pre_hooks.items() if isinstance(hook, WeightNorm)] + assert hooks + weights = [weakref.ref(getattr(module, hook.name)) for module, _, hook in hooks] + outputs = [model.decode(latent) for _ in range(3)] + assert all(ref() is None for ref in weights) + assert all(torch.equal(outputs[0], output) for output in outputs) + assert outputs[0].shape == (1, 2, 32) + for module, key, hook in hooks: + assert module._forward_pre_hooks[key] is hook + assert hook.name not in vars(module) + assert hook.name + "_g" in module._parameters + assert hook.name + "_v" in module._parameters + assert state.keys() == model.state_dict().keys() + assert all(torch.equal(value, model.state_dict()[name]) for name, value in state.items()) + model.offload() # Cleanup is also safe when the derived attribute is absent. + model.to("cpu") + model.load_state_dict(state, strict=True) + assert torch.equal(outputs[0], model.decode(latent)) + + component = tmp_path / "audio_vae" + component.mkdir() + (component / "config.json").write_text(json.dumps(model.config)) + save_file(model.state_dict(), component / "model.safetensors") + loaded = audio_module.MiniMaxH3AudioVAE.from_pretrained(tmp_path, device="cpu", cpu_offload=True) + assert set(loaded.load_report.loaded_keys) == set(state) + assert torch.equal(outputs[0], loaded.decode(latent)) + assert all(parameter.device.type == "cpu" for parameter in loaded.parameters()) + assert all(buffer.device.type == "cpu" for buffer in loaded.buffers()) + + +def test_cleanup_only_targets_legacy_hook_attributes(audio_module): + model = tiny_model(audio_module) + plain = torch.nn.Linear(2, 2) + plain.cache = torch.ones(2) + model.unrelated = plain + custom = torch.nn.Module() + custom.register_parameter("kernel", torch.nn.Parameter(torch.randn(2, 2))) + torch.nn.utils.weight_norm(custom, name="kernel", dim=1) + model.custom = custom + parameter = plain.weight + cache = plain.cache + model.offload() + assert plain.weight is parameter + assert plain.cache is cache + assert "kernel" not in vars(custom) + assert set(custom.state_dict()) == {"kernel_g", "kernel_v"} From ad3a7afacf47bf45a5993a0cd3deac0ab36dda52 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Sat, 5 Sep 2026 21:43:39 +0800 Subject: [PATCH 18/31] fix(mps): chunk MiniMax H3 SDPA queries --- configs/platforms/mps/minimax_h3_t2av.json | 1 + lightx2v/common/ops/attn/torch_sdpa.py | 33 ++++++- .../minimax_h3/infer/transformer_infer.py | 5 + .../minimax_h3/test_mps_low_memory_config.py | 1 + .../minimax_h3/test_query_chunked_sdpa.py | 96 +++++++++++++++++++ 5 files changed, 135 insertions(+), 1 deletion(-) create mode 100644 tests/models/minimax_h3/test_query_chunked_sdpa.py diff --git a/configs/platforms/mps/minimax_h3_t2av.json b/configs/platforms/mps/minimax_h3_t2av.json index 2fb44c0ad..c717d4071 100644 --- a/configs/platforms/mps/minimax_h3_t2av.json +++ b/configs/platforms/mps/minimax_h3_t2av.json @@ -29,6 +29,7 @@ "warmup": false, "attn_type": "torch_sdpa", + "mps_sdpa_query_chunk_size": 512, "rms_type": "torch_native", "rope_type": "torch_real_rope", diff --git a/lightx2v/common/ops/attn/torch_sdpa.py b/lightx2v/common/ops/attn/torch_sdpa.py index da61adbf4..15fc01a56 100644 --- a/lightx2v/common/ops/attn/torch_sdpa.py +++ b/lightx2v/common/ops/attn/torch_sdpa.py @@ -8,6 +8,33 @@ from .template import AttnWeightTemplate +def _use_h3_mps_query_chunks(q, k, v, chunk_size, scope, attn_mask, causal, drop_rate): + return ( + scope == "minimax_h3_dit" + and isinstance(chunk_size, int) + and chunk_size > 0 + and q.device.type == k.device.type == v.device.type == "mps" + and q.ndim == 4 + and q.shape == k.shape == v.shape + and q.shape[0] == 1 + and q.shape[1] == 56 + and q.shape[-1] == 128 + and q.shape[2] > 0 + and attn_mask is None + and not causal + and drop_rate == 0 + ) + + +def _query_chunked_sdpa(q, k, v, chunk_size): + # Each query still attends to every key/value. Only the query workspace is + # bounded; there is no context truncation or change to the softmax domain. + return torch.cat( + [F.scaled_dot_product_attention(q[:, :, start : start + chunk_size, :], k, v, attn_mask=None, dropout_p=0.0, is_causal=False) for start in range(0, q.shape[2], chunk_size)], + dim=2, + ) + + @ATTN_WEIGHT_REGISTER("torch_sdpa") class TorchSDPAWeight(AttnWeightTemplate): def __init__(self): @@ -44,7 +71,11 @@ def apply( enable_mem_efficient=True, ) with sdpa_ctx: - x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal) + chunk_size = kwargs.get("mps_sdpa_query_chunk_size", 0) + if _use_h3_mps_query_chunks(q, k, v, chunk_size, kwargs.get("attention_scope"), attn_mask, causal, drop_rate): + x = _query_chunked_sdpa(q, k, v, chunk_size) + else: + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal) x = x.transpose(1, 2) b, s, a, d = x.shape out = x.reshape(b, s, -1) diff --git a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py index 451981af7..0aafed967 100644 --- a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py @@ -67,6 +67,11 @@ def _attention(self, weights, hidden_states, pre_infer_out): "block_idx": self.block_idx, } if sp_state is None: + if self.config.get("attn_type") == "torch_sdpa": + attention_kwargs.update( + attention_scope="minimax_h3_dit", + mps_sdpa_query_chunk_size=self.config.get("mps_sdpa_query_chunk_size", 0), + ) seq_len = q.shape[0] cu_seqlens = torch.tensor((0, seq_len), dtype=torch.int32, device=q.device) out = weights.calculate.apply( diff --git a/tests/models/minimax_h3/test_mps_low_memory_config.py b/tests/models/minimax_h3/test_mps_low_memory_config.py index 1aaff4d72..171350da0 100644 --- a/tests/models/minimax_h3/test_mps_low_memory_config.py +++ b/tests/models/minimax_h3/test_mps_low_memory_config.py @@ -40,6 +40,7 @@ def test_mps_minimax_h3_config_enables_low_memory_streaming(): assert config["unload_modules"] is False assert config["warmup"] is False assert config["attn_type"] == "torch_sdpa" + assert config["mps_sdpa_query_chunk_size"] == 512 assert config["rms_type"] == "torch_native" assert config["rope_type"] == "torch_real_rope" assert config["vae_attn_type"] == "torch_sdpa" diff --git a/tests/models/minimax_h3/test_query_chunked_sdpa.py b/tests/models/minimax_h3/test_query_chunked_sdpa.py new file mode 100644 index 000000000..8e38e47be --- /dev/null +++ b/tests/models/minimax_h3/test_query_chunked_sdpa.py @@ -0,0 +1,96 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F + + +@pytest.fixture() +def sdpa(monkeypatch): + root = Path(__file__).parents[3] + registry = types.ModuleType("lightx2v.utils.registry_factory") + registry.ATTN_WEIGHT_REGISTER = lambda name: lambda cls: cls + monkeypatch.setitem(sys.modules, registry.__name__, registry) + package = types.ModuleType("chunked_sdpa_test") + package.__path__ = [str(root / "lightx2v/common/ops/attn")] + monkeypatch.setitem(sys.modules, package.__name__, package) + name = package.__name__ + ".torch_sdpa" + spec = importlib.util.spec_from_file_location(name, root / "lightx2v/common/ops/attn/torch_sdpa.py") + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize("length", [16, 19]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_dense_numerical_parity(sdpa, length, dtype): + torch.manual_seed(123) + q, k, v = [torch.randn(1, 4, length, 8, dtype=dtype) for _ in range(3)] + expected = F.scaled_dot_product_attention(q, k, v) + actual = sdpa._query_chunked_sdpa(q, k, v, 8) + assert actual.shape == expected.shape and actual.dtype == dtype + tolerance = 2e-2 if dtype == torch.bfloat16 else 1e-5 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + +@pytest.mark.parametrize( + "change", + [ + {}, + {"chunk_size": 0}, + {"scope": None}, + {"scope": "video_vae"}, + {"attn_mask": torch.ones(1)}, + {"causal": True}, + {"drop_rate": 0.1}, + {"device": "cpu"}, + {"device": "cuda"}, + {"shape": (1, 8, 19, 128)}, + ], +) +def test_applicability(sdpa, change): + args = {"chunk_size": 8, "scope": "minimax_h3_dit", "attn_mask": None, "causal": False, "drop_rate": 0} + device = change.get("device", "mps") + shape = change.get("shape", (1, 56, 19, 128)) + tensor = types.SimpleNamespace(device=torch.device(device), ndim=4, shape=shape) + args.update({k: v for k, v in change.items() if k not in ("device", "shape")}) + assert sdpa._use_h3_mps_query_chunks(tensor, tensor, tensor, **args) == (not change) + + +@pytest.mark.parametrize("options", [{}, {"attn_mask": torch.ones(9, 9, dtype=torch.bool)}, {"causal": True}, {"drop_rate": 0.2}]) +def test_original_path_preserved_on_cpu(sdpa, monkeypatch, options): + q, k, v = [torch.randn(9, 4, 8) for _ in range(3)] + calls = [] + original = F.scaled_dot_product_attention + + def record(*args, **kwargs): + calls.append(kwargs) + return original(*args, **kwargs) + + monkeypatch.setattr(sdpa.F, "scaled_dot_product_attention", record) + result = sdpa.TorchSDPAWeight().apply(q, k, v, attention_scope="minimax_h3_dit", mps_sdpa_query_chunk_size=512, **options) + assert result.shape == (9, 32) and result.dtype == q.dtype + assert len(calls) == 1 + assert calls[0]["is_causal"] == options.get("causal", False) + assert calls[0]["dropout_p"] == options.get("drop_rate", 0) + assert calls[0]["attn_mask"] is options.get("attn_mask") + + +def test_query_chunks_keep_complete_key_value_context(sdpa, monkeypatch): + q, k, v = [torch.randn(1, 4, 19, 8) for _ in range(3)] + lengths = [] + original = F.scaled_dot_product_attention + + def record(query, key, value, **kwargs): + assert key is k and value is v + assert kwargs == {"attn_mask": None, "dropout_p": 0.0, "is_causal": False} + lengths.append(query.shape[2]) + return original(query, key, value, **kwargs) + + monkeypatch.setattr(sdpa.F, "scaled_dot_product_attention", record) + assert sdpa._query_chunked_sdpa(q, k, v, 8).shape == q.shape + assert lengths == [8, 8, 3] From 4c82ebbab2a5d1763b08e41a8d37214aee5b2b9d Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Sun, 6 Sep 2026 11:01:20 +0800 Subject: [PATCH 19/31] fix(minimax_h3): match reference RoPE precision --- .../minimax_h3/weights/transformer_weights.py | 4 +- .../models/minimax_h3/test_rope_precision.py | 65 +++++++++++++++++++ 2 files changed, 68 insertions(+), 1 deletion(-) create mode 100644 tests/models/minimax_h3/test_rope_precision.py diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index c70a79297..75b67984b 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -99,7 +99,9 @@ def __init__(self, prefix, config, create_cuda_buffer=False): "rope", ROPE_REGISTER[config.get("rope_type", "torch_real_rope")]( layout="split_half", - compute_dtype=torch.float32, + # H3 requires BF16 Q/K; the reference rounds frequencies and + # performs the rotation at that dtype, not in FP32. + compute_dtype=torch.bfloat16, ), ) attn_type = config.get("attn_type", "flash_attn3") diff --git a/tests/models/minimax_h3/test_rope_precision.py b/tests/models/minimax_h3/test_rope_precision.py new file mode 100644 index 000000000..326d043b4 --- /dev/null +++ b/tests/models/minimax_h3/test_rope_precision.py @@ -0,0 +1,65 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch + + +@pytest.fixture() +def h3_rope(monkeypatch): + root = Path(__file__).parents[3] + spec = importlib.util.spec_from_file_location("h3_rope_weight_helpers", Path(__file__).with_name("test_transformer_disk_streaming.py")) + helper = importlib.util.module_from_spec(spec) + spec.loader.exec_module(helper) + _, weights, _ = helper.h3_modules.__wrapped__(monkeypatch) + + registry = sys.modules["lightx2v.utils.registry_factory"] + rope_registry = registry.ROPE_REGISTER + monkeypatch.setattr(registry, "ROPE_REGISTER", lambda name: lambda cls: cls) + magi = types.ModuleType("lightx2v.common.magi_custom_op_mode") + magi.use_magi_custom_ops = lambda: False + monkeypatch.setitem(sys.modules, magi.__name__, magi) + package = types.ModuleType("h3_precision_rope") + package.__path__ = [str(root / "lightx2v/common/ops/rope")] + monkeypatch.setitem(sys.modules, package.__name__, package) + spec = importlib.util.spec_from_file_location(package.__name__ + ".torch_rope", Path(package.__path__[0]) / "torch_rope.py") + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + monkeypatch.setattr(registry, "ROPE_REGISTER", rope_registry) + monkeypatch.setitem(rope_registry, "torch_real_rope", module.TorchRealRope) + monkeypatch.setitem(weights.ROPE_REGISTER, "torch_real_rope", module.TorchRealRope) + return weights.MiniMaxH3AttentionWeights("transformer_blocks.0.attn", {}).rope, module.TorchRealRope + + +def test_h3_precision_does_not_change_generic_default(h3_rope): + rope, generic = h3_rope + assert rope.layout == "split_half" + assert rope.compute_dtype == torch.bfloat16 + assert generic().compute_dtype == torch.float32 + + +@pytest.mark.parametrize("length", [1, 17, 52]) +def test_h3_bf16_rope_matches_reference_exactly(h3_rope, length): + rope, _ = h3_rope + generator = torch.Generator().manual_seed(123) + q, k = [torch.randn(length, 56, 128, generator=generator, dtype=torch.bfloat16) for _ in range(2)] + angles = torch.randn(length, 48, generator=generator, dtype=torch.float32) + angles = torch.cat((angles, angles), dim=-1) + cos, sin = angles.cos(), angles.sin() + + def reference(x): + rotary, passthrough = x[..., :96], x[..., 96:] + first, second = rotary.chunk(2, dim=-1) + rotated = torch.cat((-second, first), dim=-1) + out = rotary * cos.to(x.dtype)[:, None, :] + rotated * sin.to(x.dtype)[:, None, :] + return torch.cat((out, passthrough), dim=-1) + + actual = rope.apply(q, k, (cos, sin), rotary_dim=96) + for x, out in zip((q, k), actual): + assert out.shape == x.shape + assert out.dtype == torch.bfloat16 + assert torch.equal(out, reference(x)) + assert torch.equal(out[..., 96:], x[..., 96:]) From 6378266691f193e5bb82f43fc9d07e890f69510f Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Sun, 6 Sep 2026 21:35:48 +0800 Subject: [PATCH 20/31] feat(minimax_h3): support LoRA with disk streaming --- lightx2v/models/networks/minimax_h3/model.py | 52 ++- .../networks/minimax_h3/streaming_lora.py | 183 +++++++++ .../minimax_h3/weights/transformer_weights.py | 12 +- .../models/minimax_h3/test_streaming_lora.py | 383 ++++++++++++++++++ 4 files changed, 625 insertions(+), 5 deletions(-) create mode 100644 lightx2v/models/networks/minimax_h3/streaming_lora.py create mode 100644 tests/models/minimax_h3/test_streaming_lora.py diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 5b0288f33..2f0a2d898 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -2,6 +2,7 @@ import glob import math import os +from contextlib import nullcontext import torch import torch.distributed as dist @@ -101,8 +102,10 @@ def __init__(self, model_path, config, device, lora_path=None, lora_strength=1.0 raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support quantized DiT checkpoints yet.") if config.get("tensor_parallel", False): raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support tensor parallel inference yet.") - if lora_path is not None or config.get("lora_configs"): - raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support LoRA yet.") + if (lora_path is not None or config.get("lora_configs")) and not config.get("lora_dynamic_apply", False): + raise NotImplementedError("MiniMax-H3 dit_disk_streaming requires dynamic LoRA; dense load-time merging is not supported.") + if config.get("lora_configs") and lora_path is None: + raise ValueError("MiniMax-H3 streamed LoRA must be initialized through build_minimax_h3_model_with_lora.") if config.get("attn_type") == "sol_attn": reorder = str(config.get("sol_attn_setting", {}).get("reorder", "none")).lower() if reorder != "none": @@ -153,6 +156,7 @@ def _init_weights(self, weight_dict=None): checkpoint = self.transformer_weights.checkpoint if checkpoint.selected_reader is not None: checkpoint.selected_reader.load_modules([self.pre_weight, self.post_weight], device="cpu") + self._init_streaming_lora() return None prepost_tensor_names = _collect_declared_base_tensor_names(self.pre_weight, self.post_weight) missing = sorted(name for name in prepost_tensor_names if name not in checkpoint.weight_map) @@ -171,8 +175,33 @@ def _init_weights(self, weight_dict=None): device_module = getattr(torch, torch.device(self.device).type, None) if device_module is not None and hasattr(device_module, "empty_cache"): device_module.empty_cache() + self._init_streaming_lora() return None + def _init_streaming_lora(self): + if self.lora_path is None: + return + from lightx2v.models.networks.minimax_h3.streaming_lora import MiniMaxH3StreamingLora, streaming_target_shapes + + weights = self.transformer_weights + shapes = streaming_target_shapes(weights.checkpoint, weights.streaming_block, self.pre_weight, self.post_weight) + weights.streaming_lora = MiniMaxH3StreamingLora( + self.lora_path, + normalize_key=self._normalize_dynamic_lora_key, + target_shapes=shapes, + strength=self.lora_strength, + alpha=self.lora_alpha, + dtype=GET_DTYPE(), + ) + adapter = weights.streaming_lora + logger.info( + "Indexed MiniMax-H3 streamed LoRA: {} main-block pairs, {} pre/post pairs, ranks={}, strength={}", + sum(len(pairs) for pairs in adapter.blocks.values()), + len(adapter.resident), + sorted({pair.rank for pair in adapter.pairs.values()}), + adapter.strength, + ) + @staticmethod def _normalize_dynamic_lora_key(key): for prefix in ("base_model.model.", "model.diffusion_model.", "diffusion_model.", "transformer.", "model."): @@ -231,6 +260,8 @@ def _validate_dynamic_lora_shapes(self, source, normalized_sources, down_names): return model_keys, ranks def _load_lora_file(self, file_path, alpha=None): + if self.config.get("dit_disk_streaming", False): + raise NotImplementedError("MiniMax-H3 disk streaming uses the selective LoRA index, not the full-factor loader.") if not os.path.isfile(file_path): raise FileNotFoundError(f"MiniMax-H3 LoRA file not found: {file_path}") @@ -333,6 +364,8 @@ def _register_dynamic_lora_weights(self, lora_weights, strength): logger.info("Registered {} MiniMax-H3 dynamic LoRA branches with strength={}", len(self._pending_dynamic_lora_model_keys), strength) def _register_lora(self, lora_path, strength): + if self.config.get("dit_disk_streaming", False): + raise NotImplementedError("MiniMax-H3 streamed LoRA is configured at initialization; runtime adapter switching is not supported.") lora_weights = self._load_lora_file(lora_path) self._register_dynamic_lora_weights(lora_weights, strength) self.lora_path = lora_path @@ -342,12 +375,20 @@ def _register_lora(self, lora_path, strength): offload_manager.need_init_first_buffer = True def _remove_lora(self): + if self.config.get("dit_disk_streaming", False): + adapter = getattr(self.transformer_weights, "streaming_lora", None) + if adapter is not None: + for root in (self.pre_weight, self.post_weight, self.transformer_weights.streaming_block): + adapter.clear(root) + self.transformer_weights.streaming_lora = None super()._remove_lora() transformer_infer = getattr(self, "transformer_infer", None) if transformer_infer is not None: transformer_infer._clear_adaln_cache() def _update_lora(self, lora_path, strength, alpha=None): + if self.config.get("dit_disk_streaming", False): + raise NotImplementedError("MiniMax-H3 streamed LoRA is configured at initialization; runtime adapter switching is not supported.") if isinstance(lora_path, dict): raise NotImplementedError("MiniMax-H3 dynamic LoRA switching expects one checkpoint path, not a merged tensor dictionary") lora_weights = self._load_lora_file(lora_path, alpha=alpha) @@ -555,13 +596,16 @@ def _infer_cond_uncond(self, inputs, infer_condition=True): if not infer_condition: raise ValueError("MiniMax-H3 does not execute an unconditional pass") prompt_embeds = inputs["text_encoder_output"]["prompt_embeds"] - pre = self.pre_infer.infer(self.pre_weight, prompt_embeds) + adapter = getattr(self.transformer_weights, "streaming_lora", None) + with adapter.resident_scope(self.pre_weight) if adapter is not None else nullcontext(): + pre = self.pre_infer.infer(self.pre_weight, prompt_embeds) if self.config.get("seq_parallel", False): pre = self._seq_parallel_pre_process(pre) hidden_states = self.transformer_infer.infer(self.transformer_weights, pre) if self.config.get("seq_parallel", False): hidden_states = self._seq_parallel_post_process(hidden_states, pre) - return self.post_infer.infer(self.post_weight, hidden_states, pre) + with adapter.resident_scope(self.post_weight) if adapter is not None else nullcontext(): + return self.post_infer.infer(self.post_weight, hidden_states, pre) @torch.no_grad() def infer(self, inputs): diff --git a/lightx2v/models/networks/minimax_h3/streaming_lora.py b/lightx2v/models/networks/minimax_h3/streaming_lora.py new file mode 100644 index 000000000..76eca4447 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/streaming_lora.py @@ -0,0 +1,183 @@ +"""Selective dynamic LoRA for H3's reusable disk-streamed block. + +The index owns metadata, not factor tensors. Arithmetic and alpha/rank scaling +remain in MMWeight.register_lora/apply_lora, as in ordinary dynamic H3 LoRA. +""" + +import math +import re +from contextlib import contextmanager +from dataclasses import dataclass + +import torch +from safetensors import safe_open + + +@dataclass(frozen=True) +class LoraPair: + down_key: str + up_key: str + down_shape: tuple + up_shape: tuple + rank: int + alpha: float + + +class MiniMaxH3StreamingLora: + def __init__(self, path, *, normalize_key, target_shapes, strength, alpha=None, dtype=torch.bfloat16): + self.path = str(path) + self.dtype = dtype + self.strength = float(strength) + if not math.isfinite(self.strength): + raise ValueError("MiniMax-H3 LoRA strength must be finite") + if alpha is not None and (not math.isfinite(float(alpha)) or float(alpha) <= 0): + raise ValueError("MiniMax-H3 LoRA alpha must be finite and positive") + self.pairs = {} + self.blocks = {} + self.resident = {} + # Only pre/post factors may be cached on CPU, never main-block factors. + self._resident_cpu = {} + with safe_open(self.path, framework="pt", device="cpu") as source: + keys = {} + for key in sorted(source.keys()): + name = normalize_key(key) + if name is None: + raise ValueError(f"Unsupported MiniMax-H3 LoRA tensor: {key}") + if name in keys: + raise ValueError(f"MiniMax-H3 LoRA keys collide after normalization: {name}") + keys[name] = key + down = {name.removesuffix(".lora_down.weight") for name in keys if name.endswith(".lora_down.weight")} + up = {name.removesuffix(".lora_up.weight") for name in keys if name.endswith(".lora_up.weight")} + if not down or down != up: + raise ValueError("MiniMax-H3 LoRA has incomplete A/B pairs") + if any(name.removesuffix(".alpha") not in down for name in keys if name.endswith(".alpha")): + raise ValueError("MiniMax-H3 LoRA alpha has no matching pair") + for name in sorted(down): + weight_name = name + ".weight" + if weight_name not in target_shapes: + raise ValueError(f"Unsupported MiniMax-H3 streamed LoRA target: {weight_name}") + a_key, b_key = keys[name + ".lora_down.weight"], keys[name + ".lora_up.weight"] + a, b = source.get_slice(a_key), source.get_slice(b_key) + a_shape, b_shape = tuple(a.get_shape()), tuple(b.get_shape()) + base_shape = tuple(target_shapes[weight_name]) + if len(a_shape) != 2 or len(b_shape) != 2 or a_shape[0] <= 0 or b_shape[1] != a_shape[0] or (b_shape[0], a_shape[1]) != base_shape: + raise ValueError(f"MiniMax-H3 LoRA shape mismatch for {weight_name}: A={a_shape}, B={b_shape}, base={base_shape}") + if a.get_dtype() not in {"BF16", "F16", "F32"} or b.get_dtype() not in {"BF16", "F16", "F32"}: + raise ValueError(f"Unsupported MiniMax-H3 LoRA factor dtype: {weight_name}") + pair_alpha = alpha + if name + ".alpha" in keys: + # Only scalar metadata is read here; never materialize A/B. + alpha_key = keys[name + ".alpha"] + if math.prod(source.get_slice(alpha_key).get_shape()) != 1: + raise ValueError(f"MiniMax-H3 LoRA alpha must be scalar: {name}") + pair_alpha = float(source.get_tensor(alpha_key).item()) + if pair_alpha is None or not math.isfinite(float(pair_alpha)) or float(pair_alpha) <= 0: + raise ValueError(f"MiniMax-H3 LoRA requires finite positive alpha: {name}") + cast_alpha = torch.tensor(pair_alpha, dtype=self.dtype) + if not torch.isfinite(cast_alpha) or cast_alpha <= 0: + raise ValueError(f"MiniMax-H3 LoRA alpha is not representable in {self.dtype}: {name}") + pair = LoraPair(a_key, b_key, a_shape, b_shape, a_shape[0], float(pair_alpha)) + self.pairs[weight_name] = pair + match = re.fullmatch(r"transformer_blocks\.(\d+)\.(.+)", weight_name) + if match: + self.blocks.setdefault(int(match[1]), {})["transformer_blocks.0." + match[2]] = pair + else: + self.resident[weight_name] = pair + + @staticmethod + def weights(root): + stack, visited, weights = [root], set(), {} + while stack: + obj = stack.pop() + if obj is None or id(obj) in visited: + continue + visited.add(id(obj)) + if hasattr(obj, "register_lora") and hasattr(obj, "weight_name"): + weights[obj.weight_name] = obj + stack.extend(getattr(obj, "_modules", {}).values()) + stack.extend(getattr(obj, "_parameters", {}).values()) + return weights + + @classmethod + def clear(cls, root): + weights = [weight for weight in cls.weights(root).values() if getattr(weight, "has_lora_branch", False)] + devices = {weight.lora_down.device for weight in weights} + for device in devices: + if device.type != "cpu": + torch.get_device_module(device).synchronize() + for weight in weights: + weight.remove_lora() + + def bind(self, root, pairs): + self.clear(root) + weights = self.weights(root) + missing = pairs.keys() - weights.keys() + if missing: + raise ValueError(f"MiniMax-H3 streamed LoRA targets not found: {sorted(missing)}") + try: + with safe_open(self.path, framework="pt", device="cpu") as source: + for name, pair in pairs.items(): + weight = weights[name] + # Read only this target, with ordinary unpinned CPU storage. + tensors = self._resident_cpu.get(name) if name in self.resident else None + if tensors is None: + tensors = ( + source.get_tensor(pair.down_key).to(dtype=self.dtype, copy=True), + source.get_tensor(pair.up_key).to(dtype=self.dtype, copy=True), + torch.tensor(pair.alpha, dtype=self.dtype), + ) + if name in self.resident: + self._resident_cpu[name] = tensors + factors = dict(zip((weight.lora_down_name, weight.lora_up_name, weight.lora_alpha_name), tensors)) + weight.register_lora(factors, self.strength) + del factors, tensors + if not getattr(weight, "has_lora_branch", False): + raise RuntimeError(f"MiniMax-H3 streamed LoRA registration failed: {name}") + except Exception: + self.clear(root) + raise + + def load_block(self, block, index): + self.bind(block, self.blocks.get(index, {})) + + @contextmanager + def resident_scope(self, root): + names = self.weights(root) + pairs = {name: pair for name, pair in self.resident.items() if name in names} + if not pairs: + yield + return + self.bind(root, pairs) + try: + yield + finally: + self.clear(root) + + +def streaming_target_shapes(checkpoint, block, *resident_roots): + """Resolve supported MMWeight targets using existing base-checkpoint metadata.""" + names = set() + for name in MiniMaxH3StreamingLora.weights(block): + for index in checkpoint.block_indices: + names.add(name.replace("transformer_blocks.0.", f"transformer_blocks.{index}.", 1)) + for root in resident_roots: + names.update(MiniMaxH3StreamingLora.weights(root)) + shapes = {} + if checkpoint.selected_reader is not None: + for name in names: + spec = checkpoint.targets[name][1] + # Ordinary dynamic H3 LoRA uses BF16 factors. Sensitive FP32 + # linears require a different execution contract and are excluded. + if spec.dtype == "BF16": + shapes[name] = spec.shape + else: + by_shard = {} + for name in names: + by_shard.setdefault(checkpoint.weight_map[name], []).append(name) + for shard, shard_names in by_shard.items(): + with safe_open(checkpoint.checkpoint_dir / shard, framework="pt", device="cpu") as source: + for name in shard_names: + tensor = source.get_slice(name) + if tensor.get_dtype() == "BF16": + shapes[name] = tuple(tensor.get_shape()) + return shapes diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index 75b67984b..8e39976c4 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -165,6 +165,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): self.config = config self.num_layers = int(config.get("num_layers", 50)) self.disk_streaming = bool(config.get("dit_disk_streaming", False)) + self.streaming_lora = None if self.disk_streaming: if config.get("lazy_load", False): raise NotImplementedError( @@ -175,7 +176,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): if config.get("tensor_parallel", False): raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support tensor parallel inference yet.") if lora_path is not None: - raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support LoRA streaming yet.") + raise ValueError("Initialize MiniMax-H3 streamed LoRA through MiniMaxH3Model, not the weights constructor.") checkpoint_dir = config.get("dit_original_ckpt") if checkpoint_dir is None: @@ -227,8 +228,13 @@ def load_streaming_block(self, block_index): raise IndexError(f"MiniMax-H3 checkpoint does not contain transformer block {block_index}.") self._ensure_streaming_block() + if self.streaming_lora is not None: + # Finish the previous use before either base weights or factors change. + self.streaming_lora.clear(self.streaming_block) if self.checkpoint.selected_reader is not None: self.checkpoint.selected_reader.load_modules([self.streaming_block], device=AI_DEVICE, block_index=block_index, reusable=True) + if self.streaming_lora is not None: + self.streaming_lora.load_block(self.streaming_block, block_index) return self.streaming_block tensor_names = self.checkpoint.tensor_names_for_block(block_index) tensors = self.checkpoint.load_tensors(tensor_names, device="cpu") @@ -236,6 +242,8 @@ def load_streaming_block(self, block_index): self.streaming_block.load_state_dict(self._prepare_streaming_state_dict(tensors, block_index), block_index) finally: del tensors + if self.streaming_lora is not None: + self.streaming_lora.load_block(self.streaming_block, block_index) return self.streaming_block def _ensure_streaming_block(self): @@ -261,6 +269,8 @@ def release_disk_streaming_buffer(self): block = self.streaming_block if block is None: return + if self.streaming_lora is not None: + self.streaming_lora.clear(block) with suppress(Exception): device_module = getattr(torch, AI_DEVICE, None) if device_module is not None and hasattr(device_module, "synchronize"): diff --git a/tests/models/minimax_h3/test_streaming_lora.py b/tests/models/minimax_h3/test_streaming_lora.py new file mode 100644 index 000000000..518f0d22d --- /dev/null +++ b/tests/models/minimax_h3/test_streaming_lora.py @@ -0,0 +1,383 @@ +import ast +import importlib.util +import sys +import types +from abc import ABCMeta, abstractmethod +from pathlib import Path + +import pytest +import torch +from loguru import logger +from safetensors.torch import save_file + +ROOT = Path(__file__).parents[3] + + +def load_module(name, path, monkeypatch): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def setup(monkeypatch): + helpers = load_module("h3_stream_lora_test_helpers", Path(__file__).with_name("test_model_disk_streaming.py"), monkeypatch) + modules = helpers.h3_model_modules.__wrapped__(monkeypatch) + monkeypatch.setattr(modules[3], "AI_DEVICE", "cpu") + model = modules[-1].MiniMaxH3Model + streaming = load_module("lightx2v.models.networks.minimax_h3.streaming_lora", ROOT / "lightx2v/models/networks/minimax_h3/streaming_lora.py", monkeypatch) + # Execute the actual production MMWeight classes, excluding optional CUDA + # kernel imports. Only base checkpoint I/O is replaced by the existing tiny + # fixture; register_lora/apply/apply_lora/remove_lora remain production code. + scope = {"torch": torch, "ABCMeta": ABCMeta, "abstractmethod": abstractmethod, "logger": logger, "AI_DEVICE": "cpu"} + for file, names in [ + ("lightx2v/common/ops/utils.py", {"build_lora_and_diff_names"}), + ("lightx2v/common/ops/mm/mm_weight.py", {"MMWeightTemplate", "MMWeight"}), + ]: + tree = ast.parse((ROOT / file).read_text()) + nodes = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names] + for node in nodes: + node.decorator_list = [] + exec(compile(ast.Module(body=nodes, type_ignores=[]), file, "exec"), scope) # noqa: S102 - execute trusted local production classes + + class Linear(scope["MMWeight"]): + load = helpers._FakeLinear.load + load_state_dict = helpers._FakeLinear.load_state_dict + + registry = sys.modules["lightx2v.utils.registry_factory"].MM_WEIGHT_REGISTER + monkeypatch.setitem(registry, "Default", Linear) + monkeypatch.setitem(registry, "Default-ForceFp32", Linear) + return helpers, modules, model, streaming, Linear + + +def fixture_tensors(): + return { + "base_model.model.transformer_blocks.0.attn.to_q.lora_A.default.weight": torch.tensor([[1, 2, 3]], dtype=torch.bfloat16), + "base_model.model.transformer_blocks.0.attn.to_q.lora_B.default.weight": torch.tensor([[1], [2]], dtype=torch.bfloat16), + "transformer_blocks.1.attn.to_k.lora_down.weight": torch.tensor([[2, 1, 0], [0, 1, 2]], dtype=torch.bfloat16), + "transformer_blocks.1.attn.to_k.lora_up.weight": torch.tensor([[1, 0], [0, 1]], dtype=torch.bfloat16), + "transformer_blocks.1.attn.to_k.alpha": torch.tensor(8.0), + "token_refiner.refiner_blocks.0.attn.to_q.lora_A.weight": torch.ones(1, 3, dtype=torch.bfloat16), + "token_refiner.refiner_blocks.0.attn.to_q.lora_B.weight": torch.ones(2, 1, dtype=torch.bfloat16), + } + + +def make_index(tmp_path, setup, tensors=None, **kwargs): + _, _, model, streaming, _ = setup + tensors = fixture_tensors() if tensors is None else tensors + path = tmp_path / "lora.safetensors" + save_file(tensors, path) + shapes = {model._normalize_dynamic_lora_key(key).removesuffix(".lora_down.weight") + ".weight": (2, 3) for key in fixture_tensors() if "lora_A" in key or "lora_down" in key} + kwargs = {"alpha": 4, "strength": 0.5, **kwargs} + return streaming.MiniMaxH3StreamingLora(path, normalize_key=model._normalize_dynamic_lora_key, target_shapes=shapes, **kwargs) + + +def test_index_metadata_only_and_alpha_precedence(tmp_path, setup, monkeypatch): + streaming = setup[3] + original = streaming.safe_open + reads = [] + + class Reader: + def __init__(self, *args, **kwargs): + self.context = original(*args, **kwargs) + + def __enter__(self): + self.source = self.context.__enter__() + return self + + def __exit__(self, *args): + return self.context.__exit__(*args) + + def keys(self): + return self.source.keys() + + def get_slice(self, key): + return self.source.get_slice(key) + + def get_tensor(self, key): + reads.append(key) + assert key.endswith(".alpha"), "Index construction read a factor tensor" + return self.source.get_tensor(key) + + monkeypatch.setattr(streaming, "safe_open", Reader) + index = make_index(tmp_path, setup) + assert len(index.pairs) == 3 + assert set(index.blocks) == {0, 1} + assert len(index.resident) == 1 + assert index.pairs["transformer_blocks.0.attn.to_q.weight"].alpha == 4 + assert index.pairs["transformer_blocks.1.attn.to_k.weight"].alpha == 8 + assert {pair.rank for pair in index.pairs.values()} == {1, 2} + assert reads == ["transformer_blocks.1.attn.to_k.alpha"] + assert not index._resident_cpu + + +@pytest.mark.parametrize("problem", ["missing", "shape", "orphan_alpha", "collision", "unsupported", "alpha", "rank", "factor_dtype"]) +def test_invalid_checkpoint_rejected(tmp_path, setup, problem): + tensors = fixture_tensors() + a = next(iter(tensors)) + if problem == "missing": + tensors.pop(a) + elif problem == "shape": + tensors[a] = torch.ones(1, 4) + elif problem == "orphan_alpha": + tensors["unknown.alpha"] = torch.tensor(1.0) + elif problem == "collision": + tensors["transformer_blocks.0.attn.to_q.lora_down.weight"] = tensors[a].clone() + elif problem == "unsupported": + tensors["unknown"] = torch.tensor(1.0) + elif problem == "alpha": + tensors["transformer_blocks.1.attn.to_k.alpha"] = torch.tensor(float("nan")) + elif problem == "rank": + tensors[a] = torch.ones(2, 3) + elif problem == "factor_dtype": + tensors[a] = tensors[a].long() + with pytest.raises(ValueError): + make_index(tmp_path, setup, tensors) + + +@pytest.mark.parametrize("alpha", [None, 0, -1, float("inf")]) +def test_missing_or_invalid_config_alpha(tmp_path, setup, alpha): + with pytest.raises(ValueError, match="alpha"): + make_index(tmp_path, setup, alpha=alpha) + + +def test_selective_reads_reuse_math_release_and_refiner(tmp_path, setup, monkeypatch): + helpers, modules, Model, streaming, _ = setup + _, pre_module, post_module, transformer_module, *_ = modules + helpers._write_fake_checkpoint(tmp_path, pre_module, post_module, transformer_module) + index = make_index(tmp_path, setup) + config = helpers._config(tmp_path, lora_dynamic_apply=True) + model = Model(str(tmp_path), config, torch.device("cpu"), lora_path=index.path, lora_strength=0.5, lora_alpha=4) + weights = model.transformer_weights + index = weights.streaming_lora + assert len(index.pairs) == 3 + block = weights.streaming_block + base_pointers = {name: weight.weight.data_ptr() for name, weight in index.weights(block).items()} + original = streaming.safe_open + reads = [] + + class Reader: + def __init__(self, *args, **kwargs): + self.context = original(*args, **kwargs) + + def __enter__(self): + self.source = self.context.__enter__() + return self + + def __exit__(self, *args): + return self.context.__exit__(*args) + + def get_tensor(self, key): + reads.append(key) + return self.source.get_tensor(key) + + monkeypatch.setattr(streaming, "safe_open", Reader) + monkeypatch.setattr(torch.Tensor, "pin_memory", lambda *_args, **_kwargs: pytest.fail("streamed LoRA must not pin factors")) + for i, name in [(0, "transformer_blocks.0.attn.to_q.weight"), (1, "transformer_blocks.0.attn.to_k.weight")]: + reads.clear() + assert weights.load_streaming_block(i) is block + linears = index.weights(block) + assert {key for key, value in linears.items() if value.has_lora_branch} == {name} + assert len(reads) == 2 + assert all(f"transformer_blocks.{i}." in key for key in reads) + assert {key: value.weight.data_ptr() for key, value in linears.items()} == base_pointers + linear = linears[name] + x = torch.tensor([[1, 2, 3], [-1, 0, 2]], dtype=torch.bfloat16) + pair = next(iter(index.blocks[i].values())) + factors = fixture_tensors() + expected = x @ linear.weight + 0.5 * (pair.alpha / pair.rank) * ((x @ factors[pair.down_key].T) @ factors[pair.up_key].T) + assert torch.equal(linear.apply(x), expected) + assert not linear.lora_down.is_pinned() + refiner = model.pre_weight + # Tiny fixture's to_cuda is a no-op: explicitly activate its CPU weight. + for linear in index.weights(refiner).values(): + if getattr(linear, "weight", None) is None: + linear.weight = linear.pin_weight + reads.clear() + for _ in range(2): + with index.resident_scope(refiner): + assert sum(weight.has_lora_branch for weight in index.weights(refiner).values()) == 1 + assert not any(weight.has_lora_branch for weight in index.weights(refiner).values()) + assert len(reads) == 2 # Cached on CPU only, read once across evaluations. + assert all(tensor.device.type == "cpu" for tensors in index._resident_cpu.values() for tensor in tensors) + weights.release_disk_streaming_buffer() + assert weights.streaming_block is None + assert not any(weight.has_lora_branch for weight in index.weights(block).values()) + assert all(not hasattr(weight, "lora_down") and not hasattr(weight, "lora_scale") for weight in index.weights(block).values()) + new_block = weights.load_streaming_block(0) + assert new_block is not block + assert sum(weight.has_lora_branch for weight in index.weights(new_block).values()) == 1 + weights.release_disk_streaming_buffer() + + +def test_resident_cleanup_on_failure(tmp_path, setup): + index = make_index(tmp_path, setup) + Linear = setup[-1] + weight = Linear("token_refiner.refiner_blocks.0.attn.to_q.weight", lora_prefix="token_refiner") + weight.weight = torch.ones(3, 2, dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="pre-infer failed"), index.resident_scope(weight): + assert weight.has_lora_branch + raise RuntimeError("pre-infer failed") + assert not weight.has_lora_branch + + +def test_merged_streaming_rejected_before_loading(tmp_path, setup): + helpers, _, Model, _, _ = setup + with pytest.raises(NotImplementedError, match="dynamic LoRA"): + Model(str(tmp_path), helpers._config(tmp_path), torch.device("cpu"), lora_path="unused.safetensors") + + +def test_no_lora_streaming_unchanged(tmp_path, setup): + helpers, modules, Model, _, _ = setup + _, pre, post, transformer, *_ = modules + helpers._write_fake_checkpoint(tmp_path, pre, post, transformer) + model = Model(str(tmp_path), helpers._config(tmp_path), torch.device("cpu")) + assert model.transformer_weights.streaming_lora is None + block = model.transformer_weights.load_streaming_block(1) + assert block is model.transformer_weights.load_streaming_block(0) + model.transformer_weights.release_disk_streaming_buffer() + + +def test_official_raw_adapter_and_lora_compose(tmp_path, setup, monkeypatch): + helpers, _modules, Model, streaming, _ = setup + raw_helpers = load_module("h3_stream_lora_raw_helpers", Path(__file__).with_name("test_checkpoint_adapter.py"), monkeypatch) + _, raw, _ = raw_helpers.write_raw(tmp_path) + tensors = {} + for i in (0, 1): + tensors[f"transformer_blocks.{i}.attn.to_q.lora_A.weight"] = torch.ones(1, 3, dtype=torch.bfloat16) * (i + 1) + tensors[f"transformer_blocks.{i}.attn.to_q.lora_B.weight"] = torch.ones(4, 1, dtype=torch.bfloat16) + path = tmp_path / "raw_lora.safetensors" + save_file(tensors, path) + model = Model(str(tmp_path), helpers._config(tmp_path, lora_dynamic_apply=True), torch.device("cpu"), lora_path=str(path), lora_alpha=4) + weights = model.transformer_weights + assert weights.checkpoint.selected_reader is not None + pointer = None + for i in (0, 1, 0): + block = weights.load_streaming_block(i) + q = streaming.MiniMaxH3StreamingLora.weights(block)["transformer_blocks.0.attn.to_q.weight"] + if pointer is None: + pointer = q.weight.data_ptr() + assert q.weight.data_ptr() == pointer + raw_qkv = raw[f"blocks.{i}.attn.qkv_proj.weight"] + expected_base = torch.vstack((raw_qkv[:2], raw_qkv[6:8])).T + assert torch.equal(q.weight, expected_base) + x = torch.tensor([[1, 0, -1], [1, 2, 3]], dtype=torch.bfloat16) + a, b = tensors[f"transformer_blocks.{i}.attn.to_q.lora_A.weight"], tensors[f"transformer_blocks.{i}.attn.to_q.lora_B.weight"] + assert torch.equal(q.apply(x), x @ expected_base + 4 * ((x @ a.T) @ b.T)) + # FP32 sensitive heads are not silently given BF16 LoRA factors. + shapes = streaming.streaming_target_shapes(weights.checkpoint, block, model.pre_weight, model.post_weight) + assert "proj_in.weight" not in shapes + assert "proj_out.weight" not in shapes + weights.release_disk_streaming_buffer() + + +def test_ordinary_dynamic_loader_and_mmweight_contract(tmp_path, setup, monkeypatch): + _, _, Model, _, Linear = setup + index = make_index(tmp_path, setup) + model = object.__new__(Model) + model.config = {"dit_disk_streaming": False} + model.device = torch.device("cpu") + model.lora_alpha = 4 + model._h3_weight_shapes = {name: (2, 3) for name in index.pairs} + model.use_tp = False + # The unchanged ordinary loader pins CPU tensors; avoid requiring a GPU + # in this regression and verify that it still takes its original path. + pins = [] + monkeypatch.setattr(torch.Tensor, "pin_memory", lambda tensor: pins.append(tensor.shape) or tensor) + loaded = model._load_lora_file(index.path) + assert pins + q = Linear("transformer_blocks.0.attn.to_q.weight", lora_prefix="transformer_blocks") + q.weight = torch.ones(3, 2, dtype=torch.bfloat16) + q.register_lora(loaded, 0.5) + x = torch.ones(2, 3, dtype=torch.bfloat16) + expected = q.apply(x) + q.remove_lora() + index.load_block(q, 0) + assert torch.equal(q.apply(x), expected) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_dynamic_math_non_integer_scaling(tmp_path, setup, dtype): + _, _, _, _, Linear = setup + gen = torch.Generator().manual_seed(123) + tensors = { + "transformer_blocks.0.attn.to_q.lora_down.weight": torch.randn(3, 3, generator=gen, dtype=dtype), + "transformer_blocks.0.attn.to_q.lora_up.weight": torch.randn(2, 3, generator=gen, dtype=dtype), + } + index = make_index(tmp_path, setup, tensors, dtype=dtype, alpha=5, strength=0.7) + weight = Linear("transformer_blocks.0.attn.to_q.weight", lora_prefix="transformer_blocks") + weight.weight = torch.randn(3, 2, generator=gen, dtype=dtype) + x = torch.randn(7, 3, generator=gen, dtype=dtype) + index.load_block(weight, 0) + a, b = tensors.values() + # Preserve the released dynamic path's scalar dtype and operation order. + scale = torch.tensor(5, dtype=dtype) / 3 + expected = x @ weight.weight + 0.7 * scale * ((x @ a.T) @ b.T) + assert torch.equal(weight.apply(x), expected) + index.clear(weight) + + +def test_production_pre_infer_scope_and_error_cleanup(tmp_path, setup): + helpers, modules, Model, _streaming, _ = setup + _, pre, post, transformer, *_ = modules + helpers._write_fake_checkpoint(tmp_path, pre, post, transformer) + index = make_index(tmp_path, setup) + model = Model(str(tmp_path), helpers._config(tmp_path, lora_dynamic_apply=True), torch.device("cpu"), lora_path=index.path, lora_alpha=4) + for weight in index.weights(model.pre_weight).values(): + weight.weight = weight.pin_weight + calls = [] + + def pre_infer(root, prompt): + assert any(weight.has_lora_branch for weight in index.weights(root).values()) + calls.append("pre") + return prompt + + def transformer_infer(weights, pre): + assert not any(weight.has_lora_branch for weight in index.weights(model.pre_weight).values()) + calls.append("transformer") + return pre + + def post_infer(root, hidden, pre): + calls.append("post") + return hidden + + model.pre_infer = types.SimpleNamespace(infer=pre_infer) + model.transformer_infer = types.SimpleNamespace(infer=transformer_infer) + model.post_infer = types.SimpleNamespace(infer=post_infer) + prompt = torch.ones(1, 3, dtype=torch.bfloat16) + inputs = {"text_encoder_output": {"prompt_embeds": prompt}} + assert model._infer_cond_uncond(inputs) is prompt + assert calls == ["pre", "transformer", "post"] + + def fail(root, prompt): + pre_infer(root, prompt) + raise RuntimeError("pre-infer failed") + + model.pre_infer.infer = fail + with pytest.raises(RuntimeError, match="pre-infer failed"): + model._infer_cond_uncond(inputs) + assert not any(weight.has_lora_branch for weight in index.weights(model.pre_weight).values()) + with pytest.raises(NotImplementedError, match="selective LoRA index"): + model._load_lora_file(index.path) + with pytest.raises(NotImplementedError, match="runtime adapter switching"): + model._update_lora(index.path, 1) + model.transformer_weights.release_disk_streaming_buffer() + + +def test_ordinary_merged_lora_contract(tmp_path, setup, monkeypatch): + index = make_index(tmp_path, setup) + adapter_base = types.ModuleType("lightx2v.models.networks.lora_adapter") + adapter_base.LoraAdapter = object + monkeypatch.setitem(sys.modules, adapter_base.__name__, adapter_base) + module = load_module("h3_ordinary_lora_regression", ROOT / "lightx2v/models/networks/minimax_h3/lora.py", monkeypatch) + adapter = module.MiniMaxH3LoraAdapter() + base = {name: torch.ones(2, 3, dtype=torch.bfloat16) for name in index.pairs} + adapter.model = types.SimpleNamespace(config={"lora_merge_device": "cpu"}, use_tp=False, original_weight_dict=base) + assert adapter._merge_file(index.path, strength=0.5, alpha=4) == 3 + fixture = fixture_tensors() + for name, pair in index.pairs.items(): + expected = torch.ones(2, 3, dtype=torch.bfloat16) + expected.add_(fixture[pair.up_key] @ fixture[pair.down_key], alpha=0.5 * pair.alpha / pair.rank) + assert torch.equal(base[name], expected) From 3a65e510ee1300327c4475bbc869c046fbd6acf3 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Mon, 7 Sep 2026 08:17:03 +0800 Subject: [PATCH 21/31] feat(mps): add MiniMax H3 Turbo 4-step config --- .../mps/minimax_h3_t2av_turbo_4step.json | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 configs/platforms/mps/minimax_h3_t2av_turbo_4step.json diff --git a/configs/platforms/mps/minimax_h3_t2av_turbo_4step.json b/configs/platforms/mps/minimax_h3_t2av_turbo_4step.json new file mode 100644 index 000000000..8732e366e --- /dev/null +++ b/configs/platforms/mps/minimax_h3_t2av_turbo_4step.json @@ -0,0 +1,51 @@ +{ + "infer_steps": 5, + "target_video_length": 362, + "target_height": 768, + "target_width": 1344, + "fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "block", + "text_encoder_cpu_offload": true, + "text_encoder_offload_granularity": "block", + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "torch_sdpa", + "rms_type": "torch_native", + "rope_type": "torch_real_rope", + "feature_caching": "NoCaching", + "use_compile": false, + "video_flow_shift": 6.0, + "audio_flow_shift": 3.0, + "h3_step_update": "training_euler", + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true, + "lora_dynamic_apply": true, + "lora_configs": [ + { + "path": "lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", + "strength": 1.0, + "alpha": 128 + } + ], + "dit_prepost_resident": false, + "dit_disk_streaming": true, + "text_encoder_disk_streaming": true, + "text_encoder_host_pinned": false, + "text_encoder_release_block_offload_buffers": true, + "text_encoder_quantized": false, + "vae_use_compile": false, + "vae_attn_type": "torch_sdpa", + "video_vae_quantized": false, + "warmup": false, + "tensor_parallel": false, + "dit_quantized": false, + "dit_quant_scheme": "Default", + "mps_sdpa_query_chunk_size": 128, + "seq_parallel": false +} From f36f5b89b6b165fdf25ff6195092afd2593f30ff Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Mon, 7 Sep 2026 18:09:22 +0800 Subject: [PATCH 22/31] fix(minimax_h3): correct Video VAE QKV weight mapping --- .../video_encoders/hf/minimax_h3/weights.py | 19 +++++++-- .../test_video_vae_checkpoint_adapter.py | 40 ++++++++++++++++++- 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/weights.py b/lightx2v/models/video_encoders/hf/minimax_h3/weights.py index 8c24596c4..b6c8929e6 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/weights.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/weights.py @@ -287,6 +287,18 @@ def validate_minimax_h3_video_vae_checkpoint(module: nn.Module, component_dir: s return SafetensorsSubsetReport(Path(component_dir), files, tuple(sorted(assigned)), ignored) +def _split_video_vae_qkv(tensor: torch.Tensor, num_heads: int, head_dim: int) -> tuple[torch.Tensor, ...]: + """Extract released per-head Q/K/V components for a weight or bias.""" + if tensor.ndim not in (1, 2) or tensor.shape[0] % 3: + raise ValueError("Video VAE fused QKV must be a weight/bias with rows divisible by 3") + if num_heads <= 0 or head_dim <= 0 or tensor.shape[0] != num_heads * 3 * head_dim: + raise ValueError("Video VAE fused QKV rows do not match target num_heads * 3 * head_dim") + # Released rows flatten [num_heads, 3, head_dim, ...], not [3, num_heads, head_dim, ...]. + # Both conventions have identical shapes, so shape-only validation cannot distinguish them. + view = tensor.reshape(num_heads, 3, head_dim, *tensor.shape[1:]) + return tuple(view[:, component].reshape(num_heads * head_dim, *tensor.shape[1:]).contiguous() for component in range(3)) + + def load_minimax_h3_video_vae_checkpoint(module: nn.Module, component_dir: str | Path) -> SafetensorsSubsetReport: """Stream the released official Video VAE schema into the native module.""" report = validate_minimax_h3_video_vae_checkpoint(module, component_dir) @@ -299,9 +311,10 @@ def load_minimax_h3_video_vae_checkpoint(module: nn.Module, component_dir: str | continue tensor_slice = checkpoint.get_slice(source_key) if len(targets) == 3: - chunk_size = tensor_slice.get_shape()[0] // 3 - for index, target in enumerate(targets): - _assign_tensor(module, target, tensor_slice[index * chunk_size : (index + 1) * chunk_size]) + attention, _ = _get_parent(module, targets[0].rsplit(".", 1)[0]) + components = _split_video_vae_qkv(checkpoint.get_tensor(source_key), attention.heads, attention.dim_head) + for target, component in zip(targets, components): + _assign_tensor(module, target, component) loaded.add(target) elif ".ff.w1." in source_key: target = targets[0] diff --git a/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py b/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py index ffd3f6046..f1660bbf8 100644 --- a/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py +++ b/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py @@ -32,6 +32,10 @@ def __getitem__(self, index): parent.add_module(part, _IndexedModule()) parent = getattr(parent, part) parent.register_parameter(parts[-1], nn.Parameter(torch.empty_like(tensor, device="meta"))) + if hasattr(root, "decoder") and hasattr(root.decoder, "transformer_blocks"): + attention = root.decoder.transformer_blocks[0].attn + attention.heads = 2 + attention.dim_head = 2 return root @@ -101,8 +105,8 @@ def test_official_mapping_qkv_ffn_and_mask_token(tmp_path): assert torch.equal(state["decoder.transformer_blocks.0.attn.to_out.0.weight"], tensors["decoder.transformer_blocks.0.attn.to_out.weight"]) assert torch.equal(state["decoder.transformer_blocks.0.ff.net.2.weight"], tensors["decoder.transformer_blocks.0.ff.w2.weight"]) for index, name in enumerate(("q", "k", "v")): - assert torch.equal(state[f"decoder.transformer_blocks.0.attn.to_{name}.weight"], tensors["decoder.transformer_blocks.0.attn.to_qkv.weight"][index * 4 : (index + 1) * 4]) - assert torch.equal(state[f"decoder.transformer_blocks.0.attn.to_{name}.bias"], tensors["decoder.transformer_blocks.0.attn.to_qkv.bias"][index * 4 : (index + 1) * 4]) + assert torch.equal(state[f"decoder.transformer_blocks.0.attn.to_{name}.weight"], tensors["decoder.transformer_blocks.0.attn.to_qkv.weight"].reshape(2, 3, 2, 4)[:, index].reshape(4, 4)) + assert torch.equal(state[f"decoder.transformer_blocks.0.attn.to_{name}.bias"], tensors["decoder.transformer_blocks.0.attn.to_qkv.bias"].reshape(2, 3, 2)[:, index].reshape(4)) assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.weight"][:4] == 2) assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.weight"][4:] == 1) assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.bias"][:4] == 2) @@ -145,3 +149,35 @@ def test_legacy_subset_loader_regression(tmp_path): report = load_safetensors_subset(module, path) assert report.loaded_keys == ("layer.weight",) assert torch.equal(module.state_dict()["layer.weight"], expected["layer.weight"]) + + +@pytest.mark.parametrize("is_weight", [False, True]) +def test_qkv_per_head_components_and_reinterleave(is_weight): + heads, head_dim = 3, 2 + rows = torch.tensor([100 * h + 10 * c + d for h in range(heads) for c in range(3) for d in range(head_dim)]) + source = rows.float() + if is_weight: + source = source[:, None] * 10 + torch.arange(5) + parts = _WEIGHTS._split_video_vae_qkv(source, heads, head_dim) + for c, part in enumerate(parts): + expected = torch.tensor([100 * h + 10 * c + d for h in range(heads) for d in range(head_dim)]).float() + if is_weight: + expected = expected[:, None] * 10 + torch.arange(5) + assert torch.equal(part, expected) + assert not torch.equal(part, source.chunk(3)[c]) + assert part.dtype == source.dtype and part.device == source.device and part.is_contiguous() + rebuilt = torch.stack([p.reshape(heads, head_dim, *source.shape[1:]) for p in parts], dim=1).reshape_as(source) + assert torch.equal(rebuilt, source) + + +@pytest.mark.parametrize("shape,heads,dim", [((11,), 2, 2), ((12,), 3, 2), ((12,), 2, 0), ((12, 2, 2), 2, 2)]) +def test_qkv_invalid_geometry_rejected(shape, heads, dim): + with pytest.raises(ValueError, match="Video VAE fused QKV"): + _WEIGHTS._split_video_vae_qkv(torch.empty(shape), heads, dim) + + +def test_loader_rejects_incompatible_attention_geometry(tmp_path): + module = _parameter_module(_native_specs()) + module.decoder.transformer_blocks[0].attn.heads = 3 + with pytest.raises(ValueError, match="target num_heads"): + load_minimax_h3_video_vae_checkpoint(module, _write(tmp_path / "model.safetensors")) From 64f5194cca51445b4e3b566eec83816aa4748a35 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Mon, 7 Sep 2026 21:06:00 +0800 Subject: [PATCH 23/31] fix(mps): avoid broken temporal padding in H3 VAE --- .../video_encoders/hf/minimax_h3/video_vae.py | 11 ++++- .../minimax_h3/test_video_vae_loader.py | 41 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py index 1c68a0ab5..997e9b000 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py @@ -217,12 +217,21 @@ def __init__(self, in_channels, out_channels, kernel_size, stride=1, spatial_pad self.temporal_padding = temporal_padding self.spatial_padding_mode = spatial_padding_mode + def _pad_temporal(self, hidden_states): + if hidden_states.device.type == "mps": + # pytorch/pytorch#194922: rank-5 MPS constant padding can corrupt data. + # Fixed upstream on main, but supported stable versions may still be affected. + batch, channels, _, height, width = hidden_states.shape + zeros = hidden_states.new_zeros((batch, channels, self.temporal_padding, height, width)) + return torch.cat((zeros, hidden_states), dim=2) + return F.pad(hidden_states, (0, 0, 0, 0, self.temporal_padding, 0)) + def forward(self, hidden_states): if self.spatial_padding > 0: p = self.spatial_padding hidden_states = F.pad(hidden_states, (p, p, p, p, 0, 0), mode=self.spatial_padding_mode) if self.temporal_padding > 0: - hidden_states = F.pad(hidden_states, (0, 0, 0, 0, self.temporal_padding, 0)) + hidden_states = self._pad_temporal(hidden_states) return F.conv3d(hidden_states, self.weight, self.bias, stride=self.stride, dilation=self.dilation) diff --git a/tests/models/minimax_h3/test_video_vae_loader.py b/tests/models/minimax_h3/test_video_vae_loader.py index c4fca54c5..5a7641edb 100644 --- a/tests/models/minimax_h3/test_video_vae_loader.py +++ b/tests/models/minimax_h3/test_video_vae_loader.py @@ -242,3 +242,44 @@ def _prepare_inference_dtypes(self): model = TinyQuantizedVAE.from_pretrained(tmp_path, checkpoint_path=tmp_path / "quant.safetensors", quant_scheme="fp8-sgl", cpu_offload=True) assert model.load_report == "quant-report" assert calls == ["legacy", "pack"] + + +@pytest.mark.parametrize("device", ["cpu", "mps"]) +@pytest.mark.parametrize("dtype_name", ["float32", "float16"]) +@pytest.mark.parametrize("size", [64, 258]) +def test_causal_temporal_padding_preserves_interior(video_vae_module, device, dtype_name, size): + import torch + import torch.nn.functional as F + + if device == "mps" and not torch.backends.mps.is_available(): + pytest.skip("MPS is unavailable") + dtype = getattr(torch, dtype_name) + source = torch.randn((1, 3, 17, size, size), generator=torch.Generator().manual_seed(123)).to(dtype) + conv = video_vae_module.MiniMaxH3VideoCausalConv3d(3, 3, 1, temporal_padding=2) + actual_device = conv._pad_temporal(source.to(device)) + actual = actual_device.cpu() + expected = F.pad(source, (0, 0, 0, 0, 2, 0)) + + assert actual.shape == (1, 3, 19, size, size) + assert actual_device.device.type == device + assert actual.dtype == dtype + assert torch.count_nonzero(actual[:, :, :2]) == 0 + assert torch.equal(actual[:, :, 2:], source) + assert torch.equal(actual, expected) + assert torch.isfinite(actual).all() + + +@pytest.mark.parametrize("device", ["cpu", "mps"]) +def test_causal_conv_forward_uses_temporal_padding(video_vae_module, device): + import torch + import torch.nn.functional as F + + if device == "mps" and not torch.backends.mps.is_available(): + pytest.skip("MPS is unavailable") + conv = video_vae_module.MiniMaxH3VideoCausalConv3d(1, 1, 1, temporal_padding=2).to(device) + with torch.no_grad(): + conv.weight.fill_(1) + conv.bias.zero_() + source = torch.randn((1, 1, 3, 258, 258), generator=torch.Generator().manual_seed(123)) + actual = conv(source.to(device)).cpu() + assert torch.equal(actual, F.pad(source, (0, 0, 0, 0, 2, 0))) From bb918d15c71fac5bcc9480a3034341383489b0dd Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Tue, 8 Sep 2026 16:47:10 +0800 Subject: [PATCH 24/31] fix(minimax_h3): support short clip inference for H3 --- lightx2v/models/networks/minimax_h3/packing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightx2v/models/networks/minimax_h3/packing.py b/lightx2v/models/networks/minimax_h3/packing.py index b0a2dc950..37161b905 100644 --- a/lightx2v/models/networks/minimax_h3/packing.py +++ b/lightx2v/models/networks/minimax_h3/packing.py @@ -26,7 +26,7 @@ MAX_PIXELS = 768 * 1344 MIN_ASPECT_RATIO = 1.0 / 4.0 MAX_ASPECT_RATIO = 4.0 -MIN_DURATION = 5.0 +MIN_DURATION = 22 / FPS # Allow short clips while retaining the 17*n+5 frame alignment. MAX_DURATION = 15.0 PIXEL_MEAN = (0.485, 0.456, 0.406) PIXEL_STD = (0.229, 0.224, 0.225) From 22519573bfd502f03a01a16d9a3bf10c5cc9563e Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Tue, 8 Sep 2026 21:51:14 +0800 Subject: [PATCH 25/31] Avoid redundant tensor clone in MiniMax-H3 streaming loader --- .../models/networks/minimax_h3/checkpoint.py | 7 ++- .../minimax_h3/test_checkpoint_adapter.py | 62 ++++++++++++++++--- 2 files changed, 60 insertions(+), 9 deletions(-) diff --git a/lightx2v/models/networks/minimax_h3/checkpoint.py b/lightx2v/models/networks/minimax_h3/checkpoint.py index 00e3de4a7..7d6d05e47 100644 --- a/lightx2v/models/networks/minimax_h3/checkpoint.py +++ b/lightx2v/models/networks/minimax_h3/checkpoint.py @@ -300,7 +300,12 @@ def write_targets(self, destinations): entry = self.plan.entries[name] with self._source(name) as source: for target, start, out_start, rows in self._ranges(entry, requested): - tile = source[start : start + rows].clone() + # Source slices reference safetensors mmap-backed storage. + # Consume this temporary view immediately while source is alive; + # never cache, return, or retain it for asynchronous use. + # copy_ must finish consuming CPU data before tile is released. + # Async copy would require redesigning source lifetime management. + tile = source[start : start + rows] destination, transpose = destinations[target] if transpose: destination[:, out_start : out_start + rows].copy_(tile.t()) diff --git a/tests/models/minimax_h3/test_checkpoint_adapter.py b/tests/models/minimax_h3/test_checkpoint_adapter.py index f20f41027..b6c3561ab 100644 --- a/tests/models/minimax_h3/test_checkpoint_adapter.py +++ b/tests/models/minimax_h3/test_checkpoint_adapter.py @@ -55,10 +55,17 @@ def raw(tmp_path): return config, tensors, weight_map, plan, C.MiniMaxH3SelectedSourceReader(plan, row_chunk_size=2) -def destinations(plan, names, transpose=False): +@pytest.fixture(params=["cpu", "mps"]) +def destination_device(request): + if request.param == "mps" and not torch.backends.mps.is_available(): + pytest.skip("MPS is not available") + return request.param + + +def destinations(plan, names, transpose=False, device="cpu"): return { name: ( - torch.empty(tuple(reversed(plan.targets[name][1].shape)) if transpose else plan.targets[name][1].shape, dtype=torch.float32 if plan.targets[name][1].dtype == "F32" else torch.bfloat16), + torch.empty(tuple(reversed(plan.targets[name][1].shape)) if transpose else plan.targets[name][1].shape, dtype=torch.float32 if plan.targets[name][1].dtype == "F32" else torch.bfloat16, device=device), transpose, ) for name in names @@ -90,16 +97,16 @@ def test_mixed_format_rejected(raw, tmp_path): @pytest.mark.parametrize("prefix,target", [("blocks.0", "transformer_blocks.0"), ("token_refiner.blocks.0", "token_refiner.refiner_blocks.0")]) -def test_qkv_head_interleave_and_transpose(raw, prefix, target): +def test_qkv_head_interleave_and_transpose(raw, prefix, target, destination_device): _, tensors, _, plan, reader = raw names = [f"{target}.attn.to_{q}.weight" for q in "qkv"] raw_qkv = tensors[f"{prefix}.attn.qkv_proj.weight"] for transpose in (False, True): - dest = destinations(plan, names, transpose) + dest = destinations(plan, names, transpose, device=destination_device) reader.write_targets(dest) for component, name in enumerate(names): expected = torch.vstack([raw_qkv[component * 2 : component * 2 + 2], raw_qkv[6 + component * 2 : 8 + component * 2]]) - actual = dest[name][0].t() if transpose else dest[name][0] + actual = (dest[name][0].t() if transpose else dest[name][0]).cpu() assert actual.shape == (4, 3) assert torch.equal(actual, expected) assert not torch.equal(actual, raw_qkv.chunk(3, dim=0)[component]) @@ -151,18 +158,57 @@ def get_tensor(self, name): @pytest.mark.parametrize("prefix,target", [("blocks.0", "transformer_blocks.0"), ("token_refiner.blocks.0", "token_refiner.refiner_blocks.0")]) -def test_fc1_swap_without_full_cat(raw, monkeypatch, prefix, target): +def test_fc1_swap_without_full_cat(raw, monkeypatch, prefix, target, destination_device): _, tensors, _, plan, reader = raw name = target + ".ff.net.0.proj.weight" source = tensors[prefix + ".mlp.fc1.weight"] - dest = destinations(plan, [name], transpose=True) + dest = destinations(plan, [name], transpose=True, device=destination_device) monkeypatch.setattr(torch, "cat", lambda *a, **k: pytest.fail("no full fused cat")) reader.write_targets(dest) - actual = dest[name][0].t() + actual = dest[name][0].t().cpu() assert torch.equal(actual[:4], source[4:]) assert torch.equal(actual[4:], source[:4]) +@pytest.mark.parametrize("transpose", [False, True]) +@pytest.mark.parametrize("strided", [False, True]) +def test_written_targets_outlive_source_and_own_storage(raw, destination_device, transpose, strided): + _, tensors, _, plan, reader = raw + mapping = { + "proj_in.weight": "video_patch_proj.weight", + "context_embedder.weight": "condition_proj.weight", + "transformer_blocks.0.attn.to_out.0.weight": "blocks.0.attn.out_proj.weight", + "norm_out.norm.weight": "final_layer.norm.weight", + } + dest = {} + expected = {} + for target, source in mapping.items(): + value = tensors[source] + transposed = transpose and value.ndim == 2 + expected[target] = value.t() if transposed else value + if strided: + tensor = torch.empty((*expected[target].shape, 2), dtype=value.dtype, device=destination_device)[..., 0] + else: + tensor = torch.empty_like(value, device=destination_device) + if transposed: + tensor = tensor.t() + dest[target] = (tensor, transposed) + pointers = {name: tensor.data_ptr() for name, (tensor, _) in dest.items()} + reader.write_targets(dest) + # write_targets has closed every source context; only destinations survive. + del reader + for target, (tensor, _) in dest.items(): + assert tensor.data_ptr() == pointers[target] + assert torch.equal(tensor.cpu(), expected[target]) + tensor.zero_() + # Mutating a destination must not modify or alias the checkpoint storage. + reader = C.MiniMaxH3SelectedSourceReader(plan, row_chunk_size=2) + reader.write_targets(dest) + for target, (tensor, _) in dest.items(): + assert tensor.data_ptr() == pointers[target] + assert torch.equal(tensor.cpu(), expected[target]) + + def test_norm_refiner_and_nonblock_mapping_and_fp32(raw): _, tensors, _, plan, reader = raw names = plan.non_block_tensor_names() From 0197365ae07f692ef7d7289c769dc78a6c3c6326 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Thu, 10 Sep 2026 13:42:30 +0800 Subject: [PATCH 26/31] style: format MiniMax-H3 PR files --- .../models/networks/minimax_h3/checkpoint.py | 38 +++++++++---------- .../minimax_h3/weights/transformer_weights.py | 9 +---- .../runners/minimax_h3/minimax_h3_runner.py | 7 +--- .../video_encoders/hf/minimax_h3/weights.py | 4 +- tests/models/minimax_h3/test_checkpoint.py | 6 +-- .../minimax_h3/test_checkpoint_adapter.py | 6 ++- .../minimax_h3/test_query_chunked_sdpa.py | 4 +- .../minimax_h3/test_qwen3vl_disk_streaming.py | 7 +++- .../minimax_h3/test_runner_mps_low_memory.py | 1 + 9 files changed, 35 insertions(+), 47 deletions(-) diff --git a/lightx2v/models/networks/minimax_h3/checkpoint.py b/lightx2v/models/networks/minimax_h3/checkpoint.py index 7d6d05e47..0111cec0a 100644 --- a/lightx2v/models/networks/minimax_h3/checkpoint.py +++ b/lightx2v/models/networks/minimax_h3/checkpoint.py @@ -133,15 +133,9 @@ def __init__(self, checkpoint_dir, config=None): if not weight_map: raise ValueError(f"MiniMax-H3 safetensors index weight_map is empty: {self.index_path}") - invalid_shard_names = sorted( - name - for name, shard_name in weight_map.items() - if not isinstance(shard_name, str) or not shard_name - ) + invalid_shard_names = sorted(name for name, shard_name in weight_map.items() if not isinstance(shard_name, str) or not shard_name) if invalid_shard_names: - raise ValueError( - f"MiniMax-H3 safetensors index contains invalid shard file names for tensors: {invalid_shard_names}" - ) + raise ValueError(f"MiniMax-H3 safetensors index contains invalid shard file names for tensors: {invalid_shard_names}") self.weight_map = dict(weight_map) raw = any(n.startswith(("blocks.", "video_patch_proj.", "audio_patch_proj.", "condition_proj.", "final_layer.", "token_refiner.blocks.")) for n in weight_map) @@ -161,9 +155,21 @@ def __init__(self, checkpoint_dir, config=None): if config is not None: runtime = _native_config(config) keys = ( - "hidden_size", "num_layers", "num_attention_heads", "attention_head_dim", "ffn_dim", "time_embed_dim", - "num_refiner_layers", "freq_dim", "rope_freq_dim", "rope_theta", "time_embed_hidden_dim", "in_channels", - "audio_in_channels", "text_dim", "patch_size", + "hidden_size", + "num_layers", + "num_attention_heads", + "attention_head_dim", + "ffn_dim", + "time_embed_dim", + "num_refiner_layers", + "freq_dim", + "rope_freq_dim", + "rope_theta", + "time_embed_hidden_dim", + "in_channels", + "audio_in_channels", + "text_dim", + "patch_size", ) for key in keys: if key == "rope_theta" and key not in config: @@ -185,15 +191,7 @@ def tensor_names(self): @property def block_indices(self): pattern = re.compile(r"^blocks\.(\d+)\.") if self.format == "official_raw" else _H3_BLOCK_KEY_RE - return tuple( - sorted( - { - int(match.group(1)) - for name in self.weight_map - if (match := pattern.match(name)) is not None - } - ) - ) + return tuple(sorted({int(match.group(1)) for name in self.weight_map if (match := pattern.match(name)) is not None})) def tensor_names_for_block(self, block_index): prefix = "blocks" if self.format == "official_raw" else "transformer_blocks" diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index e3741d255..6d8aa3793 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -171,9 +171,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): self.streaming_lora = None if self.disk_streaming: if config.get("lazy_load", False): - raise NotImplementedError( - "MiniMax-H3 dit_disk_streaming reads the official sharded checkpoint directly and cannot be combined with converted lazy_load block shards." - ) + raise NotImplementedError("MiniMax-H3 dit_disk_streaming reads the official sharded checkpoint directly and cannot be combined with converted lazy_load block shards.") if config.get("dit_quantized", False): raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support quantized DiT checkpoints yet.") if config.get("tensor_parallel", False): @@ -193,10 +191,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): self.num_layers = int(config["num_layers"]) expected_block_indices = tuple(range(self.num_layers)) if self.checkpoint.block_indices != expected_block_indices: - raise ValueError( - "MiniMax-H3 dit_disk_streaming checkpoint block indices mismatch: " - f"expected {expected_block_indices}, found {self.checkpoint.block_indices}" - ) + raise ValueError(f"MiniMax-H3 dit_disk_streaming checkpoint block indices mismatch: expected {expected_block_indices}, found {self.checkpoint.block_indices}") self.blocks = WeightModuleList([]) self.streaming_block = None diff --git a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py index 1d3997613..8653b7467 100644 --- a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py +++ b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py @@ -217,12 +217,7 @@ def load_model(self): self.video_vae, self.audio_vae = self.load_vae() def _is_mps_low_memory_streaming(self): - return ( - AI_DEVICE == "mps" - and self.config.get("task") == "t2av" - and self.config.get("dit_disk_streaming", False) - and self.config.get("text_encoder_disk_streaming", False) - ) + return AI_DEVICE == "mps" and self.config.get("task") == "t2av" and self.config.get("dit_disk_streaming", False) and self.config.get("text_encoder_disk_streaming", False) def _validate_mps_low_memory_streaming_config(self): if not self.config.get("text_encoder_release_block_offload_buffers", False): diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/weights.py b/lightx2v/models/video_encoders/hf/minimax_h3/weights.py index b6c8929e6..30e837ba9 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/weights.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/weights.py @@ -210,9 +210,7 @@ def _official_video_vae_targets(source_key: str) -> tuple[str, ...] | None: return (target,) -def _official_validation_error( - *, unknown: list[str], missing: list[str], duplicates: list[str], shape_mismatches: list[str], dtype_mismatches: list[str] -) -> RuntimeError: +def _official_validation_error(*, unknown: list[str], missing: list[str], duplicates: list[str], shape_mismatches: list[str], dtype_mismatches: list[str]) -> RuntimeError: details = [] for label, values in ( ("unknown", unknown), diff --git a/tests/models/minimax_h3/test_checkpoint.py b/tests/models/minimax_h3/test_checkpoint.py index daa27fa17..d2757dfe7 100644 --- a/tests/models/minimax_h3/test_checkpoint.py +++ b/tests/models/minimax_h3/test_checkpoint.py @@ -31,11 +31,7 @@ def _write_fake_checkpoint(tmp_path): save_file(shard_1, tmp_path / "model-00001-of-00002.safetensors") save_file(shard_2, tmp_path / "model-00002-of-00002.safetensors") - weight_map = { - name: "model-00001-of-00002.safetensors" for name in shard_1 - } | { - name: "model-00002-of-00002.safetensors" for name in shard_2 - } + weight_map = {name: "model-00001-of-00002.safetensors" for name in shard_1} | {name: "model-00002-of-00002.safetensors" for name in shard_2} (tmp_path / "model.safetensors.index.json").write_text( json.dumps({"metadata": {"total_size": 0}, "weight_map": weight_map}), encoding="utf-8", diff --git a/tests/models/minimax_h3/test_checkpoint_adapter.py b/tests/models/minimax_h3/test_checkpoint_adapter.py index b6c3561ab..16ec33540 100644 --- a/tests/models/minimax_h3/test_checkpoint_adapter.py +++ b/tests/models/minimax_h3/test_checkpoint_adapter.py @@ -65,7 +65,11 @@ def destination_device(request): def destinations(plan, names, transpose=False, device="cpu"): return { name: ( - torch.empty(tuple(reversed(plan.targets[name][1].shape)) if transpose else plan.targets[name][1].shape, dtype=torch.float32 if plan.targets[name][1].dtype == "F32" else torch.bfloat16, device=device), + torch.empty( + tuple(reversed(plan.targets[name][1].shape)) if transpose else plan.targets[name][1].shape, + dtype=torch.float32 if plan.targets[name][1].dtype == "F32" else torch.bfloat16, + device=device, + ), transpose, ) for name in names diff --git a/tests/models/minimax_h3/test_query_chunked_sdpa.py b/tests/models/minimax_h3/test_query_chunked_sdpa.py index 4a7e29c60..3d59a4408 100644 --- a/tests/models/minimax_h3/test_query_chunked_sdpa.py +++ b/tests/models/minimax_h3/test_query_chunked_sdpa.py @@ -101,7 +101,5 @@ def test_gqa_fallback_preserved(sdpa): q = torch.randn(9, 4, 8, generator=generator) k, v = [torch.randn(9, 2, 8, generator=generator) for _ in range(2)] actual = sdpa.TorchSDPAWeight().apply(q, k, v, attention_scope="minimax_h3_dit", mps_sdpa_query_chunk_size=128) - expected = F.scaled_dot_product_attention( - q.transpose(0, 1), k.repeat_interleave(2, dim=1).transpose(0, 1), v.repeat_interleave(2, dim=1).transpose(0, 1) - ).transpose(0, 1).reshape(9, 32) + expected = F.scaled_dot_product_attention(q.transpose(0, 1), k.repeat_interleave(2, dim=1).transpose(0, 1), v.repeat_interleave(2, dim=1).transpose(0, 1)).transpose(0, 1).reshape(9, 32) torch.testing.assert_close(actual, expected) diff --git a/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py b/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py index 760924246..e77dcb6b0 100644 --- a/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py +++ b/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py @@ -391,8 +391,11 @@ def test_public_infer_offload_lifecycle(tmp_path, monkeypatch, qwen_module, mode monkeypatch.setattr(qwen_module, "AI_DEVICE", "cpu") monkeypatch.setattr(qwen_module, "MINIMAX_H3_TEXT_HIDDEN_SIZE", 8) backbone = qwen_module._Qwen3VLTextBackboneWeights( - encoder.config, _tiny_text_config(), num_layers=2, - block_offload=encoder.block_offload, disk_streaming=encoder.disk_streaming, + encoder.config, + _tiny_text_config(), + num_layers=2, + block_offload=encoder.block_offload, + disk_streaming=encoder.disk_streaming, ) encoder.text_encoder = backbone encoder.tokenizer = Mock(return_value={"input_ids": [0, 1, 2]}) diff --git a/tests/models/minimax_h3/test_runner_mps_low_memory.py b/tests/models/minimax_h3/test_runner_mps_low_memory.py index a6d9d016a..8bc8f35af 100644 --- a/tests/models/minimax_h3/test_runner_mps_low_memory.py +++ b/tests/models/minimax_h3/test_runner_mps_low_memory.py @@ -87,6 +87,7 @@ def __init__(self, **kwargs): self.__dict__.update(kwargs) _install_module(monkeypatch, "lightx2v.models.video_encoders.hf.ltx2.audio_vae.ops", Audio=Audio) + class _Metrics: def __getattr__(self, _name): return None From 83491d70ef3a4ed82c149a4a510f3d0c5dc7f82f Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Thu, 10 Sep 2026 14:15:38 +0800 Subject: [PATCH 27/31] fix(minimax_h3): support raw checkpoints in AdaLN cache builder --- tests/models/minimax_h3/test_adaln_cache.py | 203 ++++++++++++++++++++ tools/cache_minimax_h3_adaln/builder.py | 26 ++- 2 files changed, 224 insertions(+), 5 deletions(-) create mode 100644 tests/models/minimax_h3/test_adaln_cache.py diff --git a/tests/models/minimax_h3/test_adaln_cache.py b/tests/models/minimax_h3/test_adaln_cache.py new file mode 100644 index 000000000..e9f3b0c81 --- /dev/null +++ b/tests/models/minimax_h3/test_adaln_cache.py @@ -0,0 +1,203 @@ +"""Exercise real cache math and IO with small official/native checkpoints.""" + +import importlib.util +import json +import sys +import types +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file, save_file +from test_checkpoint_adapter import C, write_raw +from test_scheduler_layout import modules as scheduler_modules + +ROOT = Path(__file__).parents[3] + + +@pytest.fixture +def cache_modules(monkeypatch): + scheduler_modules.__wrapped__(monkeypatch) + envs = types.ModuleType("lightx2v.utils.envs") + envs.GET_DTYPE = lambda: torch.bfloat16 + monkeypatch.setitem(sys.modules, envs.__name__, envs) + + def load(name, path): + spec = importlib.util.spec_from_file_location(name, ROOT / path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, name, module) + spec.loader.exec_module(module) + return module + + monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.checkpoint", C) + load("lightx2v.models.networks.minimax_h3.infer.module_io", "lightx2v/models/networks/minimax_h3/infer/module_io.py") + load("lightx2v.models.networks.minimax_h3.infer.pre_infer", "lightx2v/models/networks/minimax_h3/infer/pre_infer.py") + load("lightx2v.models.networks.minimax_h3.adaln_cache_guide", "lightx2v/models/networks/minimax_h3/adaln_cache_guide.py") + cache = load("lightx2v.models.networks.minimax_h3.adaln_cache", "lightx2v/models/networks/minimax_h3/adaln_cache.py") + builder = load("h3_adaln_builder_under_test", "tools/cache_minimax_h3_adaln/builder.py") + return cache, builder + + +@pytest.fixture +def checkpoint(tmp_path): + raw = tmp_path / "raw" + raw.mkdir() + config, tensors, weight_map = write_raw(raw) + for name in tensors: + if name != "rope.inv_freq": + tensors[name] = tensors[name] * 0.001 + for shard in set(weight_map.values()): + save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, raw / shard) + config = C._native_config(config) + config.update(dit_original_ckpt=str(raw), use_adaln_cache=True, adaln_cache_dir=str(tmp_path / "cache"), task="fl2av", infer_steps=4) + return config, tensors, weight_map + + +def native_projections(tensors): + # Independent expected mapping for the cache's four projection families. + pairs = [("time_embedder.linear_1", "time_embedder.proj_in"), ("time_embedder.linear_2", "time_embedder.proj_out"), ("norm_out.linear", "final_layer.adaln_proj.linear")] + pairs.extend((f"transformer_blocks.{i}.adaln_proj.linear", f"blocks.{i}.adaln_proj.linear") for i in range(2)) + return {target + suffix: tensors[source + suffix] for target, source in pairs for suffix in (".weight", ".bias")} + + +@pytest.mark.parametrize("device", ["cpu", "mps"]) +def test_raw_generation_matches_native_and_loads(cache_modules, checkpoint, tmp_path, monkeypatch, device): + if device == "mps" and not torch.backends.mps.is_available(): + pytest.skip("MPS unavailable") + cache, builder = cache_modules + monkeypatch.setattr(builder, "AI_DEVICE", device) + monkeypatch.setattr(builder, "torch_device_module", getattr(torch, device)) + config, tensors, _ = checkpoint + expected = native_projections(tensors) + reader = builder._CheckpointTensors(builder._checkpoint_files(config), config=config) + for name, tensor in expected.items(): + actual = reader.get(name) + assert actual.dtype == tensor.dtype + assert torch.equal(actual, tensor) + with pytest.raises(KeyError, match="missing"): + reader.get("missing.weight") + raw_path = builder.build_persistent_adaln_cache(config) + raw_tables = load_file(raw_path / "adaln_cache.safetensors") + spec = cache._build_spec(config) + assert cache._validate_cache(raw_path, spec) + tables, norm = cache.load_persistent_adaln_cache(config, device) + for entry in spec["entries"]: + key = tuple(cache._timesteps_from_bits(entry["timestep_bits"]).tolist()) + for i in range(spec["num_layers"]): + tensor = tables[key][i] + assert tensor.device.type == device and tensor.dtype == torch.bfloat16 + assert tuple(tensor.shape) == cache._expected_table_shape(spec, entry) + assert torch.equal(tensor.cpu(), raw_tables[cache._block_key(i, entry)]) + assert norm[key].dtype == torch.bfloat16 and norm[key].device.type == device + assert tuple(norm[key].shape) == cache._expected_norm_out_shape(spec, entry) + assert torch.equal(norm[key].cpu(), raw_tables[cache._norm_out_key(entry)]) + native = tmp_path / "native" + native.mkdir() + save_file(expected, native / "model.safetensors") + for indexed in (False, True): + if indexed: + (native / "model.safetensors.index.json").write_text(json.dumps({"weight_map": {k: "model.safetensors" for k in expected}})) + native_config = dict(config, dit_original_ckpt=str(native), adaln_cache_dir=str(tmp_path / f"native-cache-{indexed}")) + native_path = builder.build_persistent_adaln_cache(native_config) + native_tables = load_file(native_path / "adaln_cache.safetensors") + assert raw_tables.keys() == native_tables.keys() + assert all(torch.equal(value, native_tables[name]) for name, value in raw_tables.items()) + with pytest.raises(FileExistsError): + builder.build_persistent_adaln_cache(config) + # Loader still rejects corrupted cache tensors. + key = next(iter(raw_tables)) + raw_tables[key] = raw_tables[key].float() + save_file(raw_tables, raw_path / "adaln_cache.safetensors") + with pytest.raises(FileNotFoundError, match="AdaLN cache not found"): + cache.load_persistent_adaln_cache(config, device) + + +@pytest.mark.parametrize("damage", ["shape", "dtype", "missing_key", "missing_shard", "runtime_config", "rope"]) +def test_invalid_raw_checkpoint_rejected(cache_modules, checkpoint, damage): + _, builder = cache_modules + config, tensors, weight_map = checkpoint + directory = Path(config["dit_original_ckpt"]) + name = "time_embedder.proj_in.weight" + if damage == "runtime_config": + config["hidden_size"] += 1 + elif damage == "missing_shard": + (directory / weight_map[name]).unlink() + elif damage == "missing_key": + del weight_map[name] + (directory / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + else: + if damage == "shape": + tensors[name] = tensors[name][:1] + elif damage == "dtype": + tensors[name] = tensors[name].to(torch.bfloat16) + else: + name = "rope.inv_freq" + tensors[name] = torch.zeros_like(tensors[name]) + shard = weight_map[name] + save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, directory / shard) + with pytest.raises((ValueError, FileNotFoundError)): + builder.build_persistent_adaln_cache(config) + assert not any(Path(config["adaln_cache_dir"]).rglob("manifest.json")) + assert not any(Path(config["adaln_cache_dir"]).rglob(".building-*")) + + +@pytest.mark.parametrize("task", ["t2av", "i2av", "l2av", "fl2av"]) +def test_base_profiles_share_raw_cache(cache_modules, checkpoint, task): + cache, builder = cache_modules + config, _, _ = checkpoint + path = builder.build_persistent_adaln_cache(config) + requested = dict(config, task=task) + assert cache._cache_path(requested) == path + assert cache._validate_cache(path, cache._build_spec(requested)) + assert cache.load_persistent_adaln_cache(requested, "cpu") + ref_config = dict(config, task="ref2av") + assert cache._cache_path(ref_config) != path + assert not cache._validate_cache(path, cache._build_spec(ref_config)) + with pytest.raises(FileNotFoundError): + cache.load_persistent_adaln_cache(ref_config, "cpu") + + +@pytest.mark.parametrize("damage", ["manifest", "profile", "shape", "missing_tensor"]) +def test_cache_compatibility_validation(cache_modules, checkpoint, damage): + cache, builder = cache_modules + config, _, _ = checkpoint + path = builder.build_persistent_adaln_cache(config) + manifest_path = path / "manifest.json" + if damage in {"manifest", "profile"}: + manifest = json.loads(manifest_path.read_text()) + if damage == "manifest": + manifest["video_flow_shift"] += 1 + else: + manifest["entries"][0]["name"] = "ref2av_video_step_000" + manifest_path.write_text(json.dumps(manifest)) + else: + tensor_path = path / "adaln_cache.safetensors" + tensors = load_file(tensor_path) + name = next(iter(tensors)) + if damage == "shape": + tensors[name] = tensors[name][:1].contiguous() + else: + del tensors[name] + save_file(tensors, tensor_path) + assert not cache._validate_cache(path, cache._build_spec(config)) + with pytest.raises(FileNotFoundError, match="AdaLN cache not found"): + cache.load_persistent_adaln_cache(config, "cpu") + + +@pytest.mark.parametrize( + "update, error", + [ + ({"use_adaln_cache": False}, ValueError), + ({"adaln_cache_dir": ""}, ValueError), + ({"dummy_model": True}, NotImplementedError), + ({"task": "unsupported"}, ValueError), + ], +) +def test_invalid_cache_config(cache_modules, checkpoint, update, error): + cache, builder = cache_modules + config, _, _ = checkpoint + config.update(update) + with pytest.raises(error): + builder.build_persistent_adaln_cache(config) + with pytest.raises(error): + cache.load_persistent_adaln_cache(config, "cpu") diff --git a/tools/cache_minimax_h3_adaln/builder.py b/tools/cache_minimax_h3_adaln/builder.py index 5aef69535..639e24528 100644 --- a/tools/cache_minimax_h3_adaln/builder.py +++ b/tools/cache_minimax_h3_adaln/builder.py @@ -30,6 +30,7 @@ _timesteps_from_bits, _validate_cache, ) +from lightx2v.models.networks.minimax_h3.checkpoint import MiniMaxH3ShardCheckpoint from lightx2v.models.networks.minimax_h3.infer.pre_infer import timestep_embedding from lightx2v_platform.base.global_var import AI_DEVICE @@ -47,9 +48,13 @@ def _checkpoint_files(config) -> list[Path]: class _CheckpointTensors: """Read individual tensors without materializing the whole checkpoint.""" - def __init__(self, files: list[Path]): + def __init__(self, files: list[Path], config=None): self.files = files - self.locations = self._find_locations() + self.adapter = None + if (files[0].parent / "model.safetensors.index.json").is_file(): + checkpoint = MiniMaxH3ShardCheckpoint(files[0].parent, config=config) + self.adapter = checkpoint.selected_reader + self.locations = self._find_locations() if self.adapter is None else {} def _find_locations(self) -> dict[str, Path]: directory = self.files[0].parent @@ -66,6 +71,17 @@ def _find_locations(self) -> dict[str, Path]: return locations def get(self, name: str) -> torch.Tensor: + if self.adapter is not None: + # Reuse the inference mapping and validation. The builder consumes + # logical (out, in) weights; _linear owns the eventual transpose. + try: + spec = self.adapter.plan.targets[name][1] + except KeyError as error: + raise KeyError(f"MiniMax-H3 checkpoint tensor is missing: {name}") from error + dtype = torch.float32 if spec.dtype == "F32" else torch.bfloat16 + tensor = torch.empty(spec.shape, dtype=dtype, device="cpu") + self.adapter.write_targets({name: (tensor, False)}) + return tensor path = self.locations.get(name) if path is None: raise KeyError(f"MiniMax-H3 checkpoint tensor is missing: {name}") @@ -94,10 +110,10 @@ def _empty_device_cache() -> None: torch_device_module.empty_cache() -def _build_cache(spec: dict, cache_path: Path, checkpoint_files: list[Path]) -> None: +def _build_cache(spec: dict, cache_path: Path, checkpoint_files: list[Path], config=None) -> None: stage_path = Path(tempfile.mkdtemp(prefix=".building-", dir=cache_path.parent)) try: - checkpoint = _CheckpointTensors(checkpoint_files) + checkpoint = _CheckpointTensors(checkpoint_files, config=config) with torch.inference_mode(): # ADALN CACHE SYNC: Keep activation placement and casts aligned with # the three online infer modules named in this file's contract. @@ -168,7 +184,7 @@ def build_persistent_adaln_cache(config) -> Path: raise FileExistsError(f"MiniMax-H3 AdaLN cache path already exists: {cache_path}") checkpoint_files = _checkpoint_files(config) logger.info("Building MiniMax-H3 AdaLN cache on {}: {}", AI_DEVICE, cache_path) - _build_cache(spec, cache_path, checkpoint_files) + _build_cache(spec, cache_path, checkpoint_files, config=config) if not _validate_cache(cache_path, spec): raise RuntimeError(f"MiniMax-H3 AdaLN cache validation failed: {cache_path}") From 31ac2271c6a8e29993e957dd62bc8c1e42da6561 Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Thu, 10 Sep 2026 14:58:34 +0800 Subject: [PATCH 28/31] perf(minimax_h3): optimize streaming chunk size by target bytes --- .../models/networks/minimax_h3/checkpoint.py | 31 +++++-- .../minimax_h3/test_checkpoint_adapter.py | 84 +++++++++++++++++++ 2 files changed, 109 insertions(+), 6 deletions(-) diff --git a/lightx2v/models/networks/minimax_h3/checkpoint.py b/lightx2v/models/networks/minimax_h3/checkpoint.py index 0111cec0a..fa96a4ed1 100644 --- a/lightx2v/models/networks/minimax_h3/checkpoint.py +++ b/lightx2v/models/networks/minimax_h3/checkpoint.py @@ -8,6 +8,8 @@ import torch from safetensors import safe_open +TARGET_CHUNK_BYTES = int(5.25 * 1024 * 1024) + _H3_BLOCK_KEY_RE = re.compile(r"^transformer_blocks\.(\d+)\.") @@ -221,13 +223,20 @@ def shards_for_sources(self, names): class MiniMaxH3SelectedSourceReader: """Bounded CPU slice staging, with no large mmap views retained across sources.""" - def __init__(self, plan, row_chunk_size=128): + def __init__(self, plan, row_chunk_size=None, *, target_chunk_bytes=TARGET_CHUNK_BYTES): + """Prefer explicit row counts, then a byte budget, then legacy 128-row chunks. + + Set target_chunk_bytes=None to disable byte-based sizing. QKV head + boundaries and one-dimensional tensors keep their dedicated slicing. + """ if plan.format != "official_raw": raise ValueError("MiniMax-H3 selected adapter requires official raw format") - if row_chunk_size < 1: - raise ValueError("row_chunk_size must be positive") + for name, value in (("row_chunk_size", row_chunk_size), ("target_chunk_bytes", target_chunk_bytes)): + if value is not None and (not isinstance(value, int) or isinstance(value, bool) or value < 1): + raise ValueError(f"{name} must be a positive integer") self.plan = plan self.row_chunk_size = row_chunk_size + self.target_chunk_bytes = target_chunk_bytes @contextmanager def _source(self, name): @@ -262,6 +271,15 @@ def validate_sources(self, names): if not torch.equal(actual.view(torch.int32), expected.view(torch.int32)): raise ValueError("MiniMax-H3 rope.inv_freq differs from native reconstruction") + def _row_chunk_size(self, entry): + if self.row_chunk_size is not None: + return self.row_chunk_size + if self.target_chunk_bytes is not None and len(entry.shape) == 2 and entry.source_name.endswith(".weight"): + row_bytes = entry.shape[1] * {"F32": 4, "BF16": 2}[entry.dtype] + # The budget is a target: at least one complete row must fit a slice. + return max(1, self.target_chunk_bytes // row_bytes) + return 128 + def _ranges(self, entry, requested): if entry.transform == "qkv_head_interleaved": dim = self.plan.config["attention_head_dim"] @@ -271,12 +289,13 @@ def _ranges(self, entry, requested): yield target.name, (head * 3 + component) * dim, head * dim, dim elif entry.transform == "swap_gate_value": half = entry.shape[0] // 2 + step = self._row_chunk_size(entry) for source_start, target_start in ((half, 0), (0, half)): - for row in range(0, half, self.row_chunk_size): - yield entry.targets[0].name, source_start + row, target_start + row, min(self.row_chunk_size, half - row) + for row in range(0, half, step): + yield entry.targets[0].name, source_start + row, target_start + row, min(step, half - row) else: size = entry.shape[0] - step = size if len(entry.shape) == 1 else self.row_chunk_size + step = size if len(entry.shape) == 1 else self._row_chunk_size(entry) for row in range(0, size, step): yield entry.targets[0].name, row, row, min(step, size - row) diff --git a/tests/models/minimax_h3/test_checkpoint_adapter.py b/tests/models/minimax_h3/test_checkpoint_adapter.py index 16ec33540..1d4dbf22b 100644 --- a/tests/models/minimax_h3/test_checkpoint_adapter.py +++ b/tests/models/minimax_h3/test_checkpoint_adapter.py @@ -76,6 +76,90 @@ def destinations(plan, names, transpose=False, device="cpu"): } +@pytest.mark.parametrize("dtype,expected", [("BF16", [2, 2, 1]), ("F32", [1, 1, 1, 1, 1])]) +def test_target_bytes_row_size_rounding_and_tail(raw, dtype, expected): + plan = raw[3] + entry = plan.entries["condition_proj.weight"]._replace(shape=(5, 3), dtype=dtype) + reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=13) + ranges = list(reader._ranges(entry, {entry.targets[0].name})) + assert [rows for _, _, _, rows in ranges] == expected + assert [(start, out) for _, start, out, _ in ranges] == [(sum(expected[:i]), sum(expected[:i])) for i in range(len(expected))] + + +@pytest.mark.parametrize("kwargs,expected", [({}, 1024), ({"target_chunk_bytes": 10752}, 2), ({"row_chunk_size": 3, "target_chunk_bytes": 1}, 3), ({"target_chunk_bytes": None}, 128)]) +def test_chunk_policy_precedence_and_legacy_fallback(raw, kwargs, expected): + plan = raw[3] + entry = plan.entries["condition_proj.weight"]._replace(shape=(2049, 2688)) + reader = C.MiniMaxH3SelectedSourceReader(plan, **kwargs) + ranges = list(reader._ranges(entry, {entry.targets[0].name})) + assert ranges[0][3] == expected + assert sum(rows for _, _, _, rows in ranges) == 2049 + if not kwargs: + fc1 = plan.entries["blocks.0.mlp.fc1.weight"]._replace(shape=(28672, 5376)) + assert next(reader._ranges(fc1, {fc1.targets[0].name}))[3] == 512 + + +def test_target_smaller_than_one_row(raw): + plan = raw[3] + entry = plan.entries["condition_proj.weight"] + reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=1) + assert [rows for _, _, _, rows in reader._ranges(entry, {entry.targets[0].name})] == [1, 1, 1] + + +@pytest.mark.parametrize("prefix", ["blocks.0", "token_refiner.blocks.0"]) +@pytest.mark.parametrize("budget,expected", [(18, [(4, 0, 3), (7, 3, 1), (0, 4, 3), (3, 7, 1)]), (1024, [(4, 0, 4), (0, 4, 4)])]) +def test_target_bytes_keeps_gate_value_halves_independent(raw, prefix, budget, expected): + plan = raw[3] + entry = plan.entries[prefix + ".mlp.fc1.weight"] + reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=budget) + assert [r[1:] for r in reader._ranges(entry, {entry.targets[0].name})] == expected + + +@pytest.mark.parametrize("prefix", ["blocks.0", "token_refiner.blocks.0"]) +def test_qkv_and_vectors_bypass_byte_policy(raw, monkeypatch, prefix): + plan = raw[3] + reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=1) + monkeypatch.setattr(reader, "_row_chunk_size", lambda *a: pytest.fail("dedicated slicing must bypass chunk policy")) + entry = plan.entries[prefix + ".attn.qkv_proj.weight"] + query = entry.targets[0].name + assert list(reader._ranges(entry, {query})) == [(query, 0, 0, 2), (query, 6, 2, 2)] + for name in ["video_patch_proj.bias", "blocks.0.norm1.weight", "token_refiner.final_norm.weight"]: + entry = plan.entries[name] + target = entry.targets[0].name + assert list(reader._ranges(entry, {target})) == [(target, 0, 0, entry.shape[0])] + + +@pytest.mark.parametrize("name", ["row_chunk_size", "target_chunk_bytes"]) +@pytest.mark.parametrize("value", [0, -1, 1.5, True]) +def test_invalid_chunk_policy_rejected(raw, name, value): + with pytest.raises(ValueError, match=name + " must be a positive integer"): + C.MiniMaxH3SelectedSourceReader(raw[3], **{name: value}) + + +@pytest.mark.parametrize("budget", [1, 13, C.TARGET_CHUNK_BYTES, None]) +@pytest.mark.parametrize("transpose", [False, True]) +def test_target_bytes_full_adapter_cpu_parity(raw, budget, transpose): + _, tensors, _, plan, _ = raw + reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=budget) + dest, expected = {}, {} + for entry in plan.entries.values(): + source = tensors[entry.source_name] + for component, target in enumerate(entry.targets): + if entry.transform == "qkv_head_interleaved": + value = source.reshape(2, 3, 2, 3)[:, component].reshape(4, 3) + elif entry.transform == "swap_gate_value": + value = torch.cat((source[4:], source[:4])) + else: + value = source + transposed = transpose and value.ndim == 2 + expected[target.name] = value.t() if transposed else value + dest[target.name] = (torch.empty_like(expected[target.name]), transposed) + reader.write_targets(dest) + for name, (tensor, _) in dest.items(): + assert tensor.dtype == expected[name].dtype + assert torch.equal(tensor, expected[name]), name + + def test_official_detection_and_complete_config_driven_plan(raw): _, tensors, _, plan, _ = raw assert plan.format == "official_raw" From 79cb146ad9382ef38be526c27c4cef6b29397e8b Mon Sep 17 00:00:00 2001 From: q6y6y6 <2282974298@qq.com> Date: Thu, 10 Sep 2026 15:50:05 +0800 Subject: [PATCH 29/31] chore(minimax_h3): remove default Turbo LoRA config --- .../mps/minimax_h3_t2av_turbo_4step.json | 51 ------------------- 1 file changed, 51 deletions(-) delete mode 100644 configs/platforms/mps/minimax_h3_t2av_turbo_4step.json diff --git a/configs/platforms/mps/minimax_h3_t2av_turbo_4step.json b/configs/platforms/mps/minimax_h3_t2av_turbo_4step.json deleted file mode 100644 index 8732e366e..000000000 --- a/configs/platforms/mps/minimax_h3_t2av_turbo_4step.json +++ /dev/null @@ -1,51 +0,0 @@ -{ - "infer_steps": 5, - "target_video_length": 362, - "target_height": 768, - "target_width": 1344, - "fps": 24, - "enable_cfg": false, - "cpu_offload": true, - "offload_granularity": "block", - "text_encoder_cpu_offload": true, - "text_encoder_offload_granularity": "block", - "vae_cpu_offload": true, - "lazy_load": false, - "unload_modules": false, - "attn_type": "torch_sdpa", - "rms_type": "torch_native", - "rope_type": "torch_real_rope", - "feature_caching": "NoCaching", - "use_compile": false, - "video_flow_shift": 6.0, - "audio_flow_shift": 3.0, - "h3_step_update": "training_euler", - "vae_spatial_scale_factor": 16, - "audio_sampling_rate": 32000, - "audio_latents_per_second": 40, - "audio_channels": 2, - "keep_latents_dtype_in_scheduler": true, - "lora_dynamic_apply": true, - "lora_configs": [ - { - "path": "lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors", - "strength": 1.0, - "alpha": 128 - } - ], - "dit_prepost_resident": false, - "dit_disk_streaming": true, - "text_encoder_disk_streaming": true, - "text_encoder_host_pinned": false, - "text_encoder_release_block_offload_buffers": true, - "text_encoder_quantized": false, - "vae_use_compile": false, - "vae_attn_type": "torch_sdpa", - "video_vae_quantized": false, - "warmup": false, - "tensor_parallel": false, - "dit_quantized": false, - "dit_quant_scheme": "Default", - "mps_sdpa_query_chunk_size": 128, - "seq_parallel": false -} From 62e29b5dba9ab4f25784a90b00abc3fb35b8008b Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Thu, 10 Sep 2026 18:42:25 +0800 Subject: [PATCH 30/31] refactor(minimax_h3)!: align checkpoint loading with upstream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 统一使用 Diffusers 权重布局，移除分支中的 raw 格式适配，保留 MPS 低内存推理所需的按需加载。 - 删除 raw DiT/VAE 参数映射、QKV/FFN 转换及配置别名兼容 - 沿用上游 safetensors 文件发现方式，移除固定索引名和参数名前缀白名单 - 将 AdaLN 缓存构建器恢复为 ModelTC 上游实现 - 在 MPS 清理缓存前同步 GPU，调整 VAE 按阶段加载与释放 - 更新模型路径、采样步数和缓存配置，新增 512×512、22 帧、4 步配置及使用说明 - 删除 tests/models/minimax_h3 下的测试文件 BREAKING CHANGE: MiniMax-H3 不再支持 raw 权重布局，请使用 Diffusers 版本。 --- configs/platforms/mps/minimax_h3_t2av.json | 4 +- .../mps/minimax_h3_t2av_4step_512_22.json | 52 ++ .../models/networks/minimax_h3/checkpoint.py | 380 +------------- lightx2v/models/networks/minimax_h3/model.py | 8 +- .../networks/minimax_h3/streaming_lora.py | 26 +- .../minimax_h3/weights/transformer_weights.py | 26 +- .../runners/minimax_h3/minimax_h3_runner.py | 40 +- .../video_encoders/hf/minimax_h3/video_vae.py | 109 +--- .../video_encoders/hf/minimax_h3/weights.py | 175 ------- scripts/platforms/mps/README.md | 40 ++ scripts/platforms/mps/run_minimax_h3_t2av.sh | 57 +-- tests/models/minimax_h3/test_adaln_cache.py | 203 -------- .../minimax_h3/test_audio_vae_offload.py | 101 ---- tests/models/minimax_h3/test_checkpoint.py | 117 ----- .../minimax_h3/test_checkpoint_adapter.py | 468 ------------------ .../minimax_h3/test_model_disk_streaming.py | 428 ---------------- .../minimax_h3/test_mps_low_memory_config.py | 98 ---- .../minimax_h3/test_query_chunked_sdpa.py | 105 ---- .../minimax_h3/test_qwen3vl_disk_streaming.py | 461 ----------------- .../models/minimax_h3/test_rope_precision.py | 65 --- .../minimax_h3/test_runner_mps_low_memory.py | 326 ------------ .../minimax_h3/test_scheduler_layout.py | 69 --- .../models/minimax_h3/test_streaming_lora.py | 383 -------------- .../test_transformer_disk_streaming.py | 419 ---------------- .../test_video_vae_checkpoint_adapter.py | 183 ------- .../minimax_h3/test_video_vae_loader.py | 295 ----------- tools/cache_minimax_h3_adaln/builder.py | 26 +- .../run_cache_minimax_h3_adaln.sh | 28 +- 28 files changed, 208 insertions(+), 4484 deletions(-) create mode 100644 configs/platforms/mps/minimax_h3_t2av_4step_512_22.json create mode 100644 scripts/platforms/mps/README.md delete mode 100644 tests/models/minimax_h3/test_adaln_cache.py delete mode 100644 tests/models/minimax_h3/test_audio_vae_offload.py delete mode 100644 tests/models/minimax_h3/test_checkpoint.py delete mode 100644 tests/models/minimax_h3/test_checkpoint_adapter.py delete mode 100644 tests/models/minimax_h3/test_model_disk_streaming.py delete mode 100644 tests/models/minimax_h3/test_mps_low_memory_config.py delete mode 100644 tests/models/minimax_h3/test_query_chunked_sdpa.py delete mode 100644 tests/models/minimax_h3/test_qwen3vl_disk_streaming.py delete mode 100644 tests/models/minimax_h3/test_rope_precision.py delete mode 100644 tests/models/minimax_h3/test_runner_mps_low_memory.py delete mode 100644 tests/models/minimax_h3/test_scheduler_layout.py delete mode 100644 tests/models/minimax_h3/test_streaming_lora.py delete mode 100644 tests/models/minimax_h3/test_transformer_disk_streaming.py delete mode 100644 tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py delete mode 100644 tests/models/minimax_h3/test_video_vae_loader.py diff --git a/configs/platforms/mps/minimax_h3_t2av.json b/configs/platforms/mps/minimax_h3_t2av.json index c717d4071..8a8d488b6 100644 --- a/configs/platforms/mps/minimax_h3_t2av.json +++ b/configs/platforms/mps/minimax_h3_t2av.json @@ -1,5 +1,5 @@ { - "infer_steps": 30, + "infer_steps": 29, "target_video_length": 124, "target_height": 480, "target_width": 480, @@ -11,6 +11,8 @@ "offload_granularity": "block", "dit_prepost_resident": false, "dit_disk_streaming": true, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln/diffusers", "text_encoder_cpu_offload": true, "text_encoder_offload_granularity": "block", diff --git a/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json b/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json new file mode 100644 index 000000000..f953292f1 --- /dev/null +++ b/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json @@ -0,0 +1,52 @@ +{ + "infer_steps": 4, + "target_video_length": 22, + "target_height": 512, + "target_width": 512, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + + "cpu_offload": true, + "offload_granularity": "block", + "dit_prepost_resident": false, + "dit_disk_streaming": true, + "use_adaln_cache": true, + "adaln_cache_dir": "~/.cache/lightx2v/adaln/diffusers", + + "text_encoder_cpu_offload": true, + "text_encoder_offload_granularity": "block", + "text_encoder_disk_streaming": true, + "text_encoder_host_pinned": false, + "text_encoder_release_block_offload_buffers": true, + "text_encoder_quantized": false, + + "vae_cpu_offload": true, + "vae_use_compile": false, + "vae_attn_type": "torch_sdpa", + "video_vae_quantized": false, + + "lazy_load": false, + "unload_modules": false, + "warmup": false, + + "attn_type": "torch_sdpa", + "mps_sdpa_query_chunk_size": 512, + "rms_type": "torch_native", + "rope_type": "torch_real_rope", + + "tensor_parallel": false, + "dit_quantized": false, + "dit_quant_scheme": "Default", + + "feature_caching": "NoCaching", + "use_compile": false, + + "video_flow_shift": 12.0, + "audio_flow_shift": 3.0, + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true +} diff --git a/lightx2v/models/networks/minimax_h3/checkpoint.py b/lightx2v/models/networks/minimax_h3/checkpoint.py index fa96a4ed1..6e764d626 100644 --- a/lightx2v/models/networks/minimax_h3/checkpoint.py +++ b/lightx2v/models/networks/minimax_h3/checkpoint.py @@ -1,190 +1,29 @@ -import json +"""Load MiniMax-H3 diffusers checkpoints by requested tensor or block.""" + import re from collections import defaultdict -from contextlib import contextmanager from pathlib import Path -from typing import NamedTuple -import torch from safetensors import safe_open -TARGET_CHUNK_BYTES = int(5.25 * 1024 * 1024) - _H3_BLOCK_KEY_RE = re.compile(r"^transformer_blocks\.(\d+)\.") -class TargetSpec(NamedTuple): - name: str - shape: tuple - dtype: str - - -class SourcePlan(NamedTuple): - source_name: str - source_shard: str - transform: str - shape: tuple - dtype: str - targets: tuple - - -def validate_mapping(entries, source_names, target_names): - """Require an exact source classification and exactly one producer per target.""" - sources, targets = [], [] - for entry in entries: - sources.append(entry.source_name) - targets.extend(target.name for target in entry.targets) - for label, actual, expected in (("source", sources, set(source_names)), ("target", targets, set(target_names))): - if len(actual) != len(set(actual)): - raise ValueError(f"MiniMax-H3 duplicate {label} coverage") - missing, unknown = sorted(expected - set(actual)), sorted(set(actual) - expected) - if missing or unknown: - raise ValueError(f"MiniMax-H3 {label} mapping mismatch: missing={missing}, unknown={unknown}") - - -def _official_schema(config): - """Logical checkpoint shapes only; runtime MM transposes belong to the writer.""" - h = config["hidden_size"] - if config.get("adaln_out_features", 18 * h) != 18 * h or config.get("final_adaln_out_features", 2 * h) != 2 * h: - raise ValueError("MiniMax-H3 incompatible AdaLN output dimensions") - inner = config["num_attention_heads"] * config["attention_head_dim"] - ffn, time = config["ffn_dim"], config["time_embed_dim"] - patch = config["in_channels"] - for axis in config["patch_size"]: - patch *= axis - schema = {} - - def add(source, target, shape, dtype="BF16", transform="rename", target_shape=None): - names = (target,) if isinstance(target, str) else target - schema[source] = (transform, tuple(shape), dtype, tuple(TargetSpec(n, tuple(target_shape or shape), dtype) for n in names)) - - for source, target, shape, dtype in ( - ("video_patch_proj", "proj_in", (h, patch), "F32"), - ("audio_patch_proj", "audio_proj_in", (h, config["audio_in_channels"]), "F32"), - ("condition_proj", "context_embedder", (h, config["text_dim"]), "BF16"), - ("time_embedder.proj_in", "time_embedder.linear_1", (config["time_embed_hidden_dim"], config["freq_dim"]), "F32"), - ("time_embedder.proj_out", "time_embedder.linear_2", (time, config["time_embed_hidden_dim"]), "F32"), - ("final_layer.adaln_proj.linear", "norm_out.linear", (2 * h, time), "BF16"), - ("final_layer.video_out", "proj_out", (patch, h), "F32"), - ("final_layer.audio_out", "audio_proj_out", (config["audio_in_channels"], h), "F32"), - ): - add(source + ".weight", target + ".weight", shape, dtype) - add(source + ".bias", target + ".bias", shape[:1], dtype) - add("token_refiner.final_norm.weight", "token_refiner.final_norm.weight", (h,), transform="identity") - add("final_layer.norm.weight", "norm_out.norm.weight", (h,)) - add("rope.inv_freq", (), (config["rope_freq_dim"],), "F32", "validate_rope") - for source_prefix, target_prefix, layers, adaln in ( - ("blocks", "transformer_blocks", config["num_layers"], True), - ("token_refiner.blocks", "token_refiner.refiner_blocks", config["num_refiner_layers"], False), - ): - for i in range(layers): - source, target = f"{source_prefix}.{i}", f"{target_prefix}.{i}" - add(source + ".attn.qkv_proj.weight", tuple(target + f".attn.to_{q}.weight" for q in "qkv"), (3 * inner, h), transform="qkv_head_interleaved", target_shape=(inner, h)) - for q in "qk": - add(source + f".attn.{q}_norm.weight", target + f".attn.norm_{q}.weight", (config["attention_head_dim"],)) - add(source + ".attn.out_proj.weight", target + ".attn.to_out.0.weight", (h, inner)) - add(source + ".mlp.fc1.weight", target + ".ff.net.0.proj.weight", (2 * ffn, h), transform="swap_gate_value") - add(source + ".mlp.fc2.weight", target + ".ff.net.2.weight", (h, ffn)) - for n in (1, 2): - add(source + f".norm{n}.weight", target + f".norm{n}.weight", (h,)) - if adaln: - add(source + ".adaln_proj.linear.weight", target + ".adaln_proj.linear.weight", (18 * h, time)) - add(source + ".adaln_proj.linear.bias", target + ".adaln_proj.linear.bias", (18 * h,)) - return schema - - -_CONFIG_ALIASES = { - "token_refiner_num_layers": "num_refiner_layers", - "ffn_hidden_size": "ffn_dim", - "latents_dim": "in_channels", - "audio_latents_dim": "audio_in_channels", - "timestep_input_dim": "freq_dim", - "time_embed_hidden_size": "time_embed_hidden_dim", - "rope_inv_freq_len": "rope_freq_dim", -} - - -def _native_config(config): - normalized = dict(config) - if "patch_size" in normalized: - normalized["patch_size"] = tuple(normalized["patch_size"]) - for source, target in _CONFIG_ALIASES.items(): - if source in config: - if target in config and config[source] != config[target]: - raise ValueError(f"MiniMax-H3 conflicting config fields: {source}, {target}") - normalized[target] = config[source] - normalized.setdefault("rope_theta", 10000.0) - return normalized - - -class MiniMaxH3CheckpointPlan: - """Index/config-only planning. No shard existence checks or tensor reads.""" - - def __init__(self, checkpoint_dir, config=None): - self.checkpoint_dir = Path(checkpoint_dir) - self.index_path = self.checkpoint_dir / "model.safetensors.index.json" - if not self.index_path.is_file(): - raise FileNotFoundError(f"MiniMax-H3 safetensors index not found: {self.index_path}") - - with self.index_path.open("r", encoding="utf-8") as handle: - index = json.load(handle) - - weight_map = index.get("weight_map") - if not isinstance(weight_map, dict): - raise ValueError(f"MiniMax-H3 safetensors index must contain a dict weight_map: {self.index_path}") - if not weight_map: - raise ValueError(f"MiniMax-H3 safetensors index weight_map is empty: {self.index_path}") - - invalid_shard_names = sorted(name for name, shard_name in weight_map.items() if not isinstance(shard_name, str) or not shard_name) - if invalid_shard_names: - raise ValueError(f"MiniMax-H3 safetensors index contains invalid shard file names for tensors: {invalid_shard_names}") +class MiniMaxH3ShardCheckpoint: + """Index safetensors headers for selective loading, using upstream file discovery.""" - self.weight_map = dict(weight_map) - raw = any(n.startswith(("blocks.", "video_patch_proj.", "audio_patch_proj.", "condition_proj.", "final_layer.", "token_refiner.blocks.")) for n in weight_map) - native = any(n.startswith(("transformer_blocks.", "proj_in.", "audio_proj_in.", "context_embedder.", "norm_out.", "token_refiner.refiner_blocks.")) for n in weight_map) - if raw and native: - raise ValueError("MiniMax-H3 mixed official raw and native checkpoint keys") - self.format = "official_raw" if raw else "native" - self.entries = {} - self.targets = {} - if raw: - signature = {"blocks.0.attn.qkv_proj.weight", "blocks.0.mlp.fc1.weight", "video_patch_proj.weight", "final_layer.video_out.weight"} - if not signature.issubset(weight_map): - raise ValueError(f"MiniMax-H3 incomplete official signature: missing {sorted(signature - weight_map.keys())}") - config_path = self.checkpoint_dir / "config.json" - self.config = _native_config(json.loads(config_path.read_text()) if config_path.is_file() else (config or {})) - schema = _official_schema(self.config) - if config is not None: - runtime = _native_config(config) - keys = ( - "hidden_size", - "num_layers", - "num_attention_heads", - "attention_head_dim", - "ffn_dim", - "time_embed_dim", - "num_refiner_layers", - "freq_dim", - "rope_freq_dim", - "rope_theta", - "time_embed_hidden_dim", - "in_channels", - "audio_in_channels", - "text_dim", - "patch_size", - ) - for key in keys: - if key == "rope_theta" and key not in config: - continue - if key in runtime and runtime[key] != self.config[key]: - raise ValueError(f"MiniMax-H3 checkpoint/runtime config mismatch: {key}") - missing, unknown = sorted(schema.keys() - weight_map.keys()), sorted(weight_map.keys() - schema.keys()) - if missing or unknown: - raise ValueError(f"MiniMax-H3 source/target mapping mismatch: missing={missing}, unknown={unknown}") - self.entries = {name: SourcePlan(name, weight_map[name], *schema[name]) for name in sorted(schema)} - expected = [target.name for _, _, _, targets in schema.values() for target in targets] - validate_mapping(self.entries.values(), weight_map, expected) - self.targets = {target.name: (entry, target) for entry in self.entries.values() for target in entry.targets} + def __init__(self, checkpoint_dir): + checkpoint = Path(checkpoint_dir) + files = sorted(checkpoint.glob("*.safetensors")) if checkpoint.is_dir() else [checkpoint] + if not files or any(not path.is_file() for path in files): + raise FileNotFoundError(f"MiniMax-H3 safetensors checkpoint not found: {checkpoint}") + self.checkpoint_dir = checkpoint if checkpoint.is_dir() else checkpoint.parent + self.weight_map = {} + # Match the upstream model loader's directory/single-file discovery. + # Read only headers here; tensor data is loaded when a block requests it. + for path in files: + with safe_open(path, framework="pt", device="cpu") as source: + self.weight_map.update(dict.fromkeys(source.keys(), path.name)) @property def tensor_names(self): @@ -192,16 +31,15 @@ def tensor_names(self): @property def block_indices(self): - pattern = re.compile(r"^blocks\.(\d+)\.") if self.format == "official_raw" else _H3_BLOCK_KEY_RE + pattern = _H3_BLOCK_KEY_RE return tuple(sorted({int(match.group(1)) for name in self.weight_map if (match := pattern.match(name)) is not None})) def tensor_names_for_block(self, block_index): - prefix = "blocks" if self.format == "official_raw" else "transformer_blocks" - block_prefix = f"{prefix}.{int(block_index)}." + block_prefix = f"transformer_blocks.{int(block_index)}." return tuple(sorted(name for name in self.weight_map if name.startswith(block_prefix))) def non_block_tensor_names(self): - pattern = re.compile(r"^blocks\.(\d+)\.") if self.format == "official_raw" else _H3_BLOCK_KEY_RE + pattern = _H3_BLOCK_KEY_RE return tuple(sorted(name for name in self.weight_map if pattern.match(name) is None)) def shard_for_tensor(self, name): @@ -210,185 +48,7 @@ def shard_for_tensor(self, name): except KeyError as error: raise KeyError(f"MiniMax-H3 checkpoint is missing requested tensor: {name}") from error - def block_names(self, block_index): - return list(self.tensor_names_for_block(block_index)) - - def non_block_names(self): - return list(self.non_block_tensor_names()) - - def shards_for_sources(self, names): - return tuple(sorted({self.shard_for_tensor(name) for name in names})) - - -class MiniMaxH3SelectedSourceReader: - """Bounded CPU slice staging, with no large mmap views retained across sources.""" - - def __init__(self, plan, row_chunk_size=None, *, target_chunk_bytes=TARGET_CHUNK_BYTES): - """Prefer explicit row counts, then a byte budget, then legacy 128-row chunks. - - Set target_chunk_bytes=None to disable byte-based sizing. QKV head - boundaries and one-dimensional tensors keep their dedicated slicing. - """ - if plan.format != "official_raw": - raise ValueError("MiniMax-H3 selected adapter requires official raw format") - for name, value in (("row_chunk_size", row_chunk_size), ("target_chunk_bytes", target_chunk_bytes)): - if value is not None and (not isinstance(value, int) or isinstance(value, bool) or value < 1): - raise ValueError(f"{name} must be a positive integer") - self.plan = plan - self.row_chunk_size = row_chunk_size - self.target_chunk_bytes = target_chunk_bytes - - @contextmanager - def _source(self, name): - entry = self.plan.entries[name] - path = self.plan.checkpoint_dir / entry.source_shard - if not path.is_file(): - raise FileNotFoundError(f"MiniMax-H3 requested shard missing for {name}: {path}") - with safe_open(path, framework="pt", device="cpu") as reader: - source = reader.get_slice(name) - if tuple(source.get_shape()) != entry.shape: - raise ValueError(f"MiniMax-H3 shape mismatch for {name}: {source.get_shape()} != {entry.shape}") - if source.get_dtype() != entry.dtype: - raise ValueError(f"MiniMax-H3 dtype mismatch for {name}: {source.get_dtype()} != {entry.dtype}") - yield source - - def read_source_slice(self, name, row_start, row_end): - entry = self.plan.entries[name] - if not 0 <= row_start < row_end <= entry.shape[0]: - raise ValueError(f"MiniMax-H3 invalid row slice for {name}: {row_start}:{row_end}") - with self._source(name) as source: - # get_slice's PyTorch result can retain the whole source mmap storage. - # Copy only the requested rows so callers cannot retain that mapping. - return source[row_start:row_end].clone() - - def validate_sources(self, names): - for name in names: - with self._source(name) as source: - if name == "rope.inv_freq": - n = self.plan.config["rope_freq_dim"] - expected = 1.0 / (self.plan.config["rope_theta"] ** (torch.arange(0, 2 * n, 2, dtype=torch.float32, device="cpu") / (2 * n))) - actual = source[:].clone() - if not torch.equal(actual.view(torch.int32), expected.view(torch.int32)): - raise ValueError("MiniMax-H3 rope.inv_freq differs from native reconstruction") - - def _row_chunk_size(self, entry): - if self.row_chunk_size is not None: - return self.row_chunk_size - if self.target_chunk_bytes is not None and len(entry.shape) == 2 and entry.source_name.endswith(".weight"): - row_bytes = entry.shape[1] * {"F32": 4, "BF16": 2}[entry.dtype] - # The budget is a target: at least one complete row must fit a slice. - return max(1, self.target_chunk_bytes // row_bytes) - return 128 - - def _ranges(self, entry, requested): - if entry.transform == "qkv_head_interleaved": - dim = self.plan.config["attention_head_dim"] - for head in range(self.plan.config["num_attention_heads"]): - for component, target in enumerate(entry.targets): - if target.name in requested: - yield target.name, (head * 3 + component) * dim, head * dim, dim - elif entry.transform == "swap_gate_value": - half = entry.shape[0] // 2 - step = self._row_chunk_size(entry) - for source_start, target_start in ((half, 0), (0, half)): - for row in range(0, half, step): - yield entry.targets[0].name, source_start + row, target_start + row, min(step, half - row) - else: - size = entry.shape[0] - step = size if len(entry.shape) == 1 else self._row_chunk_size(entry) - for row in range(0, size, step): - yield entry.targets[0].name, row, row, min(step, size - row) - - def write_targets(self, destinations): - """Write {logical_name: (runtime_tensor, transpose)} in place, never reallocating.""" - requested = set(destinations) - unknown = requested - self.plan.targets.keys() - if unknown: - raise KeyError(f"MiniMax-H3 missing target plan: {sorted(unknown)}") - sources = sorted({self.plan.targets[name][0].source_name for name in requested}) - self.validate_sources(sources) - for name, (tensor, transpose) in destinations.items(): - spec = self.plan.targets[name][1] - shape = tuple(reversed(spec.shape)) if transpose else spec.shape - dtype = torch.float32 if spec.dtype == "F32" else torch.bfloat16 - if tuple(tensor.shape) != shape or tensor.dtype != dtype: - raise ValueError(f"MiniMax-H3 destination shape/dtype mismatch for {name}: expected {shape}, {dtype}") - for name in sources: - entry = self.plan.entries[name] - with self._source(name) as source: - for target, start, out_start, rows in self._ranges(entry, requested): - # Source slices reference safetensors mmap-backed storage. - # Consume this temporary view immediately while source is alive; - # never cache, return, or retain it for asynchronous use. - # copy_ must finish consuming CPU data before tile is released. - # Async copy would require redesigning source lifetime management. - tile = source[start : start + rows] - destination, transpose = destinations[target] - if transpose: - destination[:, out_start : out_start + rows].copy_(tile.t()) - else: - destination[out_start : out_start + rows].copy_(tile) - del tile - - def load_modules(self, roots, device="cpu", block_index=None, reusable=False): - """Bind official slices to native base_attrs; transpose exactly at this boundary. - - Reusable leaves retain both their *_cuda_buffer and active tensor identity. - Pre/post CPU leaves use pin_* attributes, matching native offload semantics. - No attention, RoPE, or padding computation is changed here. - """ - destinations, bindings, visited = {}, [], set() - stack = list(roots) - while stack: - module = stack.pop() - if id(module) in visited: - continue - visited.add(id(module)) - for name, attr, transpose in getattr(module, "base_attrs", ()): - if block_index is not None: - name = _H3_BLOCK_KEY_RE.sub(f"transformer_blocks.{int(block_index)}.", name) - if name in destinations: - raise ValueError(f"MiniMax-H3 duplicate target binding: {name}") - if name not in self.plan.targets: - raise KeyError(f"MiniMax-H3 missing target plan: {name}") - spec = self.plan.targets[name][1] - dtype = torch.float32 if spec.dtype == "F32" else torch.bfloat16 - storage_attr = f"{attr}_cuda_buffer" if reusable else f"pin_{attr}" if torch.device(device).type == "cpu" else attr - tensor = getattr(module, storage_attr, None) - if tensor is None: - tensor = torch.empty(spec.shape, dtype=dtype, device=device) - if transpose: - tensor = tensor.t() - setattr(module, storage_attr, tensor) - elif tensor.device.type != torch.device(device).type or (torch.device(device).index is not None and tensor.device.index != torch.device(device).index): - raise ValueError(f"MiniMax-H3 destination device mismatch: {name}") - destinations[name] = (tensor, transpose) - bindings.append((module, attr, tensor)) - stack.extend(child for child in getattr(module, "_modules", {}).values() if child is not None) - stack.extend(child for child in getattr(module, "_parameters", {}).values() if child is not None) - self.write_targets(destinations) - for module, attr, tensor in bindings: - setattr(module, attr, tensor if reusable or tensor.device.type != "cpu" else None) - if hasattr(module, "bias_name") and module.bias_name is None: - module.bias = None - module.pin_bias = None - - -class MiniMaxH3ShardCheckpoint(MiniMaxH3CheckpointPlan): - """Production entry: retain fail-fast validation of the complete shard set.""" - - def __init__(self, checkpoint_dir, config=None): - super().__init__(checkpoint_dir, config=config) - missing = sorted(name for name in set(self.weight_map.values()) if not (self.checkpoint_dir / name).is_file()) - if missing: - raise FileNotFoundError(f"MiniMax-H3 safetensors index references missing shard files: {missing}") - self.selected_reader = MiniMaxH3SelectedSourceReader(self) if self.format == "official_raw" else None - if self.selected_reader is not None: - self.selected_reader.validate_sources(["rope.inv_freq"]) - def load_tensors(self, names, device="cpu"): - if self.selected_reader is not None: - raise ValueError("MiniMax-H3 official raw tensors must use the selected slice adapter") missing = sorted(name for name in names if name not in self.weight_map) if missing: raise KeyError(f"MiniMax-H3 checkpoint is missing requested tensors: {missing}") @@ -406,4 +66,4 @@ def load_tensors(self, names, device="cpu"): return tensors -__all__ = ["MiniMaxH3CheckpointPlan", "MiniMaxH3SelectedSourceReader", "MiniMaxH3ShardCheckpoint"] +__all__ = ["MiniMaxH3ShardCheckpoint"] diff --git a/lightx2v/models/networks/minimax_h3/model.py b/lightx2v/models/networks/minimax_h3/model.py index 973502967..779bd8cb7 100644 --- a/lightx2v/models/networks/minimax_h3/model.py +++ b/lightx2v/models/networks/minimax_h3/model.py @@ -177,17 +177,13 @@ def _init_weights(self, weight_dict=None): if not self.config.get("dit_disk_streaming", False): return super()._init_weights(weight_dict) if weight_dict is not None: - raise ValueError("MiniMax-H3 dit_disk_streaming loads weights directly from the official checkpoint; explicit weight_dict is not supported.") + raise ValueError("MiniMax-H3 dit_disk_streaming loads weights directly from the diffusers checkpoint; explicit weight_dict is not supported.") self.transformer_weights = self.transformer_weight_class(self.config) self.pre_weight = self.pre_weight_class(self.config) self.post_weight = self.post_weight_class(self.config) checkpoint = self.transformer_weights.checkpoint - if checkpoint.selected_reader is not None: - checkpoint.selected_reader.load_modules([self.pre_weight, self.post_weight], device="cpu") - self._init_streaming_lora() - return None prepost_tensor_names = _collect_declared_base_tensor_names(self.pre_weight, self.post_weight) missing = sorted(name for name in prepost_tensor_names if name not in checkpoint.weight_map) if missing: @@ -204,6 +200,8 @@ def _init_weights(self, weight_dict=None): gc.collect() device_module = getattr(torch, torch.device(self.device).type, None) if device_module is not None and hasattr(device_module, "empty_cache"): + if torch.device(self.device).type == "mps": + device_module.synchronize() device_module.empty_cache() self._init_streaming_lora() return None diff --git a/lightx2v/models/networks/minimax_h3/streaming_lora.py b/lightx2v/models/networks/minimax_h3/streaming_lora.py index 76eca4447..4774f6314 100644 --- a/lightx2v/models/networks/minimax_h3/streaming_lora.py +++ b/lightx2v/models/networks/minimax_h3/streaming_lora.py @@ -163,21 +163,13 @@ def streaming_target_shapes(checkpoint, block, *resident_roots): for root in resident_roots: names.update(MiniMaxH3StreamingLora.weights(root)) shapes = {} - if checkpoint.selected_reader is not None: - for name in names: - spec = checkpoint.targets[name][1] - # Ordinary dynamic H3 LoRA uses BF16 factors. Sensitive FP32 - # linears require a different execution contract and are excluded. - if spec.dtype == "BF16": - shapes[name] = spec.shape - else: - by_shard = {} - for name in names: - by_shard.setdefault(checkpoint.weight_map[name], []).append(name) - for shard, shard_names in by_shard.items(): - with safe_open(checkpoint.checkpoint_dir / shard, framework="pt", device="cpu") as source: - for name in shard_names: - tensor = source.get_slice(name) - if tensor.get_dtype() == "BF16": - shapes[name] = tuple(tensor.get_shape()) + by_shard = {} + for name in names: + by_shard.setdefault(checkpoint.weight_map[name], []).append(name) + for shard, shard_names in by_shard.items(): + with safe_open(checkpoint.checkpoint_dir / shard, framework="pt", device="cpu") as source: + for name in shard_names: + tensor = source.get_slice(name) + if tensor.get_dtype() == "BF16": + shapes[name] = tuple(tensor.get_shape()) return shapes diff --git a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py index 45dc5ce81..17bc085fe 100644 --- a/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py +++ b/lightx2v/models/networks/minimax_h3/weights/transformer_weights.py @@ -38,6 +38,10 @@ def _empty_device_cache(): return device_module = getattr(torch, AI_DEVICE, None) if device_module is not None and hasattr(device_module, "empty_cache"): + if AI_DEVICE == "mps": + # Drain copies before empty_cache waits while holding the GIL. + # Metal completion may need the GIL to release safetensors storage. + device_module.synchronize() device_module.empty_cache() @@ -180,7 +184,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): self.streaming_lora = None if self.disk_streaming: if config.get("lazy_load", False): - raise NotImplementedError("MiniMax-H3 dit_disk_streaming reads the official sharded checkpoint directly and cannot be combined with converted lazy_load block shards.") + raise NotImplementedError("MiniMax-H3 dit_disk_streaming reads the diffusers sharded checkpoint directly and cannot be combined with converted lazy_load block shards.") if config.get("dit_quantized", False): raise NotImplementedError("MiniMax-H3 dit_disk_streaming does not support quantized DiT checkpoints yet.") if config.get("tensor_parallel", False): @@ -190,14 +194,8 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): checkpoint_dir = config.get("dit_original_ckpt") if checkpoint_dir is None: - raise ValueError("MiniMax-H3 dit_disk_streaming requires config['dit_original_ckpt'] to point to the official transformer checkpoint directory.") - self.checkpoint = MiniMaxH3ShardCheckpoint(checkpoint_dir, config=config) - if self.checkpoint.selected_reader is not None: - # Raw config aliases (e.g. token_refiner_num_layers) must reach - # the native pre/post constructors as well as the mapping plan. - for name, value in self.checkpoint.config.items(): - config.setdefault(name, value) - self.num_layers = int(config["num_layers"]) + raise ValueError("MiniMax-H3 dit_disk_streaming requires config['dit_original_ckpt'] to point to the diffusers transformer checkpoint directory.") + self.checkpoint = MiniMaxH3ShardCheckpoint(checkpoint_dir) expected_block_indices = tuple(range(self.num_layers)) if self.checkpoint.block_indices != expected_block_indices: raise ValueError(f"MiniMax-H3 dit_disk_streaming checkpoint block indices mismatch: expected {expected_block_indices}, found {self.checkpoint.block_indices}") @@ -210,7 +208,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None): if config.get("lazy_load", False): raise NotImplementedError( - "MiniMax-H3 reads the official sharded checkpoint directly; disk lazy_load requires a converted block-sharded checkpoint and is not supported yet. Use lazy_load=false with model or block CPU offload." + "MiniMax-H3 reads the diffusers sharded checkpoint directly; disk lazy_load requires a converted block-sharded checkpoint and is not supported yet. Use lazy_load=false with model or block CPU offload." ) self.blocks = WeightModuleList([MiniMaxH3TransformerBlockWeights(i, config) for i in range(self.num_layers)]) if config.get("cpu_offload", False) and config.get("offload_granularity", "model") == "block": @@ -238,11 +236,6 @@ def load_streaming_block(self, block_index): if self.streaming_lora is not None: # Finish the previous use before either base weights or factors change. self.streaming_lora.clear(self.streaming_block) - if self.checkpoint.selected_reader is not None: - self.checkpoint.selected_reader.load_modules([self.streaming_block], device=AI_DEVICE, block_index=block_index, reusable=True) - if self.streaming_lora is not None: - self.streaming_lora.load_block(self.streaming_block, block_index) - return self.streaming_block tensor_names = self.checkpoint.tensor_names_for_block(block_index) tensors = self.checkpoint.load_tensors(tensor_names, device="cpu") try: @@ -258,9 +251,6 @@ def _ensure_streaming_block(self): return self.streaming_block = MiniMaxH3TransformerBlockWeights(0, self.config, create_cuda_buffer=True) self.add_module("streaming_block", self.streaming_block) - if self.checkpoint.selected_reader is not None: - self.checkpoint.selected_reader.load_modules([self.streaming_block], device=AI_DEVICE, block_index=0, reusable=True) - return block0_tensors = self.checkpoint.load_tensors(self.checkpoint.tensor_names_for_block(0), device="cpu") try: self.streaming_block.load(block0_tensors) diff --git a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py index 8653b7467..c0cbba53d 100644 --- a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py +++ b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py @@ -485,12 +485,31 @@ def _prepare_references(self): raise ValueError(f"MiniMax-H3 ref2av accepts at most {MAX_REFERENCE_AUDIOS} audio-bearing references") return references + def _ensure_vae_loaded(self): + if self.video_vae is None or self.audio_vae is None: + self.video_vae, self.audio_vae = self.load_vae() + + def _release_low_memory_vae(self): + if self._is_mps_low_memory_streaming() and (self.video_vae is not None or self.audio_vae is not None): + self.video_vae = None + self.audio_vae = None + gc.collect() + self.maybe_empty_cache(force=True, collect_garbage=True) + def _encode_keyframes(self, keyframes): - latents = [] - for image in keyframes: - pixels = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1)[None, :, None].float().div_(255.0) - latents.append(self.video_vae.encode_condition(pixels, video=False)) - return latents + if not keyframes: + return [] + self._ensure_vae_loaded() + try: + latents = [] + for image in keyframes: + pixels = torch.from_numpy(np.asarray(image).copy()).permute(2, 0, 1)[None, :, None].float().div_(255.0) + latents.append(self.video_vae.encode_condition(pixels, video=False)) + return latents + finally: + # Switched i2av/l2av/fl2av requests need the VAE for conditioning, + # but MPS streaming must release it again before denoising. + self._release_low_memory_vae() def _encode_references(self, references): video_latents, audio_latents = [], [] @@ -582,6 +601,8 @@ def init_run(self): elif self.config.get("offload_granularity", "model") == "model": logger.info("Moving the native MiniMax-H3 transformer to the accelerator") self.model.to_cuda() + elif self.config.get("dit_disk_streaming", False): + logger.info("MiniMax-H3 diffusers disk streaming enabled; reusing one accelerator block buffer") else: logger.info("MiniMax-H3 block offload enabled; keeping source blocks on CPU and using two accelerator buffers") torch_device_module.synchronize() @@ -635,8 +656,7 @@ def _offload_transformer(self): metrics_labels=["MiniMaxH3Runner"], ) def run_vae_decoder(self, video_rows, audio_rows): - if self._is_mps_low_memory_streaming() and (self.video_vae is None or self.audio_vae is None): - self.video_vae, self.audio_vae = self.load_vae() + self._ensure_vae_loaded() video_rows = video_rows[self.scheduler.num_condition_video_rows :] audio_rows = audio_rows[self.scheduler.num_condition_audio_rows :] video_latents = unpatchify_video_tokens( @@ -732,11 +752,7 @@ def run_main(self): with suppress(Exception): self._offload_transformer() try: - if self._is_mps_low_memory_streaming() and (self.video_vae is not None or self.audio_vae is not None): - self.video_vae = None - self.audio_vae = None - gc.collect() - self.maybe_empty_cache(force=True, collect_garbage=True) + self._release_low_memory_vae() finally: self.end_run() # Decoded FP32 video is large (roughly 1.5 GiB at the default diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py index ef9967cc9..894f69c6a 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py @@ -51,8 +51,6 @@ ) from lightx2v.models.video_encoders.hf.minimax_h3.weights import ( SafetensorsSubsetReport, - _is_official_video_vae_checkpoint, - load_minimax_h3_video_vae_checkpoint, load_safetensors_subset, ) from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER @@ -89,102 +87,6 @@ def _component_dir(model_path: str | Path, component: str) -> Path: raise FileNotFoundError(f"Cannot find MiniMax-H3 {component!r} below {model_path}") -def _resolve_video_vae_dir(model_path: str | Path) -> Path: - """Resolve official ``video_vae`` and legacy ``vae`` component layouts.""" - model_path = Path(model_path) - candidates = [ - model_path / "video_vae", - model_path / "vae", - ] - if model_path.name in {"video_vae", "vae"}: - candidates.append(model_path) - - tried = [] - for candidate in candidates: - if candidate in tried: - continue - tried.append(candidate) - if candidate.is_dir(): - return candidate - - formatted = ", ".join(str(path) for path in tried) - raise FileNotFoundError(f"Cannot find MiniMax-H3 video VAE directory. Tried: {formatted}") - - -def _read_json_file(path: Path) -> dict: - with path.open("r", encoding="utf-8") as handle: - return json.load(handle) - - -def _normalize_official_video_vae_config(wrapper_config: dict, source_config: dict | None = None) -> dict: - """Convert the official FL2VA wrapper/source config pair to native keys.""" - if source_config is None: - return dict(wrapper_config) - - config = dict(wrapper_config) - if "in_channels" in source_config: - config["in_channels"] = source_config["in_channels"] - if "out_ch" in source_config: - config["out_channels"] = source_config["out_ch"] - if "z_channels" in source_config and "latent_channels" not in config: - config["latent_channels"] = source_config["z_channels"] - if "ch" in source_config and "ch_mult" in source_config: - config["block_out_channels"] = [int(source_config["ch"]) * int(value) for value in source_config["ch_mult"]] - if "num_res_blocks" in source_config: - config["layers_per_block"] = source_config["num_res_blocks"] - if "space_down" in source_config: - config["spatial_downsample_factors"] = source_config["space_down"] - if "time_down" in source_config: - config["temporal_downsample_factors"] = source_config["time_down"] - if "padding_mode" in source_config: - config["spatial_padding_mode"] = source_config["padding_mode"] - - vit_config = source_config.get("vit_decoder_kwargs") - if isinstance(vit_config, dict): - vit_mappings = { - "num_layers": "decoder_num_layers", - "heads": "decoder_num_attention_heads", - "dim_head": "decoder_attention_head_dim", - "rope_theta": "decoder_rope_theta", - "rope_dim_ratio": "decoder_rope_dim_ratio", - } - for source_key, target_key in vit_mappings.items(): - if source_key in vit_config: - config[target_key] = vit_config[source_key] - - wrapper_mappings = { - "vae_clip_length": "clip_length", - "vae_token_drop": "token_drop", - } - for source_key, target_key in wrapper_mappings.items(): - if source_key in wrapper_config: - config[target_key] = wrapper_config[source_key] - for key in ("latent_channels", "latents_mean", "latents_std"): - if key in wrapper_config: - config[key] = wrapper_config[key] - return config - - -def _load_video_vae_config_and_weight_path( - vae_dir: Path, - checkpoint_path: str | Path | None, -) -> tuple[dict, Path | str]: - wrapper_config = _read_json_file(vae_dir / "config.json") - source_config = None - default_weight_path: Path | str = vae_dir - - source_path = wrapper_config.get("source_path") - source_safetensors_path = wrapper_config.get("source_safetensors_path") - source_dir = vae_dir / source_path if isinstance(source_path, str) else None - if source_dir is not None and (source_dir / "config.json").is_file(): - source_config = _read_json_file(source_dir / "config.json") - if source_dir is not None and isinstance(source_safetensors_path, str): - default_weight_path = source_dir / source_safetensors_path - - weight_path = checkpoint_path if checkpoint_path is not None else default_weight_path - return _normalize_official_video_vae_config(wrapper_config, source_config), weight_path - - class _SwiGLU(nn.Module): """Checkpoint-compatible SwiGLU used by the ViT decoder.""" @@ -822,10 +724,12 @@ def from_pretrained( use_compile: bool = False, attn_type: str = "torch_sdpa", ) -> "MiniMaxH3VideoVAE": - vae_dir = _resolve_video_vae_dir(model_path) + vae_dir = _component_dir(model_path, "vae") if (checkpoint_path is None) != (quant_scheme is None): raise ValueError("MiniMax-H3 video VAE checkpoint_path and quant_scheme must be configured together") - config, weight_path = _load_video_vae_config_and_weight_path(vae_dir, checkpoint_path) + with (vae_dir / "config.json").open(encoding="utf-8") as handle: + config = json.load(handle) + weight_path = checkpoint_path if checkpoint_path is not None else vae_dir if quant_scheme == "fp8-f16-accum": validate_fp8_f16_accum_checkpoint(weight_path) fallback_reason = fp8_f16_accum_mm_unavailable_reason() @@ -851,10 +755,7 @@ def from_pretrained( attn_type=attn_type, ) model._reset_runtime_buffers() - if quant_scheme is None and _is_official_video_vae_checkpoint(weight_path): - model.load_report = load_minimax_h3_video_vae_checkpoint(model, weight_path) - else: - model.load_report = load_safetensors_subset(model, weight_path) + model.load_report = load_safetensors_subset(model, weight_path) if quant_scheme is not None: # Pack only after loading the checkpoint's original Q/K/V keys. model._pack_decoder_fp8_qkv() diff --git a/lightx2v/models/video_encoders/hf/minimax_h3/weights.py b/lightx2v/models/video_encoders/hf/minimax_h3/weights.py index 30e837ba9..ceb1452e1 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/weights.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/weights.py @@ -23,7 +23,6 @@ from __future__ import annotations -import re from dataclasses import dataclass from pathlib import Path @@ -160,177 +159,3 @@ def load_safetensors_subset(module: nn.Module, component_dir: str | Path) -> Saf if missing: raise RuntimeError(f"Failed to load MiniMax-H3 tensors: {missing[:20]}") return SafetensorsSubsetReport(report.component_dir, report.files, tuple(sorted(loaded)), report.ignored_keys) - - -_OFFICIAL_VIDEO_VAE_SIGNATURE = { - "decoder.x_embedder.weight", - "decoder.transformer_blocks.0.attn.to_qkv.weight", - "decoder.transformer_blocks.0.ff.w1.weight", -} - - -def _is_official_video_vae_checkpoint(component_dir: str | Path) -> bool: - """Detect the released Video VAE schema from safetensors keys.""" - keys: set[str] = set() - for filename in _component_files(component_dir): - with safe_open(filename, framework="pt", device="cpu") as checkpoint: - keys.update(checkpoint.keys()) - return _OFFICIAL_VIDEO_VAE_SIGNATURE <= keys and "decoder.proj_in.weight" not in keys - - -def _official_video_vae_targets(source_key: str) -> tuple[str, ...] | None: - if source_key == "decoder.mask_token": - return () - - qkv = re.fullmatch(r"(decoder\.transformer_blocks\.\d+\.attn)\.to_qkv\.(weight|bias)", source_key) - if qkv: - prefix, suffix = qkv.groups() - return tuple(f"{prefix}.to_{name}.{suffix}" for name in ("q", "k", "v")) - - w1 = re.fullmatch(r"(decoder\.transformer_blocks\.\d+\.ff)\.w1\.(weight|bias)", source_key) - if w1: - prefix, suffix = w1.groups() - return (f"{prefix}.net.0.proj.{suffix}",) - - target = source_key - down_block = re.fullmatch(r"encoder\.down\.(\d+)\.block\.(\d+)\.(.+)", target) - if down_block: - stage, block, suffix = down_block.groups() - suffix = suffix.replace("nin_shortcut", "conv_shortcut") - target = f"encoder.down_blocks.{stage}.resnets.{block}.{suffix}" - else: - downsample = re.fullmatch(r"encoder\.down\.(\d+)\.downsample\.conv\.(weight|bias)", target) - if downsample: - stage, suffix = downsample.groups() - target = f"encoder.down_blocks.{stage}.downsamplers.0.conv.{suffix}" - - target = target.replace("decoder.x_embedder.", "decoder.proj_in.") - target = re.sub(r"(decoder\.transformer_blocks\.\d+\.attn)\.to_out\.", r"\1.to_out.0.", target) - target = re.sub(r"(decoder\.transformer_blocks\.\d+\.ff)\.w2\.", r"\1.net.2.", target) - return (target,) - - -def _official_validation_error(*, unknown: list[str], missing: list[str], duplicates: list[str], shape_mismatches: list[str], dtype_mismatches: list[str]) -> RuntimeError: - details = [] - for label, values in ( - ("unknown", unknown), - ("missing", missing), - ("duplicate", duplicates), - ("shape_mismatch", shape_mismatches), - ("dtype_mismatch", dtype_mismatches), - ): - if values: - details.append(f"{label}={values[:20]}{' ...' if len(values) > 20 else ''}") - return RuntimeError("Official MiniMax-H3 Video VAE checkpoint validation failed: " + ", ".join(details)) - - -def validate_minimax_h3_video_vae_checkpoint(module: nn.Module, component_dir: str | Path) -> SafetensorsSubsetReport: - """Validate official-to-native coverage using safetensors metadata only.""" - files = _component_files(component_dir) - expected = _expected_specs(module) - assigned: set[str] = set() - unknown: list[str] = [] - duplicates: list[str] = [] - shape_mismatches: list[str] = [] - dtype_mismatches: list[str] = [] - ignored = 0 - - for filename in files: - with safe_open(filename, framework="pt", device="cpu") as checkpoint: - for source_key in checkpoint.keys(): - tensor_slice = checkpoint.get_slice(source_key) - source_shape = tuple(tensor_slice.get_shape()) - source_dtype = str(tensor_slice.get_dtype()) - targets = _official_video_vae_targets(source_key) - if targets == (): - decoder_dim = expected.get("decoder.proj_in.weight", ((0,), torch.float32))[0][0] - if source_shape != (1, 1, decoder_dim): - shape_mismatches.append(f"{source_key}:{source_shape}!={(1, 1, decoder_dim)}") - elif source_dtype != _SAFETENSORS_DTYPES.get(expected["decoder.proj_in.weight"][1]): - dtype_mismatches.append(source_key) - ignored += 1 - continue - if targets is None or any(target not in expected for target in targets): - unknown.append(source_key) - continue - - if len(targets) == 3: - if not source_shape or source_shape[0] % 3: - shape_mismatches.append(source_key) - continue - mapped_shape = (source_shape[0] // 3, *source_shape[1:]) - else: - mapped_shape = source_shape - for target in targets: - expected_shape, expected_dtype = expected[target] - if mapped_shape != expected_shape: - shape_mismatches.append(f"{source_key}->{target}:{mapped_shape}!={expected_shape}") - if source_dtype != _SAFETENSORS_DTYPES.get(expected_dtype): - dtype_mismatches.append(f"{source_key}->{target}") - if target in assigned: - duplicates.append(target) - assigned.add(target) - - missing = sorted(set(expected) - assigned) - if unknown or missing or duplicates or shape_mismatches or dtype_mismatches or ignored != 1: - if ignored != 1: - unknown.append(f"ignored_count:{ignored} (expected decoder.mask_token exactly once)") - raise _official_validation_error( - unknown=unknown, - missing=missing, - duplicates=duplicates, - shape_mismatches=shape_mismatches, - dtype_mismatches=dtype_mismatches, - ) - return SafetensorsSubsetReport(Path(component_dir), files, tuple(sorted(assigned)), ignored) - - -def _split_video_vae_qkv(tensor: torch.Tensor, num_heads: int, head_dim: int) -> tuple[torch.Tensor, ...]: - """Extract released per-head Q/K/V components for a weight or bias.""" - if tensor.ndim not in (1, 2) or tensor.shape[0] % 3: - raise ValueError("Video VAE fused QKV must be a weight/bias with rows divisible by 3") - if num_heads <= 0 or head_dim <= 0 or tensor.shape[0] != num_heads * 3 * head_dim: - raise ValueError("Video VAE fused QKV rows do not match target num_heads * 3 * head_dim") - # Released rows flatten [num_heads, 3, head_dim, ...], not [3, num_heads, head_dim, ...]. - # Both conventions have identical shapes, so shape-only validation cannot distinguish them. - view = tensor.reshape(num_heads, 3, head_dim, *tensor.shape[1:]) - return tuple(view[:, component].reshape(num_heads * head_dim, *tensor.shape[1:]).contiguous() for component in range(3)) - - -def load_minimax_h3_video_vae_checkpoint(module: nn.Module, component_dir: str | Path) -> SafetensorsSubsetReport: - """Stream the released official Video VAE schema into the native module.""" - report = validate_minimax_h3_video_vae_checkpoint(module, component_dir) - loaded: set[str] = set() - for filename in report.files: - with safe_open(filename, framework="pt", device="cpu") as checkpoint: - for source_key in checkpoint.keys(): - targets = _official_video_vae_targets(source_key) - if targets == (): - continue - tensor_slice = checkpoint.get_slice(source_key) - if len(targets) == 3: - attention, _ = _get_parent(module, targets[0].rsplit(".", 1)[0]) - components = _split_video_vae_qkv(checkpoint.get_tensor(source_key), attention.heads, attention.dim_head) - for target, component in zip(targets, components): - _assign_tensor(module, target, component) - loaded.add(target) - elif ".ff.w1." in source_key: - target = targets[0] - rows = tensor_slice.get_shape()[0] - half = rows // 2 - gate = tensor_slice[:half] - reordered = torch.empty(tuple(tensor_slice.get_shape()), dtype=gate.dtype) - reordered[half:].copy_(gate) - del gate - value = tensor_slice[half:] - reordered[:half].copy_(value) - del value - _assign_tensor(module, target, reordered) - loaded.add(target) - else: - target = targets[0] - _assign_tensor(module, target, checkpoint.get_tensor(source_key)) - loaded.add(target) - if loaded != set(report.loaded_keys): - raise RuntimeError("Official MiniMax-H3 Video VAE load did not reproduce validated target coverage") - return SafetensorsSubsetReport(report.component_dir, report.files, tuple(sorted(loaded)), report.ignored_keys) diff --git a/scripts/platforms/mps/README.md b/scripts/platforms/mps/README.md new file mode 100644 index 000000000..0ff3920ad --- /dev/null +++ b/scripts/platforms/mps/README.md @@ -0,0 +1,40 @@ +# MiniMax-H3 on Apple MPS + +MiniMax-H3 只支持 diffusers 权重布局。本机的两个脚本均使用： + +```text +/Users/yongyang/Documents/x2v/models/MiniMaxAI/diffusers/MiniMax-H3 +├── transformer/ +│ ├── config.json +│ ├── diffusion_pytorch_model.safetensors.index.json +│ └── diffusion_pytorch_model-*.safetensors +├── text_encoder/ +│ ├── config.json +│ ├── model.safetensors.index.json +│ └── model-*.safetensors +├── tokenizer/ +├── vae/ +│ ├── config.json +│ └── diffusion_pytorch_model*.safetensors +└── audio_vae/ + ├── config.json + └── diffusion_pytorch_model.safetensors +``` + +文本编码器仍遵循 Transformers 的 `model.safetensors.index.json` 命名。上图展示下载模型的标准目录；DiT 权重发现沿用上游规则：枚举目录中的 `*.safetensors`，也接受单个权重文件，不依赖索引文件名或参数名前缀白名单。流式读取器只扫描文件头来记录张量位置，数据在每层需要时读取；模型仍按 diffusers 参数名取权重，不转换 raw 的 QKV/FFN、配置别名或 `video_vae/source` 布局。视频 VAE 直接从 `vae/` 加载匹配的 diffusers 张量。 + +AdaLN 缓存构建器使用 ModelTC/LightX2V 上游实现，直接读取 diffusers 参数。MPS 逐层磁盘读取器只负责按需加载已有张量，不承担权重格式转换。 + +在 LightX2V 根目录执行： + +```bash +# 首次运行先生成与配置匹配的缓存；已有缓存时无需重复构建。 +bash tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +bash scripts/platforms/mps/run_minimax_h3_t2av.sh +``` + +默认使用 `configs/platforms/mps/minimax_h3_t2av_4step_512_22.json`：512×512、22 帧、4 步、BF16，DiT 和文本编码器逐层磁盘加载，VAE 按阶段加载。输出为 `save_results/output_lightx2v_minimax_h3_t2av.mp4`。这份 4 步配置用于本机快速验证。 + +MPS 配置使用独立缓存根目录 `~/.cache/lightx2v/adaln/diffusers`，避免与之前的权重混用。换权重或改动缓存相关计算后应重新生成缓存；步数和 flow shift 必须与推理配置一致。 + +流式权重初始化在清理 MPS 缓存前显式同步 GPU，避免 safetensors 存储释放回调与 Python GIL 相互等待。正式入口已经包含该处理，无需使用之前 bench 目录中的临时包装入口。 diff --git a/scripts/platforms/mps/run_minimax_h3_t2av.sh b/scripts/platforms/mps/run_minimax_h3_t2av.sh index a7a976ba9..09c79fd62 100755 --- a/scripts/platforms/mps/run_minimax_h3_t2av.sh +++ b/scripts/platforms/mps/run_minimax_h3_t2av.sh @@ -1,50 +1,25 @@ -#!/usr/bin/env bash -set -euo pipefail +#!/bin/bash -SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) -REPO_ROOT=$(cd -- "${SCRIPT_DIR}/../../.." && pwd) - -lightx2v_path=${LIGHTX2V_PATH:-${REPO_ROOT}} -model_path=${MODEL_PATH:-} -config_json=${CONFIG_JSON:-${lightx2v_path}/configs/platforms/mps/minimax_h3_t2av.json} -output_path=${OUTPUT_PATH:-${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av.mp4} +# set paths firstly +lightx2v_path=/Users/yongyang/Documents/x2v/LightX2V +model_path=/Users/yongyang/Documents/x2v/models/MiniMaxAI/diffusers/MiniMax-H3 +# set environment variables export PLATFORM=mps export DTYPE=BF16 export SENSITIVE_LAYER_DTYPE=BF16 export TOKENIZERS_PARALLELISM=false -export PYTHONFAULTHANDLER=1 -export PYTHONUNBUFFERED=1 -export PYTHONPATH="${lightx2v_path}:${PYTHONPATH:-}" - -if [[ -z "${model_path}" ]]; then - echo "MODEL_PATH must point to the MiniMax-H3 model directory." - exit 1 -fi - -[[ -d "${model_path}" ]] || { - echo "Model directory not found: ${model_path}" - exit 1 -} - -[[ -f "${config_json}" ]] || { - echo "Config file not found: ${config_json}" - exit 1 -} - -mkdir -p "$(dirname -- "${output_path}")" +export PYTHONPATH="${lightx2v_path}:$PYTHONPATH" -prompt=${PROMPT:-A cinematic fox walking through a snowy forest} -seed=${SEED:-42} +prompt='A cinematic fox walking through a snowy forest' -echo "Starting MiniMax-H3 t2av on platform=${PLATFORM}, dtype=${DTYPE}" -echo "Config: dit_disk_streaming=true, text_encoder_disk_streaming=true, VAE lazy lifecycle active" +mkdir -p "${lightx2v_path}/save_results" -python -m lightx2v.infer \ - --model_cls minimax_h3 \ - --task t2av \ - --model_path "${model_path}" \ - --config_json "${config_json}" \ - --prompt "${prompt}" \ - --save_result_path "${output_path}" \ - --seed "${seed}" +/opt/miniconda3/envs/torch/bin/python -m lightx2v.infer \ + --model_cls minimax_h3 \ + --task t2av \ + --model_path "$model_path" \ + --config_json "${lightx2v_path}/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json" \ + --prompt "$prompt" \ + --save_result_path "${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av.mp4" \ + --seed 42 diff --git a/tests/models/minimax_h3/test_adaln_cache.py b/tests/models/minimax_h3/test_adaln_cache.py deleted file mode 100644 index e9f3b0c81..000000000 --- a/tests/models/minimax_h3/test_adaln_cache.py +++ /dev/null @@ -1,203 +0,0 @@ -"""Exercise real cache math and IO with small official/native checkpoints.""" - -import importlib.util -import json -import sys -import types -from pathlib import Path - -import pytest -import torch -from safetensors.torch import load_file, save_file -from test_checkpoint_adapter import C, write_raw -from test_scheduler_layout import modules as scheduler_modules - -ROOT = Path(__file__).parents[3] - - -@pytest.fixture -def cache_modules(monkeypatch): - scheduler_modules.__wrapped__(monkeypatch) - envs = types.ModuleType("lightx2v.utils.envs") - envs.GET_DTYPE = lambda: torch.bfloat16 - monkeypatch.setitem(sys.modules, envs.__name__, envs) - - def load(name, path): - spec = importlib.util.spec_from_file_location(name, ROOT / path) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, name, module) - spec.loader.exec_module(module) - return module - - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.checkpoint", C) - load("lightx2v.models.networks.minimax_h3.infer.module_io", "lightx2v/models/networks/minimax_h3/infer/module_io.py") - load("lightx2v.models.networks.minimax_h3.infer.pre_infer", "lightx2v/models/networks/minimax_h3/infer/pre_infer.py") - load("lightx2v.models.networks.minimax_h3.adaln_cache_guide", "lightx2v/models/networks/minimax_h3/adaln_cache_guide.py") - cache = load("lightx2v.models.networks.minimax_h3.adaln_cache", "lightx2v/models/networks/minimax_h3/adaln_cache.py") - builder = load("h3_adaln_builder_under_test", "tools/cache_minimax_h3_adaln/builder.py") - return cache, builder - - -@pytest.fixture -def checkpoint(tmp_path): - raw = tmp_path / "raw" - raw.mkdir() - config, tensors, weight_map = write_raw(raw) - for name in tensors: - if name != "rope.inv_freq": - tensors[name] = tensors[name] * 0.001 - for shard in set(weight_map.values()): - save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, raw / shard) - config = C._native_config(config) - config.update(dit_original_ckpt=str(raw), use_adaln_cache=True, adaln_cache_dir=str(tmp_path / "cache"), task="fl2av", infer_steps=4) - return config, tensors, weight_map - - -def native_projections(tensors): - # Independent expected mapping for the cache's four projection families. - pairs = [("time_embedder.linear_1", "time_embedder.proj_in"), ("time_embedder.linear_2", "time_embedder.proj_out"), ("norm_out.linear", "final_layer.adaln_proj.linear")] - pairs.extend((f"transformer_blocks.{i}.adaln_proj.linear", f"blocks.{i}.adaln_proj.linear") for i in range(2)) - return {target + suffix: tensors[source + suffix] for target, source in pairs for suffix in (".weight", ".bias")} - - -@pytest.mark.parametrize("device", ["cpu", "mps"]) -def test_raw_generation_matches_native_and_loads(cache_modules, checkpoint, tmp_path, monkeypatch, device): - if device == "mps" and not torch.backends.mps.is_available(): - pytest.skip("MPS unavailable") - cache, builder = cache_modules - monkeypatch.setattr(builder, "AI_DEVICE", device) - monkeypatch.setattr(builder, "torch_device_module", getattr(torch, device)) - config, tensors, _ = checkpoint - expected = native_projections(tensors) - reader = builder._CheckpointTensors(builder._checkpoint_files(config), config=config) - for name, tensor in expected.items(): - actual = reader.get(name) - assert actual.dtype == tensor.dtype - assert torch.equal(actual, tensor) - with pytest.raises(KeyError, match="missing"): - reader.get("missing.weight") - raw_path = builder.build_persistent_adaln_cache(config) - raw_tables = load_file(raw_path / "adaln_cache.safetensors") - spec = cache._build_spec(config) - assert cache._validate_cache(raw_path, spec) - tables, norm = cache.load_persistent_adaln_cache(config, device) - for entry in spec["entries"]: - key = tuple(cache._timesteps_from_bits(entry["timestep_bits"]).tolist()) - for i in range(spec["num_layers"]): - tensor = tables[key][i] - assert tensor.device.type == device and tensor.dtype == torch.bfloat16 - assert tuple(tensor.shape) == cache._expected_table_shape(spec, entry) - assert torch.equal(tensor.cpu(), raw_tables[cache._block_key(i, entry)]) - assert norm[key].dtype == torch.bfloat16 and norm[key].device.type == device - assert tuple(norm[key].shape) == cache._expected_norm_out_shape(spec, entry) - assert torch.equal(norm[key].cpu(), raw_tables[cache._norm_out_key(entry)]) - native = tmp_path / "native" - native.mkdir() - save_file(expected, native / "model.safetensors") - for indexed in (False, True): - if indexed: - (native / "model.safetensors.index.json").write_text(json.dumps({"weight_map": {k: "model.safetensors" for k in expected}})) - native_config = dict(config, dit_original_ckpt=str(native), adaln_cache_dir=str(tmp_path / f"native-cache-{indexed}")) - native_path = builder.build_persistent_adaln_cache(native_config) - native_tables = load_file(native_path / "adaln_cache.safetensors") - assert raw_tables.keys() == native_tables.keys() - assert all(torch.equal(value, native_tables[name]) for name, value in raw_tables.items()) - with pytest.raises(FileExistsError): - builder.build_persistent_adaln_cache(config) - # Loader still rejects corrupted cache tensors. - key = next(iter(raw_tables)) - raw_tables[key] = raw_tables[key].float() - save_file(raw_tables, raw_path / "adaln_cache.safetensors") - with pytest.raises(FileNotFoundError, match="AdaLN cache not found"): - cache.load_persistent_adaln_cache(config, device) - - -@pytest.mark.parametrize("damage", ["shape", "dtype", "missing_key", "missing_shard", "runtime_config", "rope"]) -def test_invalid_raw_checkpoint_rejected(cache_modules, checkpoint, damage): - _, builder = cache_modules - config, tensors, weight_map = checkpoint - directory = Path(config["dit_original_ckpt"]) - name = "time_embedder.proj_in.weight" - if damage == "runtime_config": - config["hidden_size"] += 1 - elif damage == "missing_shard": - (directory / weight_map[name]).unlink() - elif damage == "missing_key": - del weight_map[name] - (directory / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) - else: - if damage == "shape": - tensors[name] = tensors[name][:1] - elif damage == "dtype": - tensors[name] = tensors[name].to(torch.bfloat16) - else: - name = "rope.inv_freq" - tensors[name] = torch.zeros_like(tensors[name]) - shard = weight_map[name] - save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, directory / shard) - with pytest.raises((ValueError, FileNotFoundError)): - builder.build_persistent_adaln_cache(config) - assert not any(Path(config["adaln_cache_dir"]).rglob("manifest.json")) - assert not any(Path(config["adaln_cache_dir"]).rglob(".building-*")) - - -@pytest.mark.parametrize("task", ["t2av", "i2av", "l2av", "fl2av"]) -def test_base_profiles_share_raw_cache(cache_modules, checkpoint, task): - cache, builder = cache_modules - config, _, _ = checkpoint - path = builder.build_persistent_adaln_cache(config) - requested = dict(config, task=task) - assert cache._cache_path(requested) == path - assert cache._validate_cache(path, cache._build_spec(requested)) - assert cache.load_persistent_adaln_cache(requested, "cpu") - ref_config = dict(config, task="ref2av") - assert cache._cache_path(ref_config) != path - assert not cache._validate_cache(path, cache._build_spec(ref_config)) - with pytest.raises(FileNotFoundError): - cache.load_persistent_adaln_cache(ref_config, "cpu") - - -@pytest.mark.parametrize("damage", ["manifest", "profile", "shape", "missing_tensor"]) -def test_cache_compatibility_validation(cache_modules, checkpoint, damage): - cache, builder = cache_modules - config, _, _ = checkpoint - path = builder.build_persistent_adaln_cache(config) - manifest_path = path / "manifest.json" - if damage in {"manifest", "profile"}: - manifest = json.loads(manifest_path.read_text()) - if damage == "manifest": - manifest["video_flow_shift"] += 1 - else: - manifest["entries"][0]["name"] = "ref2av_video_step_000" - manifest_path.write_text(json.dumps(manifest)) - else: - tensor_path = path / "adaln_cache.safetensors" - tensors = load_file(tensor_path) - name = next(iter(tensors)) - if damage == "shape": - tensors[name] = tensors[name][:1].contiguous() - else: - del tensors[name] - save_file(tensors, tensor_path) - assert not cache._validate_cache(path, cache._build_spec(config)) - with pytest.raises(FileNotFoundError, match="AdaLN cache not found"): - cache.load_persistent_adaln_cache(config, "cpu") - - -@pytest.mark.parametrize( - "update, error", - [ - ({"use_adaln_cache": False}, ValueError), - ({"adaln_cache_dir": ""}, ValueError), - ({"dummy_model": True}, NotImplementedError), - ({"task": "unsupported"}, ValueError), - ], -) -def test_invalid_cache_config(cache_modules, checkpoint, update, error): - cache, builder = cache_modules - config, _, _ = checkpoint - config.update(update) - with pytest.raises(error): - builder.build_persistent_adaln_cache(config) - with pytest.raises(error): - cache.load_persistent_adaln_cache(config, "cpu") diff --git a/tests/models/minimax_h3/test_audio_vae_offload.py b/tests/models/minimax_h3/test_audio_vae_offload.py deleted file mode 100644 index 7cdde4a73..000000000 --- a/tests/models/minimax_h3/test_audio_vae_offload.py +++ /dev/null @@ -1,101 +0,0 @@ -import importlib.util -import json -import sys -import types -import weakref -from pathlib import Path - -import pytest -import torch -from safetensors.torch import save_file -from torch.nn.utils.weight_norm import WeightNorm - -REPO_ROOT = Path(__file__).parents[3] - - -@pytest.fixture() -def audio_module(monkeypatch): - def load(name, relative): - spec = importlib.util.spec_from_file_location(name, REPO_ROOT / relative) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, name, module) - spec.loader.exec_module(module) - return module - - platform = types.ModuleType("lightx2v_platform.base.global_var") - platform.AI_DEVICE = "cpu" - monkeypatch.setitem(sys.modules, platform.__name__, platform) - load("lightx2v.models.video_encoders.hf.minimax_h3.weights", "lightx2v/models/video_encoders/hf/minimax_h3/weights.py") - return load("h3_audio_offload_under_test", "lightx2v/models/audio_encoders/hf/minimax_h3/audio_vae.py") - - -def tiny_model(audio_module): - model = audio_module.MiniMaxH3AudioVAE( - { - "encoder_dim": 4, - "latent_dim": 8, - "latent_channels": 2, - "num_attention_heads": 2, - "decoder_dim": 8, - "encoder_rates": [2], - "decoder_rates": [2], - "decoder_kernel_sizes": [4], - "resblock_kernel_sizes": [3], - "resblock_dilation_sizes": [[1]], - }, - device="cpu", - cpu_offload=True, - ) - return model.eval().requires_grad_(False) - - -def test_repeated_public_decode_and_checkpoint_contract(audio_module, tmp_path): - model = tiny_model(audio_module) - latent = torch.randn(2, 2, 16) - state = {name: value.clone() for name, value in model.state_dict().items()} - hooks = [(module, key, hook) for module in model.modules() for key, hook in module._forward_pre_hooks.items() if isinstance(hook, WeightNorm)] - assert hooks - weights = [weakref.ref(getattr(module, hook.name)) for module, _, hook in hooks] - outputs = [model.decode(latent) for _ in range(3)] - assert all(ref() is None for ref in weights) - assert all(torch.equal(outputs[0], output) for output in outputs) - assert outputs[0].shape == (1, 2, 32) - for module, key, hook in hooks: - assert module._forward_pre_hooks[key] is hook - assert hook.name not in vars(module) - assert hook.name + "_g" in module._parameters - assert hook.name + "_v" in module._parameters - assert state.keys() == model.state_dict().keys() - assert all(torch.equal(value, model.state_dict()[name]) for name, value in state.items()) - model.offload() # Cleanup is also safe when the derived attribute is absent. - model.to("cpu") - model.load_state_dict(state, strict=True) - assert torch.equal(outputs[0], model.decode(latent)) - - component = tmp_path / "audio_vae" - component.mkdir() - (component / "config.json").write_text(json.dumps(model.config)) - save_file(model.state_dict(), component / "model.safetensors") - loaded = audio_module.MiniMaxH3AudioVAE.from_pretrained(tmp_path, device="cpu", cpu_offload=True) - assert set(loaded.load_report.loaded_keys) == set(state) - assert torch.equal(outputs[0], loaded.decode(latent)) - assert all(parameter.device.type == "cpu" for parameter in loaded.parameters()) - assert all(buffer.device.type == "cpu" for buffer in loaded.buffers()) - - -def test_cleanup_only_targets_legacy_hook_attributes(audio_module): - model = tiny_model(audio_module) - plain = torch.nn.Linear(2, 2) - plain.cache = torch.ones(2) - model.unrelated = plain - custom = torch.nn.Module() - custom.register_parameter("kernel", torch.nn.Parameter(torch.randn(2, 2))) - torch.nn.utils.weight_norm(custom, name="kernel", dim=1) - model.custom = custom - parameter = plain.weight - cache = plain.cache - model.offload() - assert plain.weight is parameter - assert plain.cache is cache - assert "kernel" not in vars(custom) - assert set(custom.state_dict()) == {"kernel_g", "kernel_v"} diff --git a/tests/models/minimax_h3/test_checkpoint.py b/tests/models/minimax_h3/test_checkpoint.py deleted file mode 100644 index d2757dfe7..000000000 --- a/tests/models/minimax_h3/test_checkpoint.py +++ /dev/null @@ -1,117 +0,0 @@ -import importlib.util -import json -from pathlib import Path - -import pytest -import torch -from safetensors.torch import save_file - - -def _load_checkpoint_class(): - module_path = Path(__file__).parents[3] / "lightx2v/models/networks/minimax_h3/checkpoint.py" - spec = importlib.util.spec_from_file_location("minimax_h3_checkpoint", module_path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module.MiniMaxH3ShardCheckpoint - - -MiniMaxH3ShardCheckpoint = _load_checkpoint_class() - - -def _write_fake_checkpoint(tmp_path): - shard_1 = { - "proj_in.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3), - "transformer_blocks.0.attn.to_q.weight": torch.ones((2, 2), dtype=torch.bfloat16), - "transformer_blocks.1.attn.to_q.weight": torch.full((2, 2), 3, dtype=torch.bfloat16), - } - shard_2 = { - "transformer_blocks.0.ff.net.2.weight": torch.full((2, 2), 2, dtype=torch.bfloat16), - "norm_out.linear.weight": torch.full((2, 2), 4, dtype=torch.float32), - } - save_file(shard_1, tmp_path / "model-00001-of-00002.safetensors") - save_file(shard_2, tmp_path / "model-00002-of-00002.safetensors") - - weight_map = {name: "model-00001-of-00002.safetensors" for name in shard_1} | {name: "model-00002-of-00002.safetensors" for name in shard_2} - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"metadata": {"total_size": 0}, "weight_map": weight_map}), - encoding="utf-8", - ) - return shard_1, shard_2 - - -def test_tensor_block_and_non_block_names_are_deterministic(tmp_path): - _write_fake_checkpoint(tmp_path) - checkpoint = MiniMaxH3ShardCheckpoint(tmp_path) - - assert checkpoint.tensor_names == ( - "norm_out.linear.weight", - "proj_in.weight", - "transformer_blocks.0.attn.to_q.weight", - "transformer_blocks.0.ff.net.2.weight", - "transformer_blocks.1.attn.to_q.weight", - ) - assert checkpoint.block_indices == (0, 1) - assert checkpoint.tensor_names_for_block(0) == ( - "transformer_blocks.0.attn.to_q.weight", - "transformer_blocks.0.ff.net.2.weight", - ) - assert checkpoint.non_block_tensor_names() == ( - "norm_out.linear.weight", - "proj_in.weight", - ) - assert checkpoint.block_names(0) == list(checkpoint.tensor_names_for_block(0)) - assert checkpoint.non_block_names() == list(checkpoint.non_block_tensor_names()) - assert checkpoint.shard_for_tensor("transformer_blocks.0.attn.to_q.weight") == "model-00001-of-00002.safetensors" - assert checkpoint.shard_for_tensor("transformer_blocks.0.ff.net.2.weight") == "model-00002-of-00002.safetensors" - - -def test_load_tensors_reads_requested_block_tensors_across_shards(tmp_path): - shard_1, shard_2 = _write_fake_checkpoint(tmp_path) - checkpoint = MiniMaxH3ShardCheckpoint(tmp_path) - - tensors = checkpoint.load_tensors(checkpoint.tensor_names_for_block(0)) - - assert set(tensors) == set(checkpoint.tensor_names_for_block(0)) - assert torch.equal(tensors["transformer_blocks.0.attn.to_q.weight"], shard_1["transformer_blocks.0.attn.to_q.weight"]) - assert torch.equal(tensors["transformer_blocks.0.ff.net.2.weight"], shard_2["transformer_blocks.0.ff.net.2.weight"]) - - -def test_load_tensors_rejects_unknown_tensor(tmp_path): - _write_fake_checkpoint(tmp_path) - checkpoint = MiniMaxH3ShardCheckpoint(tmp_path) - - with pytest.raises(KeyError, match="missing requested tensors"): - checkpoint.load_tensors(["transformer_blocks.9.attn.to_q.weight"]) - - -def test_shard_for_tensor_rejects_unknown_tensor(tmp_path): - _write_fake_checkpoint(tmp_path) - checkpoint = MiniMaxH3ShardCheckpoint(tmp_path) - - with pytest.raises(KeyError, match="missing requested tensor"): - checkpoint.shard_for_tensor("missing") - - -def test_missing_index_raises_file_not_found(tmp_path): - with pytest.raises(FileNotFoundError, match="safetensors index not found"): - MiniMaxH3ShardCheckpoint(tmp_path) - - -def test_invalid_weight_map_shard_name_raises_value_error(tmp_path): - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": {"proj_in.weight": ""}}), - encoding="utf-8", - ) - - with pytest.raises(ValueError, match="invalid shard file names"): - MiniMaxH3ShardCheckpoint(tmp_path) - - -def test_missing_referenced_shard_raises_file_not_found(tmp_path): - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": {"proj_in.weight": "missing.safetensors"}}), - encoding="utf-8", - ) - - with pytest.raises(FileNotFoundError, match="missing shard files"): - MiniMaxH3ShardCheckpoint(tmp_path) diff --git a/tests/models/minimax_h3/test_checkpoint_adapter.py b/tests/models/minimax_h3/test_checkpoint_adapter.py deleted file mode 100644 index 1d4dbf22b..000000000 --- a/tests/models/minimax_h3/test_checkpoint_adapter.py +++ /dev/null @@ -1,468 +0,0 @@ -import importlib.util -import json -import sys -from pathlib import Path - -import pytest -import torch -from safetensors import safe_open -from safetensors.torch import save_file -from test_model_disk_streaming import _config, h3_model_modules # noqa: F401 -from test_transformer_disk_streaming import h3_modules # noqa: F401 - -ROOT = Path(__file__).parents[3] -spec = importlib.util.spec_from_file_location("h3_checkpoint_adapter_under_test", ROOT / "lightx2v/models/networks/minimax_h3/checkpoint.py") -C = importlib.util.module_from_spec(spec) -sys.modules[spec.name] = C -spec.loader.exec_module(C) - - -def write_raw(path): - config = { - "hidden_size": 3, - "num_layers": 2, - "token_refiner_num_layers": 1, - "num_attention_heads": 2, - "attention_head_dim": 2, - "ffn_hidden_size": 4, - "latents_dim": 1, - "audio_latents_dim": 2, - "patch_size": [1, 1, 2], - "text_dim": 5, - "timestep_input_dim": 4, - "time_embed_hidden_size": 3, - "time_embed_dim": 2, - "rope_inv_freq_len": 2, - } - schema = C._official_schema(C._native_config(config)) - tensors = {} - for i, (name, (_, shape, dtype, _)) in enumerate(schema.items()): - tensor = torch.arange(torch.Size(shape).numel(), dtype=torch.float32).reshape(shape) + i - tensors[name] = tensor.to(torch.float32 if dtype == "F32" else torch.bfloat16) - tensors["rope.inv_freq"] = 1.0 / (10000.0 ** (torch.arange(0, 4, 2, dtype=torch.float32) / 4)) - weight_map = {name: f"shard-{i % 2}.safetensors" for i, name in enumerate(sorted(tensors))} - for shard in set(weight_map.values()): - save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, path / shard) - (path / "config.json").write_text(json.dumps(config)) - (path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) - return config, tensors, weight_map - - -@pytest.fixture -def raw(tmp_path): - config, tensors, weight_map = write_raw(tmp_path) - plan = C.MiniMaxH3CheckpointPlan(tmp_path) - return config, tensors, weight_map, plan, C.MiniMaxH3SelectedSourceReader(plan, row_chunk_size=2) - - -@pytest.fixture(params=["cpu", "mps"]) -def destination_device(request): - if request.param == "mps" and not torch.backends.mps.is_available(): - pytest.skip("MPS is not available") - return request.param - - -def destinations(plan, names, transpose=False, device="cpu"): - return { - name: ( - torch.empty( - tuple(reversed(plan.targets[name][1].shape)) if transpose else plan.targets[name][1].shape, - dtype=torch.float32 if plan.targets[name][1].dtype == "F32" else torch.bfloat16, - device=device, - ), - transpose, - ) - for name in names - } - - -@pytest.mark.parametrize("dtype,expected", [("BF16", [2, 2, 1]), ("F32", [1, 1, 1, 1, 1])]) -def test_target_bytes_row_size_rounding_and_tail(raw, dtype, expected): - plan = raw[3] - entry = plan.entries["condition_proj.weight"]._replace(shape=(5, 3), dtype=dtype) - reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=13) - ranges = list(reader._ranges(entry, {entry.targets[0].name})) - assert [rows for _, _, _, rows in ranges] == expected - assert [(start, out) for _, start, out, _ in ranges] == [(sum(expected[:i]), sum(expected[:i])) for i in range(len(expected))] - - -@pytest.mark.parametrize("kwargs,expected", [({}, 1024), ({"target_chunk_bytes": 10752}, 2), ({"row_chunk_size": 3, "target_chunk_bytes": 1}, 3), ({"target_chunk_bytes": None}, 128)]) -def test_chunk_policy_precedence_and_legacy_fallback(raw, kwargs, expected): - plan = raw[3] - entry = plan.entries["condition_proj.weight"]._replace(shape=(2049, 2688)) - reader = C.MiniMaxH3SelectedSourceReader(plan, **kwargs) - ranges = list(reader._ranges(entry, {entry.targets[0].name})) - assert ranges[0][3] == expected - assert sum(rows for _, _, _, rows in ranges) == 2049 - if not kwargs: - fc1 = plan.entries["blocks.0.mlp.fc1.weight"]._replace(shape=(28672, 5376)) - assert next(reader._ranges(fc1, {fc1.targets[0].name}))[3] == 512 - - -def test_target_smaller_than_one_row(raw): - plan = raw[3] - entry = plan.entries["condition_proj.weight"] - reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=1) - assert [rows for _, _, _, rows in reader._ranges(entry, {entry.targets[0].name})] == [1, 1, 1] - - -@pytest.mark.parametrize("prefix", ["blocks.0", "token_refiner.blocks.0"]) -@pytest.mark.parametrize("budget,expected", [(18, [(4, 0, 3), (7, 3, 1), (0, 4, 3), (3, 7, 1)]), (1024, [(4, 0, 4), (0, 4, 4)])]) -def test_target_bytes_keeps_gate_value_halves_independent(raw, prefix, budget, expected): - plan = raw[3] - entry = plan.entries[prefix + ".mlp.fc1.weight"] - reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=budget) - assert [r[1:] for r in reader._ranges(entry, {entry.targets[0].name})] == expected - - -@pytest.mark.parametrize("prefix", ["blocks.0", "token_refiner.blocks.0"]) -def test_qkv_and_vectors_bypass_byte_policy(raw, monkeypatch, prefix): - plan = raw[3] - reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=1) - monkeypatch.setattr(reader, "_row_chunk_size", lambda *a: pytest.fail("dedicated slicing must bypass chunk policy")) - entry = plan.entries[prefix + ".attn.qkv_proj.weight"] - query = entry.targets[0].name - assert list(reader._ranges(entry, {query})) == [(query, 0, 0, 2), (query, 6, 2, 2)] - for name in ["video_patch_proj.bias", "blocks.0.norm1.weight", "token_refiner.final_norm.weight"]: - entry = plan.entries[name] - target = entry.targets[0].name - assert list(reader._ranges(entry, {target})) == [(target, 0, 0, entry.shape[0])] - - -@pytest.mark.parametrize("name", ["row_chunk_size", "target_chunk_bytes"]) -@pytest.mark.parametrize("value", [0, -1, 1.5, True]) -def test_invalid_chunk_policy_rejected(raw, name, value): - with pytest.raises(ValueError, match=name + " must be a positive integer"): - C.MiniMaxH3SelectedSourceReader(raw[3], **{name: value}) - - -@pytest.mark.parametrize("budget", [1, 13, C.TARGET_CHUNK_BYTES, None]) -@pytest.mark.parametrize("transpose", [False, True]) -def test_target_bytes_full_adapter_cpu_parity(raw, budget, transpose): - _, tensors, _, plan, _ = raw - reader = C.MiniMaxH3SelectedSourceReader(plan, target_chunk_bytes=budget) - dest, expected = {}, {} - for entry in plan.entries.values(): - source = tensors[entry.source_name] - for component, target in enumerate(entry.targets): - if entry.transform == "qkv_head_interleaved": - value = source.reshape(2, 3, 2, 3)[:, component].reshape(4, 3) - elif entry.transform == "swap_gate_value": - value = torch.cat((source[4:], source[:4])) - else: - value = source - transposed = transpose and value.ndim == 2 - expected[target.name] = value.t() if transposed else value - dest[target.name] = (torch.empty_like(expected[target.name]), transposed) - reader.write_targets(dest) - for name, (tensor, _) in dest.items(): - assert tensor.dtype == expected[name].dtype - assert torch.equal(tensor, expected[name]), name - - -def test_official_detection_and_complete_config_driven_plan(raw): - _, tensors, _, plan, _ = raw - assert plan.format == "official_raw" - assert plan.block_indices == (0, 1) - assert set(plan.entries) == set(tensors) - C.validate_mapping(plan.entries.values(), tensors, plan.targets) - assert len(plan.entries) == 47 # Tiny config, not the release's 535/638 counts. - assert len(plan.targets) == 52 - assert plan.entries["rope.inv_freq"].targets == () - - -@pytest.mark.parametrize("name", ["transformer_blocks.0.attn.to_q.weight", "proj_in.weight"]) -def test_native_format_not_misclassified(tmp_path, name): - (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": {name: "absent.safetensors"}})) - assert C.MiniMaxH3CheckpointPlan(tmp_path).format == "native" - - -def test_mixed_format_rejected(raw, tmp_path): - mapping = dict(raw[2], **{"transformer_blocks.0.attn.to_q.weight": "shard-0.safetensors"}) - (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping})) - with pytest.raises(ValueError, match="mixed"): - C.MiniMaxH3CheckpointPlan(tmp_path) - - -@pytest.mark.parametrize("prefix,target", [("blocks.0", "transformer_blocks.0"), ("token_refiner.blocks.0", "token_refiner.refiner_blocks.0")]) -def test_qkv_head_interleave_and_transpose(raw, prefix, target, destination_device): - _, tensors, _, plan, reader = raw - names = [f"{target}.attn.to_{q}.weight" for q in "qkv"] - raw_qkv = tensors[f"{prefix}.attn.qkv_proj.weight"] - for transpose in (False, True): - dest = destinations(plan, names, transpose, device=destination_device) - reader.write_targets(dest) - for component, name in enumerate(names): - expected = torch.vstack([raw_qkv[component * 2 : component * 2 + 2], raw_qkv[6 + component * 2 : 8 + component * 2]]) - actual = (dest[name][0].t() if transpose else dest[name][0]).cpu() - assert actual.shape == (4, 3) - assert torch.equal(actual, expected) - assert not torch.equal(actual, raw_qkv.chunk(3, dim=0)[component]) - - -def test_slice_only_reading_and_owned_slice_storage(raw, monkeypatch): - _, _, _, plan, reader = raw - reads = [] - - class Slice: - def __init__(self, inner, name): - self.inner, self.name = inner, name - - def get_shape(self): - return self.inner.get_shape() - - def get_dtype(self): - return self.inner.get_dtype() - - def __getitem__(self, index): - reads.append((self.name, index.start, index.stop)) - return self.inner[index] - - class Open: - def __init__(self, *args, **kwargs): - self.inner = safe_open(*args, **kwargs) - - def __enter__(self): - self.inner.__enter__() - return self - - def __exit__(self, *args): - return self.inner.__exit__(*args) - - def get_slice(self, name): - return Slice(self.inner.get_slice(name), name) - - def get_tensor(self, name): - pytest.fail("adapter must not get_tensor") - - monkeypatch.setattr(C, "safe_open", Open) - name = "blocks.0.attn.qkv_proj.weight" - sl = reader.read_source_slice(name, 2, 4) - assert sl.untyped_storage().nbytes() == 2 * 3 * 2 - reads.clear() - dest = destinations(plan, [t.name for t in plan.entries[name].targets]) - reader.write_targets(dest) - assert reads == [(name, i, i + 2) for i in range(0, 12, 2)] - - -@pytest.mark.parametrize("prefix,target", [("blocks.0", "transformer_blocks.0"), ("token_refiner.blocks.0", "token_refiner.refiner_blocks.0")]) -def test_fc1_swap_without_full_cat(raw, monkeypatch, prefix, target, destination_device): - _, tensors, _, plan, reader = raw - name = target + ".ff.net.0.proj.weight" - source = tensors[prefix + ".mlp.fc1.weight"] - dest = destinations(plan, [name], transpose=True, device=destination_device) - monkeypatch.setattr(torch, "cat", lambda *a, **k: pytest.fail("no full fused cat")) - reader.write_targets(dest) - actual = dest[name][0].t().cpu() - assert torch.equal(actual[:4], source[4:]) - assert torch.equal(actual[4:], source[:4]) - - -@pytest.mark.parametrize("transpose", [False, True]) -@pytest.mark.parametrize("strided", [False, True]) -def test_written_targets_outlive_source_and_own_storage(raw, destination_device, transpose, strided): - _, tensors, _, plan, reader = raw - mapping = { - "proj_in.weight": "video_patch_proj.weight", - "context_embedder.weight": "condition_proj.weight", - "transformer_blocks.0.attn.to_out.0.weight": "blocks.0.attn.out_proj.weight", - "norm_out.norm.weight": "final_layer.norm.weight", - } - dest = {} - expected = {} - for target, source in mapping.items(): - value = tensors[source] - transposed = transpose and value.ndim == 2 - expected[target] = value.t() if transposed else value - if strided: - tensor = torch.empty((*expected[target].shape, 2), dtype=value.dtype, device=destination_device)[..., 0] - else: - tensor = torch.empty_like(value, device=destination_device) - if transposed: - tensor = tensor.t() - dest[target] = (tensor, transposed) - pointers = {name: tensor.data_ptr() for name, (tensor, _) in dest.items()} - reader.write_targets(dest) - # write_targets has closed every source context; only destinations survive. - del reader - for target, (tensor, _) in dest.items(): - assert tensor.data_ptr() == pointers[target] - assert torch.equal(tensor.cpu(), expected[target]) - tensor.zero_() - # Mutating a destination must not modify or alias the checkpoint storage. - reader = C.MiniMaxH3SelectedSourceReader(plan, row_chunk_size=2) - reader.write_targets(dest) - for target, (tensor, _) in dest.items(): - assert tensor.data_ptr() == pointers[target] - assert torch.equal(tensor.cpu(), expected[target]) - - -def test_norm_refiner_and_nonblock_mapping_and_fp32(raw): - _, tensors, _, plan, reader = raw - names = plan.non_block_tensor_names() - reader.validate_sources(names) - targets = [t.name for name in names for t in plan.entries[name].targets] - dest = destinations(plan, targets) - reader.write_targets(dest) - for source, target in [ - ("video_patch_proj", "proj_in"), - ("audio_patch_proj", "audio_proj_in"), - ("condition_proj", "context_embedder"), - ("time_embedder.proj_in", "time_embedder.linear_1"), - ("time_embedder.proj_out", "time_embedder.linear_2"), - ("final_layer.adaln_proj.linear", "norm_out.linear"), - ("final_layer.video_out", "proj_out"), - ("final_layer.audio_out", "audio_proj_out"), - ]: - for suffix in ("weight", "bias"): - assert torch.equal(dest[f"{target}.{suffix}"][0], tensors[f"{source}.{suffix}"]) - assert dest[f"{target}.{suffix}"][0].dtype == tensors[f"{source}.{suffix}"].dtype - for q in "qk": - assert torch.equal(dest[f"token_refiner.refiner_blocks.0.attn.norm_{q}.weight"][0], tensors[f"token_refiner.blocks.0.attn.{q}_norm.weight"]) - assert torch.equal(dest["norm_out.norm.weight"][0], tensors["final_layer.norm.weight"]) - assert torch.equal(dest["token_refiner.final_norm.weight"][0], tensors["token_refiner.final_norm.weight"]) - - -@pytest.mark.parametrize("failure", ["shape", "dtype", "rope_value"]) -def test_invalid_real_header_or_rope_fails(raw, tmp_path, failure): - _, tensors, weight_map, _, reader = raw - name = "rope.inv_freq" if failure == "rope_value" else "blocks.0.attn.qkv_proj.weight" - tensors[name] = tensors[name][:-1] if failure == "shape" else tensors[name].float() if failure == "dtype" else tensors[name] + 1 - shard = weight_map[name] - save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, tmp_path / shard) - with pytest.raises(ValueError, match="rope.inv_freq" if failure == "rope_value" else failure + " mismatch"): - reader.validate_sources([name]) - - -def test_fp32_special_parameter_must_not_be_bf16(raw, tmp_path): - _, tensors, weight_map, _, reader = raw - name = "video_patch_proj.weight" - tensors[name] = tensors[name].bfloat16() - shard = weight_map[name] - save_file({k: v for k, v in tensors.items() if weight_map[k] == shard}, tmp_path / shard) - with pytest.raises(ValueError, match="dtype mismatch"): - reader.validate_sources([name]) - - -def test_wrong_destination_orientation_or_dtype_fails_before_copy(raw): - reader = raw[4] - name = "transformer_blocks.0.attn.to_q.weight" - with pytest.raises(ValueError, match="destination shape/dtype"): - reader.write_targets({name: (torch.zeros(4, 3, dtype=torch.bfloat16), True)}) - with pytest.raises(ValueError, match="destination shape/dtype"): - reader.write_targets({name: (torch.zeros(4, 3, dtype=torch.float32), False)}) - with pytest.raises(KeyError, match="missing target plan"): - reader.write_targets({"unknown": (torch.zeros(1), False)}) - - -@pytest.mark.parametrize("failure", ["unknown", "missing"]) -def test_source_coverage_rejected(raw, tmp_path, failure): - mapping = dict(raw[2]) - if failure == "unknown": - mapping["rope.not_allowed"] = "shard-0.safetensors" - else: - del mapping["blocks.0.norm1.weight"] - (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"weight_map": mapping})) - with pytest.raises(ValueError, match=failure): - C.MiniMaxH3CheckpointPlan(tmp_path) - - -def test_duplicate_and_missing_target_coverage_rejected(raw): - plan = raw[3] - entries = list(plan.entries.values()) - first = next(i for i, e in enumerate(entries) if e.targets) - entry = entries[first] - entries[first] = entry._replace(targets=entry.targets + entry.targets) - with pytest.raises(ValueError, match="duplicate target"): - C.validate_mapping(entries, plan.entries, plan.targets) - entries[first] = entry._replace(targets=()) - with pytest.raises(ValueError, match="target mapping mismatch: missing"): - C.validate_mapping(entries, plan.entries, plan.targets) - - -def test_cross_shard_block_load_and_missing_shard_fail(raw, tmp_path): - _, _, _, plan, reader = raw - sources = plan.tensor_names_for_block(0) - assert len(sources) == 10 - assert len(plan.shards_for_sources(sources)) == 2 - dest = destinations(plan, [t.name for n in sources for t in plan.entries[n].targets]) - reader.write_targets(dest) - (tmp_path / "shard-1.safetensors").unlink() - assert C.MiniMaxH3CheckpointPlan(tmp_path).block_indices == (0, 1) - with pytest.raises(FileNotFoundError, match="requested shard missing"): - reader.write_targets(dest) - with pytest.raises(FileNotFoundError, match="missing shard files"): - C.MiniMaxH3ShardCheckpoint(tmp_path) - - -def test_official_streaming_reuses_storage_and_recreates(tmp_path, monkeypatch, h3_modules): - _, weights_module, _ = h3_modules - config, tensors, _ = write_raw(tmp_path) - config.update(C._native_config(config), dit_disk_streaming=True, dit_original_ckpt=str(tmp_path)) - monkeypatch.setattr(weights_module, "AI_DEVICE", "cpu") - weights = weights_module.MiniMaxH3TransformerWeights(config) - block = weights.streaming_block - records = {} - for name, attr, _ in weights_module._iter_base_attrs(block): - records[name] = attr - - def buffers(): - stack, result = [weights.streaming_block], {} - while stack: - module = stack.pop() - for name, attr, _ in getattr(module, "base_attrs", ()): - tensor = getattr(module, attr) - assert tensor is getattr(module, attr + "_cuda_buffer") - result[name] = (id(tensor), tensor.data_ptr()) - stack.extend(getattr(module, "_modules", {}).values()) - return result - - before = buffers() - assert weights.load_streaming_block(1) is block - assert buffers() == before - assert torch.equal(block.ff.in_proj.weight[:, :4].t(), tensors["blocks.1.mlp.fc1.weight"][4:]) - weights.release_disk_streaming_buffer() - assert weights.streaming_block is None and block.attn.to_q.weight is None - new = weights.load_streaming_block(0) - assert new is not block - before = buffers() - assert weights.load_streaming_block(1) is new and buffers() == before - - -def test_raw_config_drives_native_dimensions_without_runtime_defaults(tmp_path, monkeypatch, h3_modules): - _, weights_module, _ = h3_modules - write_raw(tmp_path) - monkeypatch.setattr(weights_module, "AI_DEVICE", "cpu") - weights = weights_module.MiniMaxH3TransformerWeights({"dit_disk_streaming": True, "dit_original_ckpt": str(tmp_path)}) - assert weights.num_layers == 2 - assert weights.config["hidden_size"] == 3 - assert weights.config["num_refiner_layers"] == 1 - - -def test_official_model_prepost_initialization(tmp_path, monkeypatch, h3_model_modules): - _, _, _, transformer, _, _, model_module = h3_model_modules - config, tensors, _ = write_raw(tmp_path) - monkeypatch.setattr(transformer, "AI_DEVICE", "cpu") - config = _config(tmp_path, **config) - del config["num_refiner_layers"] - model = model_module.MiniMaxH3Model(str(tmp_path), config, torch.device("cpu")) - assert model.config["num_refiner_layers"] == 1 - assert model.config["freq_dim"] == 4 - assert model.pre_weight.proj_in.pin_weight.dtype == torch.float32 - assert model.post_weight.proj_out.pin_weight.dtype == torch.float32 - assert torch.equal(model.pre_weight.context_embedder.pin_weight.t(), tensors["condition_proj.weight"]) - assert torch.equal(model.post_weight.proj_out.pin_weight.t(), tensors["final_layer.video_out.weight"]) - - -def test_quantized_streaming_remains_rejected(tmp_path, h3_modules): - _, weights_module, _ = h3_modules - with pytest.raises(NotImplementedError, match="quantized"): - weights_module.MiniMaxH3TransformerWeights({"dit_disk_streaming": True, "dit_quantized": True}) - - -def test_quantized_nonstreaming_keeps_base_loader(tmp_path, h3_model_modules): - model_module = h3_model_modules[-1] - with pytest.raises(AssertionError, match="full checkpoint loading"): - model_module.MiniMaxH3Model( - str(tmp_path), _config(tmp_path, dit_disk_streaming=False, dit_quantized=True, dit_quant_scheme="int8-torchao", dit_quantized_ckpt=str(tmp_path)), torch.device("cpu") - ) diff --git a/tests/models/minimax_h3/test_model_disk_streaming.py b/tests/models/minimax_h3/test_model_disk_streaming.py deleted file mode 100644 index 9af3ed460..000000000 --- a/tests/models/minimax_h3/test_model_disk_streaming.py +++ /dev/null @@ -1,428 +0,0 @@ -import importlib.util -import json -import sys -import types -from pathlib import Path - -import pytest -import torch -from safetensors.torch import save_file - -REPO_ROOT = Path(__file__).parents[3] - - -class _FakeWeightModule: - def __init__(self): - self._modules = {} - self._parameters = {} - - def add_module(self, name, module): - self._modules[name] = module - setattr(self, name, module) - - def load(self, weight_dict): - for module in self._modules.values(): - if hasattr(module, "load"): - module.load(weight_dict) - - def load_state_dict(self, destination, block_index, adapter_block_index=None): - for module in self._modules.values(): - if hasattr(module, "load_state_dict"): - module.load_state_dict(destination, block_index, adapter_block_index) - - def to_cuda(self): - pass - - def to_cpu(self): - pass - - -class _FakeWeightModuleList(_FakeWeightModule): - def __init__(self, modules=None): - super().__init__() - self._list = [] - if modules is not None: - for module in modules: - self.append(module) - - def append(self, module): - self._list.append(module) - self.add_module(str(len(self._list) - 1), module) - - def __getitem__(self, index): - return self._list[index] - - def __len__(self): - return len(self._list) - - def __iter__(self): - return iter(self._list) - - -def _resolve_block_name(name, block_index): - parts = name.split(".", 2) - if len(parts) == 3 and parts[0] == "transformer_blocks" and parts[1].isdigit(): - return f"transformer_blocks.{int(block_index)}.{parts[2]}" - return name - - -class _FakeLinear: - def __init__(self, weight_name, bias_name=None, create_cuda_buffer=False, **_kwargs): - self.weight_name = weight_name - self.bias_name = bias_name - self.create_cuda_buffer = create_cuda_buffer - self.base_attrs = [(weight_name, "weight", True)] - if bias_name is not None: - self.base_attrs.append((bias_name, "bias", False)) - - def load(self, weight_dict): - for name, attr_name, transpose in self.base_attrs: - tensor = weight_dict[name] - if transpose: - tensor = tensor.t() - if self.create_cuda_buffer: - setattr(self, f"{attr_name}_cuda_buffer", tensor.clone()) - else: - setattr(self, attr_name, tensor.clone()) - if tensor.device.type == "cpu": - setattr(self, f"pin_{attr_name}", tensor.clone()) - del weight_dict[name] - - def load_state_dict(self, destination, block_index, adapter_block_index=None): - for name, attr_name, _transpose in self.base_attrs: - actual_name = _resolve_block_name(name, block_index) - if actual_name in destination: - buffer = getattr(self, f"{attr_name}_cuda_buffer") - setattr(self, attr_name, buffer.copy_(destination[actual_name])) - - -class _FakeRMS: - def __init__(self, weight_name, create_cuda_buffer=False, **_kwargs): - self.weight_name = weight_name - self.create_cuda_buffer = create_cuda_buffer - self.base_attrs = [(weight_name, "weight", False)] - - def load(self, weight_dict): - tensor = weight_dict[self.weight_name] - if self.create_cuda_buffer: - self.weight_cuda_buffer = tensor.clone() - else: - self.weight = tensor.clone() - self.pin_weight = tensor.clone() - if tensor.device.type == "cpu": - del weight_dict[self.weight_name] - - def load_state_dict(self, destination, block_index, adapter_block_index=None): - actual_name = _resolve_block_name(self.weight_name, block_index) - if actual_name in destination: - self.weight = self.weight_cuda_buffer.copy_(destination[actual_name]) - - -class _FakeLeaf: - base_attrs = () - - def __init__(self, *_args, **_kwargs): - pass - - def set_config(self, *_args, **_kwargs): - pass - - -class _FakeBaseTransformerModel: - load_ckpt_called = False - init_weights_called = False - - def __init__(self, model_path, config, device, model_type=None, lora_path=None, lora_strength=1.0): - self.device = torch.device(device) - self.model_path = model_path - self.config = config - self.lora_path = lora_path - self.lora_strength = lora_strength - self.model_type = model_type - self.cpu_offload = config.get("cpu_offload", False) - self.offload_granularity = config.get("offload_granularity", "block") - self.lazy_load = config.get("lazy_load", False) - self.dit_quantized = config.get("dit_quantized", False) - self.use_tp = config.get("tensor_parallel", False) - self.tp_size = 1 - self.tp_rank = 0 - self.seq_p_group = None - self.sensitive_layer = {} - - def _init_weights(self, weight_dict=None): - _FakeBaseTransformerModel.init_weights_called = True - if weight_dict is None: - self._load_ckpt(False, {}) - - def _load_ckpt(self, unified_dtype, sensitive_layer): - _FakeBaseTransformerModel.load_ckpt_called = True - raise AssertionError("full checkpoint loading must not run in disk streaming") - - def _apply_weights(self, weight_dict=None): - pass - - def _init_offload_manager(self): - raise AssertionError("WeightAsyncStreamManager/offload manager must not initialize in disk streaming") - - -def _load_module(module_name, relative_path): - module_path = REPO_ROOT / relative_path - spec = importlib.util.spec_from_file_location(module_name, module_path) - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) - return module - - -@pytest.fixture() -def h3_model_modules(monkeypatch): - for package_name in [ - "lightx2v", - "lightx2v.common", - "lightx2v.common.modules", - "lightx2v.models", - "lightx2v.models.networks", - "lightx2v.models.networks.minimax_h3", - "lightx2v.models.networks.minimax_h3.infer", - "lightx2v.models.networks.minimax_h3.weights", - "lightx2v.utils", - ]: - package = types.ModuleType(package_name) - package.__path__ = [] - monkeypatch.setitem(sys.modules, package_name, package) - - # Load main's FP8 helpers without importing the full runtime package. - for module_name in ( - "lightx2v.common.ops.mm.fp8_f16_accum", - "lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy", - ): - spec = importlib.util.spec_from_file_location(module_name, REPO_ROOT / (module_name.replace(".", "/") + ".py")) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, module_name, module) - spec.loader.exec_module(module) - - guide = _load_module("h3_adaln_cache_guide_under_test", "lightx2v/models/networks/minimax_h3/adaln_cache_guide.py") - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.adaln_cache_guide", guide) - - # Cache IO is outside these isolated model/streaming tests. - cache = types.ModuleType("lightx2v.models.networks.minimax_h3.adaln_cache") - cache.validate_adaln_cache_config = lambda config: None - cache.load_persistent_adaln_cache = lambda config, device: ({}, {}) - monkeypatch.setitem(sys.modules, cache.__name__, cache) - - weight_module = types.ModuleType("lightx2v.common.modules.weight_module") - weight_module.WeightModule = _FakeWeightModule - weight_module.WeightModuleList = _FakeWeightModuleList - monkeypatch.setitem(sys.modules, "lightx2v.common.modules.weight_module", weight_module) - - registry = types.ModuleType("lightx2v.utils.registry_factory") - registry.MM_WEIGHT_REGISTER = {"Default": _FakeLinear, "Default-ForceFp32": _FakeLinear} - registry.RMS_WEIGHT_REGISTER = {"torch_native": _FakeRMS} - registry.ROPE_REGISTER = {"torch_real_rope": _FakeLeaf} - registry.ATTN_WEIGHT_REGISTER = {"flash_attn3": _FakeLeaf} - monkeypatch.setitem(sys.modules, "lightx2v.utils.registry_factory", registry) - - envs = types.ModuleType("lightx2v.utils.envs") - envs.GET_DTYPE = lambda: torch.bfloat16 - monkeypatch.setitem(sys.modules, "lightx2v.utils.envs", envs) - - base_model = types.ModuleType("lightx2v.models.networks.base_model") - base_model.BaseTransformerModel = _FakeBaseTransformerModel - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.base_model", base_model) - - triton_ops = types.ModuleType("lightx2v.models.networks.minimax_h3.infer.triton_ops") - triton_ops.MiniMaxH3TritonRope = _FakeLeaf - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.infer.triton_ops", triton_ops) - - infer_module = types.ModuleType("lightx2v.models.networks.minimax_h3.infer.transformer_infer") - - class MiniMaxH3TransformerInfer: - def __init__(self, config): - self.config = config - - infer_module.MiniMaxH3TransformerInfer = MiniMaxH3TransformerInfer - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.infer.transformer_infer", infer_module) - - offload_module = types.ModuleType("lightx2v.models.networks.minimax_h3.infer.offload") - - class MiniMaxH3OffloadTransformerInfer: - pass - - offload_module.MiniMaxH3OffloadTransformerInfer = MiniMaxH3OffloadTransformerInfer - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.infer.offload", offload_module) - - for module_name, class_name in [ - ("module_io", "MiniMaxH3SequenceParallelState"), - ("post_infer", "MiniMaxH3PostInfer"), - ("pre_infer", "MiniMaxH3PreInfer"), - ]: - module = types.ModuleType(f"lightx2v.models.networks.minimax_h3.infer.{module_name}") - setattr(module, class_name, _FakeLeaf) - monkeypatch.setitem(sys.modules, f"lightx2v.models.networks.minimax_h3.infer.{module_name}", module) - - tensor_parallel = types.ModuleType("lightx2v.models.networks.minimax_h3.weights.tensor_parallel") - tensor_parallel.unwrap_tp_linear = lambda obj: obj - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.weights.tensor_parallel", tensor_parallel) - - checkpoint_module = _load_module( - "lightx2v.models.networks.minimax_h3.checkpoint", - "lightx2v/models/networks/minimax_h3/checkpoint.py", - ) - pre_module = _load_module( - "minimax_h3_pre_weights_for_model_test", - "lightx2v/models/networks/minimax_h3/weights/pre_weights.py", - ) - post_module = _load_module( - "minimax_h3_post_weights_for_model_test", - "lightx2v/models/networks/minimax_h3/weights/post_weights.py", - ) - transformer_module = _load_module( - "minimax_h3_transformer_weights_for_model_test", - "lightx2v/models/networks/minimax_h3/weights/transformer_weights.py", - ) - - weights_package = sys.modules["lightx2v.models.networks.minimax_h3.weights"] - weights_package.MiniMaxH3PreWeights = pre_module.MiniMaxH3PreWeights - weights_package.MiniMaxH3PostWeights = post_module.MiniMaxH3PostWeights - weights_package.MiniMaxH3TransformerWeights = transformer_module.MiniMaxH3TransformerWeights - - model_module = _load_module( - "minimax_h3_model_under_test", - "lightx2v/models/networks/minimax_h3/model.py", - ) - return checkpoint_module, pre_module, post_module, transformer_module, infer_module, offload_module, model_module - - -def _iter_base_attrs(module): - if hasattr(module, "base_attrs"): - yield from module.base_attrs - for child in getattr(module, "_modules", {}).values(): - yield from _iter_base_attrs(child) - - -def _tensors_from_roots(roots, block_index=None, value=1): - tensors = {} - for root in roots: - for name, _attr_name, transpose in _iter_base_attrs(root): - actual_name = _resolve_block_name(name, block_index) if block_index is not None else name - if transpose: - tensor = torch.full((2, 3), value, dtype=torch.bfloat16) - elif actual_name.endswith(".bias"): - tensor = torch.full((2,), value, dtype=torch.bfloat16) - else: - tensor = torch.full((3,), value, dtype=torch.bfloat16) - tensors[actual_name] = tensor - return tensors - - -def _write_fake_checkpoint(tmp_path, pre_module, post_module, transformer_module, num_layers=2): - config = {"num_layers": num_layers, "num_refiner_layers": 1} - tensors = {} - tensors.update(_tensors_from_roots([pre_module.MiniMaxH3PreWeights(config)], value=3)) - tensors.update(_tensors_from_roots([post_module.MiniMaxH3PostWeights(config)], value=4)) - block_template = transformer_module.MiniMaxH3TransformerBlockWeights(0, config) - for block_index in range(num_layers): - tensors.update(_tensors_from_roots([block_template], block_index=block_index, value=block_index + 1)) - - names = sorted(tensors) - shard_1_names = set(names[::2]) - shard_1 = {name: tensors[name] for name in names if name in shard_1_names} - shard_2 = {name: tensors[name] for name in names if name not in shard_1_names} - save_file(shard_1, tmp_path / "model-00001-of-00002.safetensors") - save_file(shard_2, tmp_path / "model-00002-of-00002.safetensors") - weight_map = { - **{name: "model-00001-of-00002.safetensors" for name in shard_1}, - **{name: "model-00002-of-00002.safetensors" for name in shard_2}, - } - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"metadata": {"total_size": 0}, "weight_map": weight_map}), - encoding="utf-8", - ) - return tensors - - -def _config(tmp_path, **overrides): - config = { - "seq_parallel": False, - "cfg_parallel": False, - "enable_cfg": False, - "cpu_offload": True, - "use_adaln_cache": True, - "adaln_cache_dir": str(tmp_path / "adaln-cache"), - "offload_granularity": "block", - "dit_disk_streaming": True, - "dit_original_ckpt": str(tmp_path), - "lazy_load": False, - "dit_quantized": False, - "dit_quant_scheme": "Default", - "tensor_parallel": False, - "num_layers": 2, - "num_refiner_layers": 1, - } - config.update(overrides) - return config - - -def test_non_disk_streaming_uses_base_init_weights(tmp_path, h3_model_modules): - _checkpoint_module, _pre_module, _post_module, _transformer_module, _infer_module, _offload_module, model_module = h3_model_modules - _FakeBaseTransformerModel.init_weights_called = False - _FakeBaseTransformerModel.load_ckpt_called = False - - with pytest.raises(AssertionError, match="full checkpoint loading"): - model_module.MiniMaxH3Model(str(tmp_path), _config(tmp_path, dit_disk_streaming=False), torch.device("cpu")) - - assert _FakeBaseTransformerModel.init_weights_called is True - assert _FakeBaseTransformerModel.load_ckpt_called is True - - -def test_disk_streaming_model_init_skips_full_checkpoint_load(tmp_path, monkeypatch, h3_model_modules): - checkpoint_module, pre_module, post_module, transformer_module, infer_module, offload_module, model_module = h3_model_modules - _write_fake_checkpoint(tmp_path, pre_module, post_module, transformer_module, num_layers=2) - calls = [] - - class SpyCheckpoint(checkpoint_module.MiniMaxH3ShardCheckpoint): - def load_tensors(self, names, device="cpu"): - calls.append(tuple(names)) - return super().load_tensors(names, device=device) - - monkeypatch.setattr(transformer_module, "MiniMaxH3ShardCheckpoint", SpyCheckpoint) - _FakeBaseTransformerModel.init_weights_called = False - _FakeBaseTransformerModel.load_ckpt_called = False - - model = model_module.MiniMaxH3Model(str(tmp_path), _config(tmp_path), torch.device("cpu")) - - assert _FakeBaseTransformerModel.load_ckpt_called is False - assert model.transformer_weights.disk_streaming is True - assert len(model.transformer_weights.blocks) == 0 - assert model.transformer_weights.streaming_block is not None - assert model.pre_weight.proj_in.pin_weight is not None - assert model.post_weight.proj_out.pin_weight is not None - assert model.transformer_infer_class is infer_module.MiniMaxH3TransformerInfer - assert model.transformer_infer_class is not offload_module.MiniMaxH3OffloadTransformerInfer - assert not hasattr(model.transformer_infer, "offload_manager") - - block0_names = model.transformer_weights.checkpoint.tensor_names_for_block(0) - block1_names = model.transformer_weights.checkpoint.tensor_names_for_block(1) - prepost_names = model_module._collect_declared_base_tensor_names(model.pre_weight, model.post_weight) - assert calls == [block0_names, prepost_names] - assert not any(set(call) == set(block1_names) for call in calls) - - -def test_disk_streaming_rejects_cpu_offload_false(tmp_path, h3_model_modules): - _checkpoint_module, _pre_module, _post_module, _transformer_module, _infer_module, _offload_module, model_module = h3_model_modules - - with pytest.raises(ValueError, match="requires cpu_offload=true"): - model_module.MiniMaxH3Model(str(tmp_path), _config(tmp_path, cpu_offload=False), torch.device("cpu")) - - -def test_disk_streaming_rejects_non_block_offload(tmp_path, h3_model_modules): - _checkpoint_module, _pre_module, _post_module, _transformer_module, _infer_module, _offload_module, model_module = h3_model_modules - - with pytest.raises(ValueError, match="requires offload_granularity='block'"): - model_module.MiniMaxH3Model(str(tmp_path), _config(tmp_path, offload_granularity="model"), torch.device("cpu")) - - -def test_cpu_offload_requires_persistent_cache(tmp_path, h3_model_modules): - model_module = h3_model_modules[-1] - with pytest.raises(ValueError, match="cpu_offload=true requires use_adaln_cache=true"): - model_module.MiniMaxH3Model(str(tmp_path), _config(tmp_path, use_adaln_cache=False), torch.device("cpu")) diff --git a/tests/models/minimax_h3/test_mps_low_memory_config.py b/tests/models/minimax_h3/test_mps_low_memory_config.py deleted file mode 100644 index 171350da0..000000000 --- a/tests/models/minimax_h3/test_mps_low_memory_config.py +++ /dev/null @@ -1,98 +0,0 @@ -import importlib.util -import json -import os -import subprocess -from pathlib import Path - -import pytest - -REPO_ROOT = Path(__file__).parents[3] -CONFIG_PATH = REPO_ROOT / "configs/platforms/mps/minimax_h3_t2av.json" -LAUNCHER_PATH = REPO_ROOT / "scripts/platforms/mps/run_minimax_h3_t2av.sh" - - -def _load_config(): - return json.loads(CONFIG_PATH.read_text(encoding="utf-8")) - - -@pytest.fixture() -def runner_module(monkeypatch): - helper_path = REPO_ROOT / "tests/models/minimax_h3/test_runner_mps_low_memory.py" - spec = importlib.util.spec_from_file_location("minimax_h3_runner_low_memory_test_helpers", helper_path) - helper = importlib.util.module_from_spec(spec) - spec.loader.exec_module(helper) - return helper._load_runner_module(monkeypatch) - - -def test_mps_minimax_h3_config_enables_low_memory_streaming(): - config = _load_config() - - assert config["cpu_offload"] is True - assert config["offload_granularity"] == "block" - assert config["dit_disk_streaming"] is True - assert config["text_encoder_cpu_offload"] is True - assert config["text_encoder_offload_granularity"] == "block" - assert config["text_encoder_disk_streaming"] is True - assert config["text_encoder_host_pinned"] is False - assert config["text_encoder_release_block_offload_buffers"] is True - assert config["vae_cpu_offload"] is True - assert config["lazy_load"] is False - assert config["unload_modules"] is False - assert config["warmup"] is False - assert config["attn_type"] == "torch_sdpa" - assert config["mps_sdpa_query_chunk_size"] == 512 - assert config["rms_type"] == "torch_native" - assert config["rope_type"] == "torch_real_rope" - assert config["vae_attn_type"] == "torch_sdpa" - assert config["dit_quantized"] is False - assert config["dit_quant_scheme"] == "Default" - assert config["text_encoder_quantized"] is False - assert config["video_vae_quantized"] is False - assert config["tensor_parallel"] is False - assert config["use_compile"] is False - assert config["vae_use_compile"] is False - assert "dit_original_ckpt" not in config - - -def test_real_mps_config_triggers_runner_low_memory_load_model(runner_module): - config = _load_config() - config.update({"task": "t2av", "model_path": "/tmp/minimax-h3"}) - runner = object.__new__(runner_module.MiniMaxH3Runner) - runner.config = config - calls = [] - runner.load_transformer = lambda: calls.append("transformer") or object() - runner.load_text_encoder = lambda: calls.append("text_encoder") or [object()] - runner.load_vae = lambda: calls.append("vae") or (object(), object()) - - assert runner._is_mps_low_memory_streaming() is True - runner.load_model() - - assert calls == ["transformer", "text_encoder"] - assert runner.video_vae is None - assert runner.audio_vae is None - - -def test_mps_launcher_has_safe_static_defaults(): - text = LAUNCHER_PATH.read_text(encoding="utf-8") - - assert "export PLATFORM=mps" in text - assert "export DTYPE=BF16" in text - assert "PYTORCH_ENABLE_MPS_FALLBACK" not in text - assert "CUDA" not in text - assert "PYTORCH_CUDA_ALLOC_CONF" not in text - - -def test_mps_launcher_fails_fast_without_model_path(): - env = os.environ.copy() - env.pop("MODEL_PATH", None) - result = subprocess.run( - [str(LAUNCHER_PATH)], - cwd=REPO_ROOT, - env=env, - check=False, - capture_output=True, - text=True, - ) - - assert result.returncode == 1 - assert "MODEL_PATH must point to the MiniMax-H3 model directory." in result.stdout diff --git a/tests/models/minimax_h3/test_query_chunked_sdpa.py b/tests/models/minimax_h3/test_query_chunked_sdpa.py deleted file mode 100644 index 3d59a4408..000000000 --- a/tests/models/minimax_h3/test_query_chunked_sdpa.py +++ /dev/null @@ -1,105 +0,0 @@ -import importlib.util -import sys -import types -from pathlib import Path - -import pytest -import torch -import torch.nn.functional as F - - -@pytest.fixture() -def sdpa(monkeypatch): - root = Path(__file__).parents[3] - registry = types.ModuleType("lightx2v.utils.registry_factory") - registry.ATTN_WEIGHT_REGISTER = lambda name: lambda cls: cls - monkeypatch.setitem(sys.modules, registry.__name__, registry) - package = types.ModuleType("chunked_sdpa_test") - package.__path__ = [str(root / "lightx2v/common/ops/attn")] - monkeypatch.setitem(sys.modules, package.__name__, package) - name = package.__name__ + ".torch_sdpa" - spec = importlib.util.spec_from_file_location(name, root / "lightx2v/common/ops/attn/torch_sdpa.py") - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, name, module) - spec.loader.exec_module(module) - return module - - -@pytest.mark.parametrize("length", [16, 19]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_dense_numerical_parity(sdpa, length, dtype): - torch.manual_seed(123) - q, k, v = [torch.randn(1, 4, length, 8, dtype=dtype) for _ in range(3)] - expected = F.scaled_dot_product_attention(q, k, v) - actual = sdpa._query_chunked_sdpa(q, k, v, 8) - assert actual.shape == expected.shape and actual.dtype == dtype - tolerance = 2e-2 if dtype == torch.bfloat16 else 1e-5 - torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) - - -@pytest.mark.parametrize( - "change", - [ - {}, - {"chunk_size": 0}, - {"scope": None}, - {"scope": "video_vae"}, - {"attn_mask": torch.ones(1)}, - {"causal": True}, - {"drop_rate": 0.1}, - {"device": "cpu"}, - {"device": "cuda"}, - {"shape": (1, 8, 19, 128)}, - ], -) -def test_applicability(sdpa, change): - args = {"chunk_size": 8, "scope": "minimax_h3_dit", "attn_mask": None, "causal": False, "drop_rate": 0} - device = change.get("device", "mps") - shape = change.get("shape", (1, 56, 19, 128)) - tensor = types.SimpleNamespace(device=torch.device(device), ndim=4, shape=shape) - args.update({k: v for k, v in change.items() if k not in ("device", "shape")}) - assert sdpa._use_h3_mps_query_chunks(tensor, tensor, tensor, **args) == (not change) - - -@pytest.mark.parametrize("options", [{}, {"attn_mask": torch.ones(9, 9, dtype=torch.bool)}, {"causal": True}, {"drop_rate": 0.2}]) -def test_original_path_preserved_on_cpu(sdpa, monkeypatch, options): - q, k, v = [torch.randn(9, 4, 8) for _ in range(3)] - calls = [] - original = F.scaled_dot_product_attention - - def record(*args, **kwargs): - calls.append(kwargs) - return original(*args, **kwargs) - - monkeypatch.setattr(sdpa.F, "scaled_dot_product_attention", record) - result = sdpa.TorchSDPAWeight().apply(q, k, v, attention_scope="minimax_h3_dit", mps_sdpa_query_chunk_size=512, **options) - assert result.shape == (9, 32) and result.dtype == q.dtype - assert len(calls) == 1 - assert calls[0]["is_causal"] == options.get("causal", False) - assert calls[0]["dropout_p"] == options.get("drop_rate", 0) - assert calls[0]["attn_mask"] is options.get("attn_mask") - - -def test_query_chunks_keep_complete_key_value_context(sdpa, monkeypatch): - q, k, v = [torch.randn(1, 4, 19, 8) for _ in range(3)] - lengths = [] - original = F.scaled_dot_product_attention - - def record(query, key, value, **kwargs): - assert key is k and value is v - assert kwargs == {"attn_mask": None, "dropout_p": 0.0, "is_causal": False} - lengths.append(query.shape[2]) - return original(query, key, value, **kwargs) - - monkeypatch.setattr(sdpa.F, "scaled_dot_product_attention", record) - assert sdpa._query_chunked_sdpa(q, k, v, 8).shape == q.shape - assert lengths == [8, 8, 3] - - -def test_gqa_fallback_preserved(sdpa): - generator = torch.Generator().manual_seed(123) - q = torch.randn(9, 4, 8, generator=generator) - k, v = [torch.randn(9, 2, 8, generator=generator) for _ in range(2)] - actual = sdpa.TorchSDPAWeight().apply(q, k, v, attention_scope="minimax_h3_dit", mps_sdpa_query_chunk_size=128) - expected = F.scaled_dot_product_attention(q.transpose(0, 1), k.repeat_interleave(2, dim=1).transpose(0, 1), v.repeat_interleave(2, dim=1).transpose(0, 1)).transpose(0, 1).reshape(9, 32) - torch.testing.assert_close(actual, expected) diff --git a/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py b/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py deleted file mode 100644 index e77dcb6b0..000000000 --- a/tests/models/minimax_h3/test_qwen3vl_disk_streaming.py +++ /dev/null @@ -1,461 +0,0 @@ -import importlib.util -import json -import sys -import types -from pathlib import Path -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F -from safetensors.torch import save_file - -REPO_ROOT = Path(__file__).parents[3] - - -class _FakeWeightModule: - def __init__(self): - self._modules = {} - self._parameters = {} - - def add_module(self, name, module): - self._modules[name] = module - setattr(self, name, module) - - def load(self, weight_dict): - for module in self._modules.values(): - if hasattr(module, "load"): - module.load(weight_dict) - - def load_state_dict(self, destination, block_index, adapter_block_index=None): - for module in self._modules.values(): - if hasattr(module, "load_state_dict"): - module.load_state_dict(destination, block_index, adapter_block_index) - - -class _FakeWeightModuleList(_FakeWeightModule): - def __init__(self, modules=None): - super().__init__() - self._list = [] - if modules is not None: - for module in modules: - self.append(module) - - def append(self, module): - self._list.append(module) - self.add_module(str(len(self._list) - 1), module) - - def __getitem__(self, index): - return self._list[index] - - def __len__(self): - return len(self._list) - - def __iter__(self): - return iter(self._list) - - -def _resolve_qwen_layer_name(name, layer_index): - prefix = "model.language_model.layers." - if not name.startswith(prefix): - return name - parts = name.split(".", 4) - if len(parts) == 5 and parts[3].isdigit(): - return f"{prefix}{int(layer_index)}.{parts[4]}" - return name - - -class _FakeLinear: - def __init__(self, weight_name, bias_name=None, create_cuda_buffer=False, **_kwargs): - self.weight_name = weight_name - self.bias_name = bias_name - self.create_cuda_buffer = create_cuda_buffer - self.base_attrs = [(weight_name, "weight", True)] - if bias_name is not None: - self.base_attrs.append((bias_name, "bias", False)) - - def load(self, weight_dict): - for name, attr_name, transpose in self.base_attrs: - tensor = weight_dict[name] - if transpose: - tensor = tensor.t() - if self.create_cuda_buffer: - setattr(self, f"{attr_name}_cuda_buffer", tensor.clone()) - else: - setattr(self, attr_name, tensor.clone()) - - def load_state_dict(self, destination, block_index, adapter_block_index=None): - for name, attr_name, _transpose in self.base_attrs: - actual_name = _resolve_qwen_layer_name(name, block_index) - if actual_name in destination: - buffer = getattr(self, f"{attr_name}_cuda_buffer") - setattr(self, attr_name, buffer.copy_(destination[actual_name])) - - -class _FakeRMS: - def __init__(self, weight_name, create_cuda_buffer=False, **_kwargs): - self.weight_name = weight_name - self.create_cuda_buffer = create_cuda_buffer - self.base_attrs = [(weight_name, "weight", False)] - - def load(self, weight_dict): - tensor = weight_dict[self.weight_name] - if self.create_cuda_buffer: - self.weight_cuda_buffer = tensor.clone() - else: - self.weight = tensor.clone() - - def load_state_dict(self, destination, block_index, adapter_block_index=None): - actual_name = _resolve_qwen_layer_name(self.weight_name, block_index) - if actual_name in destination: - self.weight = self.weight_cuda_buffer.copy_(destination[actual_name]) - - -class _FakeEmbedding: - def __init__(self, weight_name, *_args, **_kwargs): - self.weight_name = weight_name - self.weight = None - - def apply(self, input_indices): - return F.embedding(input_indices, self.weight) - - -class _FakeAttention: - def __init__(self, *_args, **_kwargs): - pass - - def apply(self, q, *_args, **_kwargs): - return q.reshape(q.shape[0], -1) - - -class _FakeLeaf: - def __init__(self, *_args, **_kwargs): - pass - - -class _FakeAttnWeightTemplate: - def __init__(self, *_args, **_kwargs): - pass - - -def _load_qwen_module(monkeypatch): - for package_name in [ - "lightx2v", - "lightx2v.common", - "lightx2v.common.modules", - "lightx2v.common.offload", - "lightx2v.common.ops", - "lightx2v.common.ops.attn", - "lightx2v.common.ops.embedding", - "lightx2v.common.ops.mm", - "lightx2v.common.ops.norm", - "lightx2v.models", - "lightx2v.models.input_encoders", - "lightx2v.models.input_encoders.hf", - "lightx2v.models.input_encoders.hf.minimax_h3", - "lightx2v.models.networks", - "lightx2v.models.networks.minimax_h3", - "lightx2v.models.networks.minimax_h3.weights", - "lightx2v.utils", - "lightx2v_platform", - "lightx2v_platform.base", - ]: - package = types.ModuleType(package_name) - package.__path__ = [] - monkeypatch.setitem(sys.modules, package_name, package) - - weight_module = types.ModuleType("lightx2v.common.modules.weight_module") - weight_module.WeightModule = _FakeWeightModule - weight_module.WeightModuleList = _FakeWeightModuleList - monkeypatch.setitem(sys.modules, "lightx2v.common.modules.weight_module", weight_module) - - event_manager = types.ModuleType("lightx2v.common.offload.event_manager") - - class RaisingEventSlotWeightAsyncStreamManager: - def __init__(self, *_args, **_kwargs): - raise AssertionError("EventSlotWeightAsyncStreamManager must not be instantiated") - - event_manager.EventSlotWeightAsyncStreamManager = RaisingEventSlotWeightAsyncStreamManager - monkeypatch.setitem(sys.modules, "lightx2v.common.offload.event_manager", event_manager) - - attn_template = types.ModuleType("lightx2v.common.ops.attn.template") - attn_template.AttnWeightTemplate = _FakeAttnWeightTemplate - monkeypatch.setitem(sys.modules, "lightx2v.common.ops.attn.template", attn_template) - for module_name, class_name in [ - ("lightx2v.common.ops.attn.torch_sdpa", "TorchSDPAWeight"), - ("lightx2v.common.ops.embedding.embedding_weight", "EmbeddingWeight"), - ("lightx2v.common.ops.mm.mm_weight", "MMWeight"), - ("lightx2v.common.ops.norm.rms_norm_weight", "RMSWeightFP32Qwen"), - ]: - module = types.ModuleType(module_name) - setattr(module, class_name, _FakeLeaf) - monkeypatch.setitem(sys.modules, module_name, module) - - vision = types.ModuleType("lightx2v.models.input_encoders.hf.minimax_h3.qwen3vl_vision") - vision.MiniMaxH3Qwen3VLVisionTower = _FakeLeaf - monkeypatch.setitem(sys.modules, "lightx2v.models.input_encoders.hf.minimax_h3.qwen3vl_vision", vision) - - packing = types.ModuleType("lightx2v.models.networks.minimax_h3.packing") - packing.VIDEO_TAG = 2 - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.packing", packing) - - packing_ref = types.ModuleType("lightx2v.models.networks.minimax_h3.packing_ref2av") - packing_ref.build_ref2av_presentation = lambda *_args, **_kwargs: None - packing_ref.sample_reference_video_frames = lambda *_args, **_kwargs: None - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.packing_ref2av", packing_ref) - - tp = types.ModuleType("lightx2v.models.networks.minimax_h3.weights.tensor_parallel") - tp.MiniMaxH3TensorParallelLinear = _FakeLinear - tp.unwrap_tp_linear = lambda obj: obj - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.weights.tensor_parallel", tp) - - envs = types.ModuleType("lightx2v.utils.envs") - envs.GET_DTYPE = lambda: torch.bfloat16 - monkeypatch.setitem(sys.modules, "lightx2v.utils.envs", envs) - - registry = types.ModuleType("lightx2v.utils.registry_factory") - registry.ATTN_WEIGHT_REGISTER = {"torch_sdpa": _FakeAttention} - registry.EMBEDDING_WEIGHT_REGISTER = {"Default": _FakeEmbedding} - registry.MM_WEIGHT_REGISTER = {"Default": _FakeLinear} - registry.RMS_WEIGHT_REGISTER = {"fp32_variance_qwen": _FakeRMS} - monkeypatch.setitem(sys.modules, "lightx2v.utils.registry_factory", registry) - - global_var = types.ModuleType("lightx2v_platform.base.global_var") - global_var.AI_DEVICE = "cpu" - monkeypatch.setitem(sys.modules, "lightx2v_platform.base.global_var", global_var) - - spec = importlib.util.spec_from_file_location( - "qwen3vl_disk_streaming_under_test", - REPO_ROOT / "lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py", - ) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - module.AI_DEVICE = "cpu" - module.torch_device_module = SimpleNamespace(synchronize=lambda: None) - return module - - -@pytest.fixture() -def qwen_module(monkeypatch): - return _load_qwen_module(monkeypatch) - - -def _tiny_text_config(): - return { - "hidden_size": 8, - "intermediate_size": 16, - "num_hidden_layers": 2, - "num_attention_heads": 2, - "num_key_value_heads": 1, - "head_dim": 4, - "vocab_size": 32, - "rms_norm_eps": 1e-6, - "rope_theta": 10000.0, - } - - -def _make_backbone(qwen_module): - return qwen_module._Qwen3VLTextBackboneWeights( - {"text_encoder_disk_streaming": True}, - _tiny_text_config(), - num_layers=2, - attn_type="torch_sdpa", - block_offload=False, - disk_streaming=True, - ) - - -def _tensors_for_names(names, value): - tensors = {} - for name in names: - if name.endswith("embed_tokens.weight"): - tensors[name] = torch.full((32, 8), value, dtype=torch.bfloat16) - elif name.endswith(".weight") and any(proj in name for proj in (".q_proj.", ".k_proj.", ".v_proj.", ".o_proj.", ".gate_proj.", ".up_proj.", ".down_proj.")): - tensors[name] = torch.full((2, 3), value, dtype=torch.bfloat16) - else: - tensors[name] = torch.full((3,), value, dtype=torch.bfloat16) - return tensors - - -def _write_fake_checkpoint(tmp_path, backbone, qwen_module): - template_layer = backbone.layers[0] - embedding_name = backbone.embed_tokens.weight_name - layer0_names = qwen_module._Qwen3VLTextBackboneWeights._layer_tensor_names(template_layer, 0) - layer1_names = qwen_module._Qwen3VLTextBackboneWeights._layer_tensor_names(template_layer, 1) - tensors = {} - tensors.update(_tensors_for_names((embedding_name,), 5)) - tensors.update(_tensors_for_names(layer0_names, 1)) - tensors.update(_tensors_for_names(layer1_names, 2)) - - names = sorted(tensors) - shard_1_names = set(names[::2]) - shard_1 = {name: tensors[name] for name in names if name in shard_1_names} - shard_2 = {name: tensors[name] for name in names if name not in shard_1_names} - save_file(shard_1, tmp_path / "model-00001-of-00002.safetensors") - save_file(shard_2, tmp_path / "model-00002-of-00002.safetensors") - weight_map = { - **{name: "model-00001-of-00002.safetensors" for name in shard_1}, - **{name: "model-00002-of-00002.safetensors" for name in shard_2}, - } - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"metadata": {"total_size": 0}, "weight_map": weight_map}), - encoding="utf-8", - ) - return embedding_name, layer0_names, layer1_names, weight_map - - -def test_qwen3vl_disk_streaming_reuses_one_layer_buffer(tmp_path, monkeypatch, qwen_module): - backbone = _make_backbone(qwen_module) - embedding_name, layer0_names, layer1_names, weight_map = _write_fake_checkpoint(tmp_path, backbone, qwen_module) - calls = [] - original_loader = qwen_module._load_selected_checkpoint_tensors - - def spy_loader(text_encoder_path, weight_map_arg, names): - calls.append(tuple(names)) - return original_loader(text_encoder_path, weight_map_arg, names) - - monkeypatch.setattr(qwen_module, "_load_selected_checkpoint_tensors", spy_loader) - - backbone.init_disk_streaming(tmp_path, weight_map) - assert calls == [layer0_names] - assert backbone.offload_manager is None - layer0 = backbone.load_streaming_layer(0) - layer_id = id(layer0) - buffer = layer0.self_attn.q_proj.weight_cuda_buffer - buffer_id = id(buffer) - assert torch.all(layer0.self_attn.q_proj.weight == 1) - - layer1 = backbone.load_streaming_layer(1) - assert id(layer1) == layer_id - assert id(layer1.self_attn.q_proj.weight_cuda_buffer) == buffer_id - assert torch.all(layer1.self_attn.q_proj.weight == 2) - - hidden_states = backbone._forward_streaming_embedding(torch.tensor([0, 1], dtype=torch.long)) - assert hidden_states.dtype == torch.bfloat16 - assert tuple(hidden_states.shape) == (2, 8) - assert getattr(backbone.embed_tokens, "weight", None) is None - assert getattr(backbone.embed_tokens, "pin_weight", None) is None - - assert calls == [layer0_names, layer0_names, layer1_names, (embedding_name,)] - assert all(set(call) in [set(layer0_names), set(layer1_names), {embedding_name}] for call in calls) - - -def test_qwen3vl_disk_streaming_rejects_vision_inputs(tmp_path, qwen_module): - backbone = _make_backbone(qwen_module) - _embedding_name, _layer0_names, _layer1_names, weight_map = _write_fake_checkpoint(tmp_path, backbone, qwen_module) - backbone.init_disk_streaming(tmp_path, weight_map) - - with pytest.raises(NotImplementedError, match="text-only t2av"): - backbone.forward( - torch.tensor([0], dtype=torch.long), - vision_mask=torch.tensor([True]), - vision_embeds=torch.zeros((1, 8), dtype=torch.bfloat16), - ) - - -def test_qwen3vl_release_disk_streaming_buffer_clears_device_refs(tmp_path, qwen_module): - backbone = _make_backbone(qwen_module) - _embedding_name, _layer0_names, _layer1_names, weight_map = _write_fake_checkpoint(tmp_path, backbone, qwen_module) - backbone.init_disk_streaming(tmp_path, weight_map) - old_layer = backbone.streaming_layer - assert old_layer.self_attn.q_proj.weight_cuda_buffer is not None - - backbone.release_disk_streaming_buffer() - - assert backbone.streaming_layer is None - assert old_layer.self_attn.q_proj.weight is None - assert old_layer.self_attn.q_proj.weight_cuda_buffer is None - assert old_layer.input_layernorm.weight is None - assert old_layer.input_layernorm.weight_cuda_buffer is None - - -@pytest.mark.parametrize("mode", ["disk", "resident", "block"]) -@pytest.mark.parametrize("release_buffers", [True, False]) -@pytest.mark.parametrize("fail_forward", [False, True]) -def test_public_infer_offload_lifecycle(tmp_path, monkeypatch, qwen_module, mode, release_buffers, fail_forward): - from unittest.mock import Mock - - # Exercise the real constructor's mode flags without needing an accelerator. - monkeypatch.setattr(qwen_module, "AI_DEVICE", "mps") - encoder = qwen_module.MiniMaxH3Qwen3VLTextEncoder( - { - "task": "t2av", - "text_encoder_cpu_offload": True, - "text_encoder_offload_granularity": "model" if mode == "resident" else "block", - "text_encoder_disk_streaming": mode == "disk", - "text_encoder_release_block_offload_buffers": release_buffers, - "text_encoder_load_on_init": False, - } - ) - monkeypatch.setattr(qwen_module, "AI_DEVICE", "cpu") - monkeypatch.setattr(qwen_module, "MINIMAX_H3_TEXT_HIDDEN_SIZE", 8) - backbone = qwen_module._Qwen3VLTextBackboneWeights( - encoder.config, - _tiny_text_config(), - num_layers=2, - block_offload=encoder.block_offload, - disk_streaming=encoder.disk_streaming, - ) - encoder.text_encoder = backbone - encoder.tokenizer = Mock(return_value={"input_ids": [0, 1, 2]}) - assert (encoder.cpu_offload, encoder.block_offload, encoder.disk_streaming) == (True, mode == "block", mode == "disk") - events = [] - - def layer_forward(self, hidden_states, position_embeddings): - events.append("layer") - if fail_forward and events.count("layer") == 2: - raise RuntimeError("injected forward failure") - return hidden_states + 1 - - monkeypatch.setattr(qwen_module._Qwen3VLDecoderLayerWeights, "forward", layer_forward) - if mode == "disk": - _, _, _, weight_map = _write_fake_checkpoint(tmp_path, backbone, qwen_module) - backbone.init_disk_streaming(tmp_path, weight_map) - else: - backbone.embed_tokens.weight = torch.full((32, 8), 5, dtype=torch.bfloat16) - - def block_forward(input_ids, *args): - events.append("block") - if fail_forward: - raise RuntimeError("injected forward failure") - return torch.full((input_ids.numel(), 8), 7, dtype=torch.bfloat16) - - monkeypatch.setattr(backbone, "_forward_with_block_offload", Mock(side_effect=block_forward)) - for name in ("to_cuda", "to_cpu", "init_block_offload", "release_block_offload_buffers"): - monkeypatch.setattr(backbone, name, Mock()) - for name in ("_forward_streaming_embedding", "load_streaming_layer", "release_disk_streaming_buffer"): - monkeypatch.setattr(backbone, name, Mock(wraps=getattr(backbone, name))) - - prompt = "A cat walking on the grass." - if fail_forward: - with pytest.raises(RuntimeError, match="injected forward failure"): - encoder.infer(prompt) - else: - result = encoder.infer(prompt) - assert set(result) == {"prompt_embeds", "text_token_tags"} - embeds, tags = result["prompt_embeds"], result["text_token_tags"] - assert isinstance(embeds, torch.Tensor) and isinstance(tags, torch.Tensor) - assert embeds.shape == (3, 8) and embeds.dtype == torch.bfloat16 - assert embeds.device.type == "cpu" and embeds.is_contiguous() - assert torch.isfinite(embeds).all() and torch.all(embeds == 7) - assert tags.shape == (3,) and tags.dtype == torch.long - assert tags.device == embeds.device - assert torch.all(tags == qwen_module.MINIMAX_H3_TEXT_TAG) - - encoder.tokenizer.assert_called_once_with(prompt, add_special_tokens=False) - assert backbone.to_cuda.call_count == int(mode == "resident") - assert backbone.to_cpu.call_count == int(mode == "resident") - assert backbone.init_block_offload.call_count == int(mode == "block") - assert backbone.release_block_offload_buffers.call_count == int(mode == "block" and release_buffers) - assert backbone.release_disk_streaming_buffer.call_count == int(mode == "disk" and release_buffers) - assert backbone._forward_with_block_offload.call_count == int(mode == "block") - if mode == "disk": - backbone._forward_streaming_embedding.assert_called_once() - assert [call.args[0] for call in backbone.load_streaming_layer.call_args_list] == [0, 1] - assert (backbone.streaming_layer is None) == release_buffers - assert backbone.embed_tokens.weight is None - assert not hasattr(backbone.embed_tokens, "pin_weight") - else: - backbone._forward_streaming_embedding.assert_not_called() - backbone.load_streaming_layer.assert_not_called() diff --git a/tests/models/minimax_h3/test_rope_precision.py b/tests/models/minimax_h3/test_rope_precision.py deleted file mode 100644 index 326d043b4..000000000 --- a/tests/models/minimax_h3/test_rope_precision.py +++ /dev/null @@ -1,65 +0,0 @@ -import importlib.util -import sys -import types -from pathlib import Path - -import pytest -import torch - - -@pytest.fixture() -def h3_rope(monkeypatch): - root = Path(__file__).parents[3] - spec = importlib.util.spec_from_file_location("h3_rope_weight_helpers", Path(__file__).with_name("test_transformer_disk_streaming.py")) - helper = importlib.util.module_from_spec(spec) - spec.loader.exec_module(helper) - _, weights, _ = helper.h3_modules.__wrapped__(monkeypatch) - - registry = sys.modules["lightx2v.utils.registry_factory"] - rope_registry = registry.ROPE_REGISTER - monkeypatch.setattr(registry, "ROPE_REGISTER", lambda name: lambda cls: cls) - magi = types.ModuleType("lightx2v.common.magi_custom_op_mode") - magi.use_magi_custom_ops = lambda: False - monkeypatch.setitem(sys.modules, magi.__name__, magi) - package = types.ModuleType("h3_precision_rope") - package.__path__ = [str(root / "lightx2v/common/ops/rope")] - monkeypatch.setitem(sys.modules, package.__name__, package) - spec = importlib.util.spec_from_file_location(package.__name__ + ".torch_rope", Path(package.__path__[0]) / "torch_rope.py") - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, spec.name, module) - spec.loader.exec_module(module) - monkeypatch.setattr(registry, "ROPE_REGISTER", rope_registry) - monkeypatch.setitem(rope_registry, "torch_real_rope", module.TorchRealRope) - monkeypatch.setitem(weights.ROPE_REGISTER, "torch_real_rope", module.TorchRealRope) - return weights.MiniMaxH3AttentionWeights("transformer_blocks.0.attn", {}).rope, module.TorchRealRope - - -def test_h3_precision_does_not_change_generic_default(h3_rope): - rope, generic = h3_rope - assert rope.layout == "split_half" - assert rope.compute_dtype == torch.bfloat16 - assert generic().compute_dtype == torch.float32 - - -@pytest.mark.parametrize("length", [1, 17, 52]) -def test_h3_bf16_rope_matches_reference_exactly(h3_rope, length): - rope, _ = h3_rope - generator = torch.Generator().manual_seed(123) - q, k = [torch.randn(length, 56, 128, generator=generator, dtype=torch.bfloat16) for _ in range(2)] - angles = torch.randn(length, 48, generator=generator, dtype=torch.float32) - angles = torch.cat((angles, angles), dim=-1) - cos, sin = angles.cos(), angles.sin() - - def reference(x): - rotary, passthrough = x[..., :96], x[..., 96:] - first, second = rotary.chunk(2, dim=-1) - rotated = torch.cat((-second, first), dim=-1) - out = rotary * cos.to(x.dtype)[:, None, :] + rotated * sin.to(x.dtype)[:, None, :] - return torch.cat((out, passthrough), dim=-1) - - actual = rope.apply(q, k, (cos, sin), rotary_dim=96) - for x, out in zip((q, k), actual): - assert out.shape == x.shape - assert out.dtype == torch.bfloat16 - assert torch.equal(out, reference(x)) - assert torch.equal(out[..., 96:], x[..., 96:]) diff --git a/tests/models/minimax_h3/test_runner_mps_low_memory.py b/tests/models/minimax_h3/test_runner_mps_low_memory.py deleted file mode 100644 index 8bc8f35af..000000000 --- a/tests/models/minimax_h3/test_runner_mps_low_memory.py +++ /dev/null @@ -1,326 +0,0 @@ -import importlib.util -import sys -import types -from pathlib import Path -from types import SimpleNamespace - -import pytest -import torch - -REPO_ROOT = Path(__file__).parents[3] - - -class _Profiler: - def __init__(self, *_args, **_kwargs): - pass - - def __call__(self, func): - return func - - def __enter__(self): - return self - - def __exit__(self, *_args): - return False - - -def _install_module(monkeypatch, name, **attrs): - module = types.ModuleType(name) - for attr_name, value in attrs.items(): - setattr(module, attr_name, value) - monkeypatch.setitem(sys.modules, name, module) - return module - - -def _load_runner_module(monkeypatch): - for package_name in [ - "lightx2v", - "lightx2v.models", - "lightx2v.models.audio_encoders", - "lightx2v.models.audio_encoders.hf", - "lightx2v.models.input_encoders", - "lightx2v.models.input_encoders.hf", - "lightx2v.models.networks", - "lightx2v.models.networks.minimax_h3", - "lightx2v.models.runners", - "lightx2v.models.runners.default_runner", - "lightx2v.models.schedulers", - "lightx2v.models.video_encoders", - "lightx2v.models.video_encoders.hf", - "lightx2v.models.video_encoders.hf.ltx2", - "lightx2v.models.video_encoders.hf.ltx2.audio_vae", - "lightx2v.server", - "lightx2v.utils", - "lightx2v_platform", - "lightx2v_platform.base", - ]: - package = types.ModuleType(package_name) - package.__path__ = [] - monkeypatch.setitem(sys.modules, package_name, package) - - request_name = "lightx2v.models.runners.request_fields" - request_spec = importlib.util.spec_from_file_location(request_name, REPO_ROOT / "lightx2v/models/runners/request_fields.py") - request_module = importlib.util.module_from_spec(request_spec) - monkeypatch.setitem(sys.modules, request_name, request_module) - request_spec.loader.exec_module(request_module) - - class DefaultRunner: - def __init__(self, config): - self.config = config - - def maybe_empty_cache(self, **_kwargs): - return False - - def end_run(self): - pass - - _install_module(monkeypatch, "lightx2v.models.runners.default_runner", DefaultRunner=DefaultRunner) - _install_module(monkeypatch, "lightx2v.models.audio_encoders.hf.minimax_h3", MiniMaxH3AudioVAE=object) - _install_module(monkeypatch, "lightx2v.models.input_encoders.hf.minimax_h3", MiniMaxH3Qwen3VLTextEncoder=object) - _install_module(monkeypatch, "lightx2v.models.networks.minimax_h3.lora", MiniMaxH3LoraAdapter=object) - _install_module(monkeypatch, "lightx2v.models.networks.minimax_h3.model", MiniMaxH3Model=object) - _install_module(monkeypatch, "lightx2v.models.schedulers.minimax_h3", MiniMaxH3Scheduler=object) - _install_module(monkeypatch, "lightx2v.models.video_encoders.hf.minimax_h3", MiniMaxH3VideoVAE=object) - - class Audio: - def __init__(self, **kwargs): - self.__dict__.update(kwargs) - - _install_module(monkeypatch, "lightx2v.models.video_encoders.hf.ltx2.audio_vae.ops", Audio=Audio) - - class _Metrics: - def __getattr__(self, _name): - return None - - _install_module(monkeypatch, "lightx2v.server.metrics", monitor_cli=_Metrics()) - _install_module(monkeypatch, "lightx2v.utils.envs", DTYPE_MAP={"fp32": torch.float32}, GET_RECORDER_MODE=lambda: None) - _install_module( - monkeypatch, - "lightx2v.utils.input_info", - INPUT_INFO_TYPES={task: object for task in ("t2av", "i2av", "l2av", "fl2av", "ref2av")}, - FL2AVInputInfo=object, - I2AVInputInfo=object, - L2AVInputInfo=object, - Ref2AVInputInfo=object, - T2AVInputInfo=object, - ) - _install_module(monkeypatch, "lightx2v.utils.ltx2_media_io", encode_video=lambda **_kwargs: None) - _install_module(monkeypatch, "lightx2v.utils.profiler", ProfilingContext4DebugL1=_Profiler, ProfilingContext4DebugL2=_Profiler) - _install_module(monkeypatch, "lightx2v.utils.registry_factory", RUNNER_REGISTER=lambda _name: lambda cls: cls) - _install_module(monkeypatch, "lightx2v_platform.base.global_var", AI_DEVICE="mps") - - packing_names = { - "TEXT_TAG": 1, - "align_num_frames": lambda value: value, - "prepare_keyframe_image": lambda image, *_args, **_kwargs: image, - "resolve_canvas_size": lambda width, height: (height, width), - "unpack_audio_tokens": lambda rows, *_args, **_kwargs: rows, - "unpatchify_video_tokens": lambda rows, *_args, **_kwargs: rows, - "validate_t2av_geometry": lambda *_args, **_kwargs: None, - } - _install_module(monkeypatch, "lightx2v.models.networks.minimax_h3.packing", **packing_names) - _install_module( - monkeypatch, - "lightx2v.models.networks.minimax_h3.packing_ref2av", - DEFAULT_REFERENCE_IMAGE_RESIZE_MODE="contain", - MAX_REFERENCES=12, - MAX_REFERENCE_AUDIOS=3, - MAX_REFERENCE_IMAGES=9, - MAX_REFERENCE_VIDEOS=3, - REFERENCE_IMAGE_RESIZE_MODES=("contain",), - MiniMaxH3PreparedReference=object, - decode_reference_audio=lambda *_args, **_kwargs: None, - decode_reference_video=lambda *_args, **_kwargs: None, - prepare_reference_frames=lambda frames, *_args, **_kwargs: frames, - prepare_reference_image=lambda image, *_args, **_kwargs: image, - prepare_reference_waveform=lambda waveform, *_args, **_kwargs: waveform, - resample_reference_frames=lambda frames, *_args, **_kwargs: frames, - resolve_reference_image_size=lambda width, height, **_kwargs: (height, width), - trim_reference_num_frames=lambda value: value, - ) - - module_path = REPO_ROOT / "lightx2v/models/runners/minimax_h3/minimax_h3_runner.py" - spec = importlib.util.spec_from_file_location("minimax_h3_runner_under_test", module_path) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - module.torch_device_module = SimpleNamespace(synchronize=lambda: None) - return module - - -@pytest.fixture() -def runner_module(monkeypatch): - return _load_runner_module(monkeypatch) - - -def _low_memory_config(**overrides): - config = { - "task": "t2av", - "dit_disk_streaming": True, - "text_encoder_disk_streaming": True, - "text_encoder_release_block_offload_buffers": True, - "warmup": False, - "cpu_offload": True, - } - config.update(overrides) - return config - - -def _make_runner(runner_module, config): - runner = object.__new__(runner_module.MiniMaxH3Runner) - runner.config = config - return runner - - -def test_low_memory_load_model_defers_vae(runner_module): - runner = _make_runner(runner_module, _low_memory_config()) - calls = [] - runner.load_transformer = lambda: calls.append("transformer") or object() - runner.load_text_encoder = lambda: calls.append("text_encoder") or [object()] - runner.load_vae = lambda: calls.append("vae") or (object(), object()) - - runner.load_model() - - assert calls == ["transformer", "text_encoder"] - assert runner.video_vae is None - assert runner.audio_vae is None - - -def test_non_low_memory_load_model_keeps_eager_vae(runner_module): - runner = _make_runner(runner_module, _low_memory_config(text_encoder_disk_streaming=False)) - video_vae = object() - audio_vae = object() - calls = [] - runner.load_transformer = lambda: calls.append("transformer") or object() - runner.load_text_encoder = lambda: calls.append("text_encoder") or [object()] - runner.load_vae = lambda: calls.append("vae") or (video_vae, audio_vae) - - runner.load_model() - - assert calls == ["transformer", "text_encoder", "vae"] - assert runner.video_vae is video_vae - assert runner.audio_vae is audio_vae - - -@pytest.mark.parametrize( - ("override", "message"), - [ - ({"text_encoder_release_block_offload_buffers": False}, "text_encoder_release_block_offload_buffers=true"), - ({"warmup": True}, "warmup=false"), - ], -) -def test_low_memory_load_model_rejects_unsupported_first_version_configs(runner_module, override, message): - runner = _make_runner(runner_module, _low_memory_config(**override)) - runner.load_transformer = lambda: object() - runner.load_text_encoder = lambda: [object()] - runner.load_vae = lambda: (object(), object()) - - with pytest.raises(ValueError, match=message): - runner.load_model() - - -def test_offload_transformer_releases_disk_streaming_buffer(runner_module): - runner = _make_runner(runner_module, _low_memory_config()) - calls = [] - runner.maybe_empty_cache = lambda **kwargs: calls.append(("empty", kwargs)) - runner.model = SimpleNamespace( - block_offload=True, - prepost_resident=False, - pre_weight=SimpleNamespace(to_cpu=lambda: calls.append("pre_cpu")), - post_weight=SimpleNamespace(to_cpu=lambda: calls.append("post_cpu")), - transformer_weights=SimpleNamespace(release_disk_streaming_buffer=lambda: calls.append("release_dit")), - ) - - runner._offload_transformer() - - assert calls == ["pre_cpu", "post_cpu", "release_dit", ("empty", {"force": True, "collect_garbage": True})] - - -def test_offload_transformer_preserves_regular_block_offload_behavior(runner_module): - runner = _make_runner(runner_module, _low_memory_config(dit_disk_streaming=False)) - calls = [] - runner.maybe_empty_cache = lambda **kwargs: calls.append(("empty", kwargs)) - runner.model = SimpleNamespace( - block_offload=True, - prepost_resident=False, - pre_weight=SimpleNamespace(to_cpu=lambda: calls.append("pre_cpu")), - post_weight=SimpleNamespace(to_cpu=lambda: calls.append("post_cpu")), - transformer_weights=SimpleNamespace(release_disk_streaming_buffer=lambda: calls.append("release_dit")), - ) - - runner._offload_transformer() - - assert calls == ["pre_cpu", "post_cpu", ("empty", {"force": True, "collect_garbage": True})] - - -def test_run_vae_decoder_lazy_loads_once(runner_module): - runner = _make_runner(runner_module, _low_memory_config()) - calls = [] - video_vae = SimpleNamespace( - decode_parallel=False, - decode=lambda latents: ("video", latents), - ) - audio_vae = SimpleNamespace(decode=lambda latents: ("audio", latents)) - runner.load_vae = lambda: calls.append("load_vae") or (video_vae, audio_vae) - runner.video_vae = None - runner.audio_vae = None - runner._vae_decode_tile_shapes = {} - runner.scheduler = SimpleNamespace( - num_condition_video_rows=0, - num_condition_audio_rows=0, - num_latent_frames=1, - latent_height=1, - latent_width=1, - num_audio_latents=1, - ) - - first = runner.run_vae_decoder(torch.tensor([1]), torch.tensor([2])) - second = runner.run_vae_decoder(torch.tensor([3]), torch.tensor([4])) - - assert calls == ["load_vae"] - assert first[0][0] == "video" - assert first[1][0] == "audio" - assert torch.equal(first[0][1], torch.tensor([1])) - assert torch.equal(first[1][1], torch.tensor([2])) - assert second[0][0] == "video" - assert second[1][0] == "audio" - assert torch.equal(second[0][1], torch.tensor([3])) - assert torch.equal(second[1][1], torch.tensor([4])) - - -def test_run_main_releases_vae_after_processing_result(runner_module): - runner = _make_runner(runner_module, _low_memory_config()) - calls = [] - runner.maybe_empty_cache = lambda **kwargs: calls.append(("empty", kwargs)) - runner.init_run = lambda: calls.append("init") - runner.run_segment = lambda _segment: calls.append("dit") or ("video_rows", "audio_rows") - runner._offload_transformer = lambda: calls.append("offload_dit") - runner.run_vae_decoder = lambda *_args: calls.append("decode") or ("decoded_video", "decoded_audio") - - def process(): - calls.append(("process", runner.video_vae, runner.audio_vae)) - return "result" - - runner.process_images_after_vae_decoder = process - runner.end_run = lambda: calls.append("end_run") - runner.video_vae = object() - runner.audio_vae = object() - - assert runner.run_main() == "result" - - assert [call if isinstance(call, str) else call[0] for call in calls] == [ - "init", - "dit", - "offload_dit", - "decode", - "process", - "empty", - "end_run", - ] - process_call = calls[4] - assert process_call[1] is not None - assert process_call[2] is not None - assert runner.video_vae is None - assert runner.audio_vae is None - assert runner.gen_video is None - assert runner.gen_audio is None diff --git a/tests/models/minimax_h3/test_scheduler_layout.py b/tests/models/minimax_h3/test_scheduler_layout.py deleted file mode 100644 index e917987e3..000000000 --- a/tests/models/minimax_h3/test_scheduler_layout.py +++ /dev/null @@ -1,69 +0,0 @@ -import importlib.util -import sys -import types -from pathlib import Path - -import pytest -import torch - -REPO_ROOT = Path(__file__).parents[3] - - -@pytest.fixture -def modules(monkeypatch): - # Keep geometry tests independent of accelerator initialization and media dependencies. - for name in ( - "lightx2v", - "lightx2v.models", - "lightx2v.models.networks", - "lightx2v.models.networks.minimax_h3", - "lightx2v.models.schedulers", - "lightx2v.models.schedulers.minimax_h3", - "lightx2v_platform", - "lightx2v_platform.base", - ): - module = types.ModuleType(name) - module.__path__ = [] - monkeypatch.setitem(sys.modules, name, module) - for name, attrs in ( - ("lightx2v.models.schedulers.scheduler", {"BaseScheduler": object}), - ("lightx2v_platform.base.global_var", {"AI_DEVICE": "cpu"}), - ("lightx2v.models.networks.minimax_h3.packing_ref2av", {"build_ref2av_packed_sequence": None}), - ): - module = types.ModuleType(name) - module.__dict__.update(attrs) - monkeypatch.setitem(sys.modules, name, module) - - def load(name): - spec = importlib.util.spec_from_file_location(name, REPO_ROOT / (name.replace(".", "/") + ".py")) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, name, module) - spec.loader.exec_module(module) - return module - - return load("lightx2v.models.networks.minimax_h3.packing"), load("lightx2v.models.schedulers.minimax_h3.scheduler") - - -@pytest.mark.parametrize("anchors", [(), ("first", "last")]) -@pytest.mark.parametrize("device", ["cpu", pytest.param("mps", marks=pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS unavailable"))]) -def test_layout_device_boundary_preserves_geometry_and_metadata(modules, anchors, device): - packing, scheduler = modules - layout = packing.build_packed_sequence(torch.ones(1, dtype=torch.long), 37, 2, 2, 207, keyframe_anchors=anchors) - assert layout.position_ids.device.type == "cpu" - assert layout.position_ids.dtype == torch.float64 - expected = layout.position_ids.to(torch.float32) - - actual = scheduler._layout_to_device(layout, device) - - assert actual.position_ids.device.type == device - assert actual.position_ids.dtype == torch.float32 - assert torch.isfinite(actual.position_ids).all() - assert torch.equal(actual.position_ids.cpu(), expected) - assert layout.position_ids.dtype == torch.float64 - for name in ("token_tags", "video_indices", "audio_indices", "text_indices"): - value = getattr(actual, name) - assert value.device.type == device - assert value.dtype == torch.long - assert torch.equal(value.cpu(), getattr(layout, name)) - for name in ("sequence_length", "num_condition_video_rows", "num_condition_audio_rows"): - assert getattr(actual, name) == getattr(layout, name) diff --git a/tests/models/minimax_h3/test_streaming_lora.py b/tests/models/minimax_h3/test_streaming_lora.py deleted file mode 100644 index 518f0d22d..000000000 --- a/tests/models/minimax_h3/test_streaming_lora.py +++ /dev/null @@ -1,383 +0,0 @@ -import ast -import importlib.util -import sys -import types -from abc import ABCMeta, abstractmethod -from pathlib import Path - -import pytest -import torch -from loguru import logger -from safetensors.torch import save_file - -ROOT = Path(__file__).parents[3] - - -def load_module(name, path, monkeypatch): - spec = importlib.util.spec_from_file_location(name, path) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, name, module) - spec.loader.exec_module(module) - return module - - -@pytest.fixture -def setup(monkeypatch): - helpers = load_module("h3_stream_lora_test_helpers", Path(__file__).with_name("test_model_disk_streaming.py"), monkeypatch) - modules = helpers.h3_model_modules.__wrapped__(monkeypatch) - monkeypatch.setattr(modules[3], "AI_DEVICE", "cpu") - model = modules[-1].MiniMaxH3Model - streaming = load_module("lightx2v.models.networks.minimax_h3.streaming_lora", ROOT / "lightx2v/models/networks/minimax_h3/streaming_lora.py", monkeypatch) - # Execute the actual production MMWeight classes, excluding optional CUDA - # kernel imports. Only base checkpoint I/O is replaced by the existing tiny - # fixture; register_lora/apply/apply_lora/remove_lora remain production code. - scope = {"torch": torch, "ABCMeta": ABCMeta, "abstractmethod": abstractmethod, "logger": logger, "AI_DEVICE": "cpu"} - for file, names in [ - ("lightx2v/common/ops/utils.py", {"build_lora_and_diff_names"}), - ("lightx2v/common/ops/mm/mm_weight.py", {"MMWeightTemplate", "MMWeight"}), - ]: - tree = ast.parse((ROOT / file).read_text()) - nodes = [node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in names] - for node in nodes: - node.decorator_list = [] - exec(compile(ast.Module(body=nodes, type_ignores=[]), file, "exec"), scope) # noqa: S102 - execute trusted local production classes - - class Linear(scope["MMWeight"]): - load = helpers._FakeLinear.load - load_state_dict = helpers._FakeLinear.load_state_dict - - registry = sys.modules["lightx2v.utils.registry_factory"].MM_WEIGHT_REGISTER - monkeypatch.setitem(registry, "Default", Linear) - monkeypatch.setitem(registry, "Default-ForceFp32", Linear) - return helpers, modules, model, streaming, Linear - - -def fixture_tensors(): - return { - "base_model.model.transformer_blocks.0.attn.to_q.lora_A.default.weight": torch.tensor([[1, 2, 3]], dtype=torch.bfloat16), - "base_model.model.transformer_blocks.0.attn.to_q.lora_B.default.weight": torch.tensor([[1], [2]], dtype=torch.bfloat16), - "transformer_blocks.1.attn.to_k.lora_down.weight": torch.tensor([[2, 1, 0], [0, 1, 2]], dtype=torch.bfloat16), - "transformer_blocks.1.attn.to_k.lora_up.weight": torch.tensor([[1, 0], [0, 1]], dtype=torch.bfloat16), - "transformer_blocks.1.attn.to_k.alpha": torch.tensor(8.0), - "token_refiner.refiner_blocks.0.attn.to_q.lora_A.weight": torch.ones(1, 3, dtype=torch.bfloat16), - "token_refiner.refiner_blocks.0.attn.to_q.lora_B.weight": torch.ones(2, 1, dtype=torch.bfloat16), - } - - -def make_index(tmp_path, setup, tensors=None, **kwargs): - _, _, model, streaming, _ = setup - tensors = fixture_tensors() if tensors is None else tensors - path = tmp_path / "lora.safetensors" - save_file(tensors, path) - shapes = {model._normalize_dynamic_lora_key(key).removesuffix(".lora_down.weight") + ".weight": (2, 3) for key in fixture_tensors() if "lora_A" in key or "lora_down" in key} - kwargs = {"alpha": 4, "strength": 0.5, **kwargs} - return streaming.MiniMaxH3StreamingLora(path, normalize_key=model._normalize_dynamic_lora_key, target_shapes=shapes, **kwargs) - - -def test_index_metadata_only_and_alpha_precedence(tmp_path, setup, monkeypatch): - streaming = setup[3] - original = streaming.safe_open - reads = [] - - class Reader: - def __init__(self, *args, **kwargs): - self.context = original(*args, **kwargs) - - def __enter__(self): - self.source = self.context.__enter__() - return self - - def __exit__(self, *args): - return self.context.__exit__(*args) - - def keys(self): - return self.source.keys() - - def get_slice(self, key): - return self.source.get_slice(key) - - def get_tensor(self, key): - reads.append(key) - assert key.endswith(".alpha"), "Index construction read a factor tensor" - return self.source.get_tensor(key) - - monkeypatch.setattr(streaming, "safe_open", Reader) - index = make_index(tmp_path, setup) - assert len(index.pairs) == 3 - assert set(index.blocks) == {0, 1} - assert len(index.resident) == 1 - assert index.pairs["transformer_blocks.0.attn.to_q.weight"].alpha == 4 - assert index.pairs["transformer_blocks.1.attn.to_k.weight"].alpha == 8 - assert {pair.rank for pair in index.pairs.values()} == {1, 2} - assert reads == ["transformer_blocks.1.attn.to_k.alpha"] - assert not index._resident_cpu - - -@pytest.mark.parametrize("problem", ["missing", "shape", "orphan_alpha", "collision", "unsupported", "alpha", "rank", "factor_dtype"]) -def test_invalid_checkpoint_rejected(tmp_path, setup, problem): - tensors = fixture_tensors() - a = next(iter(tensors)) - if problem == "missing": - tensors.pop(a) - elif problem == "shape": - tensors[a] = torch.ones(1, 4) - elif problem == "orphan_alpha": - tensors["unknown.alpha"] = torch.tensor(1.0) - elif problem == "collision": - tensors["transformer_blocks.0.attn.to_q.lora_down.weight"] = tensors[a].clone() - elif problem == "unsupported": - tensors["unknown"] = torch.tensor(1.0) - elif problem == "alpha": - tensors["transformer_blocks.1.attn.to_k.alpha"] = torch.tensor(float("nan")) - elif problem == "rank": - tensors[a] = torch.ones(2, 3) - elif problem == "factor_dtype": - tensors[a] = tensors[a].long() - with pytest.raises(ValueError): - make_index(tmp_path, setup, tensors) - - -@pytest.mark.parametrize("alpha", [None, 0, -1, float("inf")]) -def test_missing_or_invalid_config_alpha(tmp_path, setup, alpha): - with pytest.raises(ValueError, match="alpha"): - make_index(tmp_path, setup, alpha=alpha) - - -def test_selective_reads_reuse_math_release_and_refiner(tmp_path, setup, monkeypatch): - helpers, modules, Model, streaming, _ = setup - _, pre_module, post_module, transformer_module, *_ = modules - helpers._write_fake_checkpoint(tmp_path, pre_module, post_module, transformer_module) - index = make_index(tmp_path, setup) - config = helpers._config(tmp_path, lora_dynamic_apply=True) - model = Model(str(tmp_path), config, torch.device("cpu"), lora_path=index.path, lora_strength=0.5, lora_alpha=4) - weights = model.transformer_weights - index = weights.streaming_lora - assert len(index.pairs) == 3 - block = weights.streaming_block - base_pointers = {name: weight.weight.data_ptr() for name, weight in index.weights(block).items()} - original = streaming.safe_open - reads = [] - - class Reader: - def __init__(self, *args, **kwargs): - self.context = original(*args, **kwargs) - - def __enter__(self): - self.source = self.context.__enter__() - return self - - def __exit__(self, *args): - return self.context.__exit__(*args) - - def get_tensor(self, key): - reads.append(key) - return self.source.get_tensor(key) - - monkeypatch.setattr(streaming, "safe_open", Reader) - monkeypatch.setattr(torch.Tensor, "pin_memory", lambda *_args, **_kwargs: pytest.fail("streamed LoRA must not pin factors")) - for i, name in [(0, "transformer_blocks.0.attn.to_q.weight"), (1, "transformer_blocks.0.attn.to_k.weight")]: - reads.clear() - assert weights.load_streaming_block(i) is block - linears = index.weights(block) - assert {key for key, value in linears.items() if value.has_lora_branch} == {name} - assert len(reads) == 2 - assert all(f"transformer_blocks.{i}." in key for key in reads) - assert {key: value.weight.data_ptr() for key, value in linears.items()} == base_pointers - linear = linears[name] - x = torch.tensor([[1, 2, 3], [-1, 0, 2]], dtype=torch.bfloat16) - pair = next(iter(index.blocks[i].values())) - factors = fixture_tensors() - expected = x @ linear.weight + 0.5 * (pair.alpha / pair.rank) * ((x @ factors[pair.down_key].T) @ factors[pair.up_key].T) - assert torch.equal(linear.apply(x), expected) - assert not linear.lora_down.is_pinned() - refiner = model.pre_weight - # Tiny fixture's to_cuda is a no-op: explicitly activate its CPU weight. - for linear in index.weights(refiner).values(): - if getattr(linear, "weight", None) is None: - linear.weight = linear.pin_weight - reads.clear() - for _ in range(2): - with index.resident_scope(refiner): - assert sum(weight.has_lora_branch for weight in index.weights(refiner).values()) == 1 - assert not any(weight.has_lora_branch for weight in index.weights(refiner).values()) - assert len(reads) == 2 # Cached on CPU only, read once across evaluations. - assert all(tensor.device.type == "cpu" for tensors in index._resident_cpu.values() for tensor in tensors) - weights.release_disk_streaming_buffer() - assert weights.streaming_block is None - assert not any(weight.has_lora_branch for weight in index.weights(block).values()) - assert all(not hasattr(weight, "lora_down") and not hasattr(weight, "lora_scale") for weight in index.weights(block).values()) - new_block = weights.load_streaming_block(0) - assert new_block is not block - assert sum(weight.has_lora_branch for weight in index.weights(new_block).values()) == 1 - weights.release_disk_streaming_buffer() - - -def test_resident_cleanup_on_failure(tmp_path, setup): - index = make_index(tmp_path, setup) - Linear = setup[-1] - weight = Linear("token_refiner.refiner_blocks.0.attn.to_q.weight", lora_prefix="token_refiner") - weight.weight = torch.ones(3, 2, dtype=torch.bfloat16) - with pytest.raises(RuntimeError, match="pre-infer failed"), index.resident_scope(weight): - assert weight.has_lora_branch - raise RuntimeError("pre-infer failed") - assert not weight.has_lora_branch - - -def test_merged_streaming_rejected_before_loading(tmp_path, setup): - helpers, _, Model, _, _ = setup - with pytest.raises(NotImplementedError, match="dynamic LoRA"): - Model(str(tmp_path), helpers._config(tmp_path), torch.device("cpu"), lora_path="unused.safetensors") - - -def test_no_lora_streaming_unchanged(tmp_path, setup): - helpers, modules, Model, _, _ = setup - _, pre, post, transformer, *_ = modules - helpers._write_fake_checkpoint(tmp_path, pre, post, transformer) - model = Model(str(tmp_path), helpers._config(tmp_path), torch.device("cpu")) - assert model.transformer_weights.streaming_lora is None - block = model.transformer_weights.load_streaming_block(1) - assert block is model.transformer_weights.load_streaming_block(0) - model.transformer_weights.release_disk_streaming_buffer() - - -def test_official_raw_adapter_and_lora_compose(tmp_path, setup, monkeypatch): - helpers, _modules, Model, streaming, _ = setup - raw_helpers = load_module("h3_stream_lora_raw_helpers", Path(__file__).with_name("test_checkpoint_adapter.py"), monkeypatch) - _, raw, _ = raw_helpers.write_raw(tmp_path) - tensors = {} - for i in (0, 1): - tensors[f"transformer_blocks.{i}.attn.to_q.lora_A.weight"] = torch.ones(1, 3, dtype=torch.bfloat16) * (i + 1) - tensors[f"transformer_blocks.{i}.attn.to_q.lora_B.weight"] = torch.ones(4, 1, dtype=torch.bfloat16) - path = tmp_path / "raw_lora.safetensors" - save_file(tensors, path) - model = Model(str(tmp_path), helpers._config(tmp_path, lora_dynamic_apply=True), torch.device("cpu"), lora_path=str(path), lora_alpha=4) - weights = model.transformer_weights - assert weights.checkpoint.selected_reader is not None - pointer = None - for i in (0, 1, 0): - block = weights.load_streaming_block(i) - q = streaming.MiniMaxH3StreamingLora.weights(block)["transformer_blocks.0.attn.to_q.weight"] - if pointer is None: - pointer = q.weight.data_ptr() - assert q.weight.data_ptr() == pointer - raw_qkv = raw[f"blocks.{i}.attn.qkv_proj.weight"] - expected_base = torch.vstack((raw_qkv[:2], raw_qkv[6:8])).T - assert torch.equal(q.weight, expected_base) - x = torch.tensor([[1, 0, -1], [1, 2, 3]], dtype=torch.bfloat16) - a, b = tensors[f"transformer_blocks.{i}.attn.to_q.lora_A.weight"], tensors[f"transformer_blocks.{i}.attn.to_q.lora_B.weight"] - assert torch.equal(q.apply(x), x @ expected_base + 4 * ((x @ a.T) @ b.T)) - # FP32 sensitive heads are not silently given BF16 LoRA factors. - shapes = streaming.streaming_target_shapes(weights.checkpoint, block, model.pre_weight, model.post_weight) - assert "proj_in.weight" not in shapes - assert "proj_out.weight" not in shapes - weights.release_disk_streaming_buffer() - - -def test_ordinary_dynamic_loader_and_mmweight_contract(tmp_path, setup, monkeypatch): - _, _, Model, _, Linear = setup - index = make_index(tmp_path, setup) - model = object.__new__(Model) - model.config = {"dit_disk_streaming": False} - model.device = torch.device("cpu") - model.lora_alpha = 4 - model._h3_weight_shapes = {name: (2, 3) for name in index.pairs} - model.use_tp = False - # The unchanged ordinary loader pins CPU tensors; avoid requiring a GPU - # in this regression and verify that it still takes its original path. - pins = [] - monkeypatch.setattr(torch.Tensor, "pin_memory", lambda tensor: pins.append(tensor.shape) or tensor) - loaded = model._load_lora_file(index.path) - assert pins - q = Linear("transformer_blocks.0.attn.to_q.weight", lora_prefix="transformer_blocks") - q.weight = torch.ones(3, 2, dtype=torch.bfloat16) - q.register_lora(loaded, 0.5) - x = torch.ones(2, 3, dtype=torch.bfloat16) - expected = q.apply(x) - q.remove_lora() - index.load_block(q, 0) - assert torch.equal(q.apply(x), expected) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_dynamic_math_non_integer_scaling(tmp_path, setup, dtype): - _, _, _, _, Linear = setup - gen = torch.Generator().manual_seed(123) - tensors = { - "transformer_blocks.0.attn.to_q.lora_down.weight": torch.randn(3, 3, generator=gen, dtype=dtype), - "transformer_blocks.0.attn.to_q.lora_up.weight": torch.randn(2, 3, generator=gen, dtype=dtype), - } - index = make_index(tmp_path, setup, tensors, dtype=dtype, alpha=5, strength=0.7) - weight = Linear("transformer_blocks.0.attn.to_q.weight", lora_prefix="transformer_blocks") - weight.weight = torch.randn(3, 2, generator=gen, dtype=dtype) - x = torch.randn(7, 3, generator=gen, dtype=dtype) - index.load_block(weight, 0) - a, b = tensors.values() - # Preserve the released dynamic path's scalar dtype and operation order. - scale = torch.tensor(5, dtype=dtype) / 3 - expected = x @ weight.weight + 0.7 * scale * ((x @ a.T) @ b.T) - assert torch.equal(weight.apply(x), expected) - index.clear(weight) - - -def test_production_pre_infer_scope_and_error_cleanup(tmp_path, setup): - helpers, modules, Model, _streaming, _ = setup - _, pre, post, transformer, *_ = modules - helpers._write_fake_checkpoint(tmp_path, pre, post, transformer) - index = make_index(tmp_path, setup) - model = Model(str(tmp_path), helpers._config(tmp_path, lora_dynamic_apply=True), torch.device("cpu"), lora_path=index.path, lora_alpha=4) - for weight in index.weights(model.pre_weight).values(): - weight.weight = weight.pin_weight - calls = [] - - def pre_infer(root, prompt): - assert any(weight.has_lora_branch for weight in index.weights(root).values()) - calls.append("pre") - return prompt - - def transformer_infer(weights, pre): - assert not any(weight.has_lora_branch for weight in index.weights(model.pre_weight).values()) - calls.append("transformer") - return pre - - def post_infer(root, hidden, pre): - calls.append("post") - return hidden - - model.pre_infer = types.SimpleNamespace(infer=pre_infer) - model.transformer_infer = types.SimpleNamespace(infer=transformer_infer) - model.post_infer = types.SimpleNamespace(infer=post_infer) - prompt = torch.ones(1, 3, dtype=torch.bfloat16) - inputs = {"text_encoder_output": {"prompt_embeds": prompt}} - assert model._infer_cond_uncond(inputs) is prompt - assert calls == ["pre", "transformer", "post"] - - def fail(root, prompt): - pre_infer(root, prompt) - raise RuntimeError("pre-infer failed") - - model.pre_infer.infer = fail - with pytest.raises(RuntimeError, match="pre-infer failed"): - model._infer_cond_uncond(inputs) - assert not any(weight.has_lora_branch for weight in index.weights(model.pre_weight).values()) - with pytest.raises(NotImplementedError, match="selective LoRA index"): - model._load_lora_file(index.path) - with pytest.raises(NotImplementedError, match="runtime adapter switching"): - model._update_lora(index.path, 1) - model.transformer_weights.release_disk_streaming_buffer() - - -def test_ordinary_merged_lora_contract(tmp_path, setup, monkeypatch): - index = make_index(tmp_path, setup) - adapter_base = types.ModuleType("lightx2v.models.networks.lora_adapter") - adapter_base.LoraAdapter = object - monkeypatch.setitem(sys.modules, adapter_base.__name__, adapter_base) - module = load_module("h3_ordinary_lora_regression", ROOT / "lightx2v/models/networks/minimax_h3/lora.py", monkeypatch) - adapter = module.MiniMaxH3LoraAdapter() - base = {name: torch.ones(2, 3, dtype=torch.bfloat16) for name in index.pairs} - adapter.model = types.SimpleNamespace(config={"lora_merge_device": "cpu"}, use_tp=False, original_weight_dict=base) - assert adapter._merge_file(index.path, strength=0.5, alpha=4) == 3 - fixture = fixture_tensors() - for name, pair in index.pairs.items(): - expected = torch.ones(2, 3, dtype=torch.bfloat16) - expected.add_(fixture[pair.up_key] @ fixture[pair.down_key], alpha=0.5 * pair.alpha / pair.rank) - assert torch.equal(base[name], expected) diff --git a/tests/models/minimax_h3/test_transformer_disk_streaming.py b/tests/models/minimax_h3/test_transformer_disk_streaming.py deleted file mode 100644 index 1ba80ddd9..000000000 --- a/tests/models/minimax_h3/test_transformer_disk_streaming.py +++ /dev/null @@ -1,419 +0,0 @@ -import importlib.util -import json -import sys -import types -from pathlib import Path -from types import SimpleNamespace - -import pytest -import torch -from safetensors.torch import save_file - -REPO_ROOT = Path(__file__).parents[3] - - -class _FakeWeightModule: - def __init__(self): - self._modules = {} - self._parameters = {} - - def add_module(self, name, module): - self._modules[name] = module - setattr(self, name, module) - - def load(self, weight_dict): - for module in self._modules.values(): - if hasattr(module, "load"): - module.load(weight_dict) - - def load_state_dict(self, destination, block_index, adapter_block_index=None): - for module in self._modules.values(): - if hasattr(module, "load_state_dict"): - module.load_state_dict(destination, block_index, adapter_block_index) - - -class _FakeWeightModuleList(_FakeWeightModule): - def __init__(self, modules=None): - super().__init__() - self._list = [] - if modules is not None: - for module in modules: - self.append(module) - - def append(self, module): - self._list.append(module) - self.add_module(str(len(self._list) - 1), module) - - def __getitem__(self, index): - return self._list[index] - - def __len__(self): - return len(self._list) - - def __iter__(self): - return iter(self._list) - - -def _resolve_block_name(name, block_index): - parts = name.split(".", 2) - if len(parts) == 3 and parts[0] == "transformer_blocks" and parts[1].isdigit(): - return f"transformer_blocks.{int(block_index)}.{parts[2]}" - return name - - -class _FakeLinear: - def __init__(self, weight_name, bias_name=None, create_cuda_buffer=False, **_kwargs): - self.weight_name = weight_name - self.bias_name = bias_name - self.create_cuda_buffer = create_cuda_buffer - self.base_attrs = [(weight_name, "weight", True)] - if bias_name is not None: - self.base_attrs.append((bias_name, "bias", False)) - - def load(self, weight_dict): - for name, attr_name, transpose in self.base_attrs: - tensor = weight_dict[name] - if transpose: - tensor = tensor.t() - if self.create_cuda_buffer: - setattr(self, f"{attr_name}_cuda_buffer", tensor.clone()) - else: - setattr(self, attr_name, tensor.clone()) - - def load_state_dict(self, destination, block_index, adapter_block_index=None): - for name, attr_name, _transpose in self.base_attrs: - actual_name = _resolve_block_name(name, block_index) - if actual_name in destination: - buffer = getattr(self, f"{attr_name}_cuda_buffer") - setattr(self, attr_name, buffer.copy_(destination[actual_name])) - - -class _FakeRMS: - def __init__(self, weight_name, create_cuda_buffer=False, **_kwargs): - self.weight_name = weight_name - self.create_cuda_buffer = create_cuda_buffer - self.base_attrs = [(weight_name, "weight", False)] - - def load(self, weight_dict): - tensor = weight_dict[self.weight_name] - if self.create_cuda_buffer: - self.weight_cuda_buffer = tensor.clone() - else: - self.weight = tensor.clone() - - def load_state_dict(self, destination, block_index, adapter_block_index=None): - actual_name = _resolve_block_name(self.weight_name, block_index) - if actual_name in destination: - self.weight = self.weight_cuda_buffer.copy_(destination[actual_name]) - - -class _FakeLeaf: - base_attrs = () - - def __init__(self, *_args, **_kwargs): - pass - - def set_config(self, *_args, **_kwargs): - pass - - -def _load_module(module_name, relative_path): - module_path = REPO_ROOT / relative_path - spec = importlib.util.spec_from_file_location(module_name, module_path) - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) - return module - - -@pytest.fixture() -def h3_modules(monkeypatch): - for package_name in [ - "lightx2v", - "lightx2v.common", - "lightx2v.common.modules", - "lightx2v.common.transformer_infer", - "lightx2v.models", - "lightx2v.models.networks", - "lightx2v.models.networks.minimax_h3", - "lightx2v.models.networks.minimax_h3.infer", - "lightx2v.utils", - ]: - package = types.ModuleType(package_name) - package.__path__ = [] - monkeypatch.setitem(sys.modules, package_name, package) - - # Load main's FP8 helpers without importing the full runtime package. - for module_name in ( - "lightx2v.common.ops.mm.fp8_f16_accum", - "lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy", - ): - spec = importlib.util.spec_from_file_location(module_name, REPO_ROOT / (module_name.replace(".", "/") + ".py")) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, module_name, module) - spec.loader.exec_module(module) - - # Cache IO is outside these isolated model/streaming tests. - cache = types.ModuleType("lightx2v.models.networks.minimax_h3.adaln_cache") - cache.validate_adaln_cache_config = lambda config: None - cache.load_persistent_adaln_cache = lambda config, device: ({}, {}) - monkeypatch.setitem(sys.modules, cache.__name__, cache) - - weight_module = types.ModuleType("lightx2v.common.modules.weight_module") - weight_module.WeightModule = _FakeWeightModule - weight_module.WeightModuleList = _FakeWeightModuleList - monkeypatch.setitem(sys.modules, "lightx2v.common.modules.weight_module", weight_module) - - registry = types.ModuleType("lightx2v.utils.registry_factory") - registry.MM_WEIGHT_REGISTER = {"Default": _FakeLinear} - registry.RMS_WEIGHT_REGISTER = {"torch_native": _FakeRMS} - registry.ROPE_REGISTER = {"torch_real_rope": _FakeLeaf} - registry.ATTN_WEIGHT_REGISTER = {"flash_attn3": _FakeLeaf} - monkeypatch.setitem(sys.modules, "lightx2v.utils.registry_factory", registry) - - triton_ops = types.ModuleType("lightx2v.models.networks.minimax_h3.infer.triton_ops") - triton_ops.MiniMaxH3TritonRope = _FakeLeaf - monkeypatch.setitem(sys.modules, "lightx2v.models.networks.minimax_h3.infer.triton_ops", triton_ops) - - checkpoint_module = _load_module( - "lightx2v.models.networks.minimax_h3.checkpoint", - "lightx2v/models/networks/minimax_h3/checkpoint.py", - ) - weights_module = _load_module( - "minimax_h3_transformer_weights_under_test", - "lightx2v/models/networks/minimax_h3/weights/transformer_weights.py", - ) - - base_infer = types.ModuleType("lightx2v.common.transformer_infer.transformer_infer") - - class _BaseTransformerInfer: - def init_compile(self, config): - self.use_compile = config.get("use_compile", False) - - def run_block(self, block_idx, block, *args): - return self.infer_block(block, *args) - - base_infer.BaseTransformerInfer = _BaseTransformerInfer - monkeypatch.setitem(sys.modules, "lightx2v.common.transformer_infer.transformer_infer", base_infer) - - envs = types.ModuleType("lightx2v.utils.envs") - envs.GET_DTYPE = lambda: torch.float32 - monkeypatch.setitem(sys.modules, "lightx2v.utils.envs", envs) - infer_module = _load_module( - "minimax_h3_transformer_infer_under_test", - "lightx2v/models/networks/minimax_h3/infer/transformer_infer.py", - ) - - return checkpoint_module, weights_module, infer_module - - -def _iter_base_attrs(module): - if hasattr(module, "base_attrs"): - yield from module.base_attrs - for child in getattr(module, "_modules", {}).values(): - yield from _iter_base_attrs(child) - - -def _block_tensors_from_template(block, block_index, value): - tensors = {} - for name, _attr_name, transpose in _iter_base_attrs(block): - actual_name = _resolve_block_name(name, block_index) - if transpose: - tensor = torch.full((2, 3), value, dtype=torch.float32) - elif actual_name.endswith(".bias"): - tensor = torch.full((2,), value, dtype=torch.float32) - else: - tensor = torch.full((3,), value, dtype=torch.float32) - tensors[actual_name] = tensor - return tensors - - -def _write_fake_official_checkpoint(tmp_path, weights_module, num_layers=2): - template = weights_module.MiniMaxH3TransformerBlockWeights(0, {"num_layers": num_layers}) - block_tensors = {} - for block_index in range(num_layers): - block_tensors.update(_block_tensors_from_template(template, block_index, block_index + 1)) - - names = sorted(block_tensors) - shard_1_names = set(names[::2]) - shard_1 = {name: block_tensors[name] for name in names if name in shard_1_names} - shard_2 = {name: block_tensors[name] for name in names if name not in shard_1_names} - save_file(shard_1, tmp_path / "model-00001-of-00002.safetensors") - save_file(shard_2, tmp_path / "model-00002-of-00002.safetensors") - - weight_map = { - **{name: "model-00001-of-00002.safetensors" for name in shard_1}, - **{name: "model-00002-of-00002.safetensors" for name in shard_2}, - } - (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"metadata": {"total_size": 0}, "weight_map": weight_map}), - encoding="utf-8", - ) - return block_tensors - - -def test_transformer_weights_stream_official_shards_one_block_at_a_time(tmp_path, monkeypatch, h3_modules): - checkpoint_module, weights_module, _infer_module = h3_modules - _write_fake_official_checkpoint(tmp_path, weights_module, num_layers=2) - requested_names = [] - - class SpyCheckpoint(checkpoint_module.MiniMaxH3ShardCheckpoint): - def load_tensors(self, names, device="cpu"): - requested_names.append(tuple(names)) - return super().load_tensors(names, device=device) - - monkeypatch.setattr(weights_module, "MiniMaxH3ShardCheckpoint", SpyCheckpoint) - - weights = weights_module.MiniMaxH3TransformerWeights( - { - "dit_disk_streaming": True, - "dit_original_ckpt": str(tmp_path), - "num_layers": 2, - "dit_quantized": False, - "tensor_parallel": False, - } - ) - - assert weights.disk_streaming is True - assert len(weights.blocks) == 0 - assert weights.checkpoint.block_indices == (0, 1) - assert weights.streaming_block_indices == (0, 1) - assert weights.streaming_block.attn.to_q.weight_cuda_buffer.device.type == "cpu" - - block0_names = weights.checkpoint.tensor_names_for_block(0) - block1_names = weights.checkpoint.tensor_names_for_block(1) - assert requested_names == [block0_names] - - block0 = weights.load_streaming_block(0) - block_id = id(block0) - buffer = block0.attn.to_q.weight_cuda_buffer - buffer_id = id(buffer) - assert torch.all(block0.attn.to_q.weight == 1) - assert requested_names[-1] == block0_names - - block1 = weights.load_streaming_block(1) - assert id(block1) == block_id - assert id(block1.attn.to_q.weight_cuda_buffer) == buffer_id - assert torch.all(block1.attn.to_q.weight == 2) - assert block1.attn.to_q.weight.shape == (3, 2) - assert requested_names[-1] == block1_names - assert all(set(names) in [set(block0_names), set(block1_names)] for names in requested_names) - - -def test_transformer_weights_release_and_reinitialize_streaming_block(tmp_path, h3_modules): - _checkpoint_module, weights_module, _infer_module = h3_modules - _write_fake_official_checkpoint(tmp_path, weights_module, num_layers=2) - - weights = weights_module.MiniMaxH3TransformerWeights( - { - "dit_disk_streaming": True, - "dit_original_ckpt": str(tmp_path), - "num_layers": 2, - "dit_quantized": False, - "tensor_parallel": False, - } - ) - old_block = weights.streaming_block - assert old_block.attn.to_q.weight_cuda_buffer is not None - - weights.release_disk_streaming_buffer() - - assert weights.streaming_block is None - assert old_block.attn.to_q.weight is None - assert old_block.attn.to_q.weight_cuda_buffer is None - assert old_block.adaln.bias is None - assert old_block.adaln.bias_cuda_buffer is None - assert old_block.norm1.weight is None - assert old_block.norm1.weight_cuda_buffer is None - - block0 = weights.load_streaming_block(0) - assert block0 is weights.streaming_block - assert block0 is not old_block - assert torch.all(block0.attn.to_q.weight == 1) - assert block0.attn.to_q.weight_cuda_buffer is not None - - weights.release_disk_streaming_buffer() - assert weights.streaming_block is None - - block1 = weights.load_streaming_block(1) - assert block1 is weights.streaming_block - assert block1 is not block0 - assert torch.all(block1.attn.to_q.weight == 2) - assert block1.attn.to_q.weight_cuda_buffer is not None - - -def test_transformer_disk_streaming_rejects_missing_block(tmp_path, h3_modules): - _checkpoint_module, weights_module, _infer_module = h3_modules - _write_fake_official_checkpoint(tmp_path, weights_module, num_layers=1) - - with pytest.raises(ValueError, match="checkpoint block indices mismatch"): - weights_module.MiniMaxH3TransformerWeights( - { - "dit_disk_streaming": True, - "dit_original_ckpt": str(tmp_path), - "num_layers": 2, - "dit_quantized": False, - "tensor_parallel": False, - } - ) - - -def test_transformer_infer_dispatches_to_disk_streaming(h3_modules): - _checkpoint_module, _weights_module, infer_module = h3_modules - infer = infer_module.MiniMaxH3TransformerInfer({"num_attention_heads": 1, "use_adaln_cache": False}) - loaded = [] - ran = [] - - class FakeBlockWeights: - disk_streaming = True - checkpoint = SimpleNamespace(block_indices=(0, 1)) - - def load_streaming_block(self, block_index): - loaded.append(block_index) - return f"block-{block_index}" - - def run_block(block_index, block, hidden_states, pre_infer_out): - ran.append((block_index, block, hidden_states)) - return hidden_states + block_index + 1 - - infer.run_block = run_block - pre_infer_out = SimpleNamespace(hidden_states=0) - - assert infer.infer(FakeBlockWeights(), pre_infer_out) == 3 - assert loaded == [0, 1] - assert ran == [(0, "block-0", 0), (1, "block-1", 1)] - - -def test_persistent_cache_prepared_before_streaming_and_survives_clear(h3_modules, monkeypatch): - _, weights_module, infer_module = h3_modules - monkeypatch.setattr(infer_module, "AI_DEVICE", "mps") - assert infer_module.MiniMaxH3TransformerInfer._cache_device() == torch.device("mps") - infer = infer_module.MiniMaxH3TransformerInfer({"num_attention_heads": 1, "use_adaln_cache": True}) - infer.scheduler = SimpleNamespace(unique_timesteps_cpu=torch.tensor([0.5])) - tables = [torch.tensor([1.0]), torch.tensor([2.0])] - norm_out = torch.tensor([3.0]) - infer._adaln_cache[(0.5,)] = tables - infer._norm_out_cache[(0.5,)] = norm_out - pre = SimpleNamespace(hidden_states=0, temb=None) - - class Blocks: - disk_streaming = True - checkpoint = SimpleNamespace(block_indices=(0, 1)) - - def load_streaming_block(self, index): - assert pre.norm_out_modulation is norm_out - assert infer._get_cached_adaln(index) is tables[index] - return index - - infer.run_block = lambda index, block, hidden, pre: hidden + block + 1 - assert infer.infer(Blocks(), pre) == 3 - infer._clear_adaln_cache() - assert infer._adaln_cache[(0.5,)] is tables - assert infer.infer(Blocks(), pre) == 3 - infer.scheduler.unique_timesteps_cpu = torch.tensor([0.25]) - with pytest.raises(KeyError, match="no entry"): - infer.infer(Blocks(), pre) - - cached_block = weights_module.MiniMaxH3TransformerBlockWeights(0, {"use_adaln_cache": True}) - assert not hasattr(cached_block, "adaln") - assert hasattr(cached_block, "ff") diff --git a/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py b/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py deleted file mode 100644 index f1660bbf8..000000000 --- a/tests/models/minimax_h3/test_video_vae_checkpoint_adapter.py +++ /dev/null @@ -1,183 +0,0 @@ -import importlib.util -import sys -from pathlib import Path - -import pytest -import torch -import torch.nn as nn -from safetensors.torch import save_file - -REPO_ROOT = Path(__file__).parents[3] -_SPEC = importlib.util.spec_from_file_location("minimax_h3_weights_under_test", REPO_ROOT / "lightx2v/models/video_encoders/hf/minimax_h3/weights.py") -_WEIGHTS = importlib.util.module_from_spec(_SPEC) -sys.modules[_SPEC.name] = _WEIGHTS -_SPEC.loader.exec_module(_WEIGHTS) -_is_official_video_vae_checkpoint = _WEIGHTS._is_official_video_vae_checkpoint -load_minimax_h3_video_vae_checkpoint = _WEIGHTS.load_minimax_h3_video_vae_checkpoint -load_safetensors_subset = _WEIGHTS.load_safetensors_subset -validate_minimax_h3_video_vae_checkpoint = _WEIGHTS.validate_minimax_h3_video_vae_checkpoint - - -def _parameter_module(specs): - class _IndexedModule(nn.Module): - def __getitem__(self, index): - return self._modules[str(index)] - - root = _IndexedModule() - for key, tensor in specs.items(): - parent = root - parts = key.split(".") - for part in parts[:-1]: - if not hasattr(parent, part): - parent.add_module(part, _IndexedModule()) - parent = getattr(parent, part) - parent.register_parameter(parts[-1], nn.Parameter(torch.empty_like(tensor, device="meta"))) - if hasattr(root, "decoder") and hasattr(root.decoder, "transformer_blocks"): - attention = root.decoder.transformer_blocks[0].attn - attention.heads = 2 - attention.dim_head = 2 - return root - - -def _official_tensors(): - qkv_weight = torch.arange(48, dtype=torch.float32).reshape(12, 4) - qkv_bias = torch.arange(12, dtype=torch.float32) - w1_weight = torch.cat((torch.ones(4, 4), torch.full((4, 4), 2.0))) - w1_bias = torch.cat((torch.ones(4), torch.full((4,), 2.0))) - return { - "encoder.conv_in.weight": torch.full((2, 3, 1, 1, 1), 3.0), - "encoder.down.0.block.0.nin_shortcut.weight": torch.full((2, 2, 1, 1, 1), 4.0), - "encoder.down.0.downsample.conv.bias": torch.full((2,), 5.0), - "decoder.x_embedder.weight": torch.full((4, 2), 6.0), - "decoder.transformer_blocks.0.attn.to_out.weight": torch.full((4, 4), 7.0), - "decoder.transformer_blocks.0.attn.to_qkv.weight": qkv_weight, - "decoder.transformer_blocks.0.attn.to_qkv.bias": qkv_bias, - "decoder.transformer_blocks.0.ff.w1.weight": w1_weight, - "decoder.transformer_blocks.0.ff.w1.bias": w1_bias, - "decoder.transformer_blocks.0.ff.w2.weight": torch.full((4, 4), 8.0), - "decoder.mask_token": torch.zeros(1, 1, 4), - } - - -def _native_specs(): - tensors = _official_tensors() - return { - "encoder.conv_in.weight": tensors["encoder.conv_in.weight"], - "encoder.down_blocks.0.resnets.0.conv_shortcut.weight": tensors["encoder.down.0.block.0.nin_shortcut.weight"], - "encoder.down_blocks.0.downsamplers.0.conv.bias": tensors["encoder.down.0.downsample.conv.bias"], - "decoder.proj_in.weight": tensors["decoder.x_embedder.weight"], - "decoder.transformer_blocks.0.attn.to_out.0.weight": tensors["decoder.transformer_blocks.0.attn.to_out.weight"], - "decoder.transformer_blocks.0.attn.to_q.weight": torch.empty(4, 4), - "decoder.transformer_blocks.0.attn.to_k.weight": torch.empty(4, 4), - "decoder.transformer_blocks.0.attn.to_v.weight": torch.empty(4, 4), - "decoder.transformer_blocks.0.attn.to_q.bias": torch.empty(4), - "decoder.transformer_blocks.0.attn.to_k.bias": torch.empty(4), - "decoder.transformer_blocks.0.attn.to_v.bias": torch.empty(4), - "decoder.transformer_blocks.0.ff.net.0.proj.weight": tensors["decoder.transformer_blocks.0.ff.w1.weight"], - "decoder.transformer_blocks.0.ff.net.0.proj.bias": tensors["decoder.transformer_blocks.0.ff.w1.bias"], - "decoder.transformer_blocks.0.ff.net.2.weight": tensors["decoder.transformer_blocks.0.ff.w2.weight"], - } - - -def _write(path: Path, tensors=None): - save_file(tensors or _official_tensors(), path) - return path - - -def test_official_detection_uses_strict_key_signature(tmp_path): - official = _write(tmp_path / "official.safetensors") - legacy = tmp_path / "legacy.safetensors" - save_file({"decoder.proj_in.weight": torch.zeros(4, 2)}, legacy) - assert _is_official_video_vae_checkpoint(official) - assert not _is_official_video_vae_checkpoint(legacy) - - -def test_official_mapping_qkv_ffn_and_mask_token(tmp_path): - tensors = _official_tensors() - module = _parameter_module(_native_specs()) - report = load_minimax_h3_video_vae_checkpoint(module, _write(tmp_path / "model.safetensors", tensors)) - state = module.state_dict() - - assert len(report.loaded_keys) == len(_native_specs()) - assert report.ignored_keys == 1 - assert torch.equal(state["encoder.down_blocks.0.resnets.0.conv_shortcut.weight"], tensors["encoder.down.0.block.0.nin_shortcut.weight"]) - assert torch.equal(state["decoder.proj_in.weight"], tensors["decoder.x_embedder.weight"]) - assert torch.equal(state["decoder.transformer_blocks.0.attn.to_out.0.weight"], tensors["decoder.transformer_blocks.0.attn.to_out.weight"]) - assert torch.equal(state["decoder.transformer_blocks.0.ff.net.2.weight"], tensors["decoder.transformer_blocks.0.ff.w2.weight"]) - for index, name in enumerate(("q", "k", "v")): - assert torch.equal(state[f"decoder.transformer_blocks.0.attn.to_{name}.weight"], tensors["decoder.transformer_blocks.0.attn.to_qkv.weight"].reshape(2, 3, 2, 4)[:, index].reshape(4, 4)) - assert torch.equal(state[f"decoder.transformer_blocks.0.attn.to_{name}.bias"], tensors["decoder.transformer_blocks.0.attn.to_qkv.bias"].reshape(2, 3, 2)[:, index].reshape(4)) - assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.weight"][:4] == 2) - assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.weight"][4:] == 1) - assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.bias"][:4] == 2) - assert torch.all(state["decoder.transformer_blocks.0.ff.net.0.proj.bias"][4:] == 1) - - -def test_mask_token_shape_is_validated(tmp_path): - tensors = _official_tensors() - tensors["decoder.mask_token"] = torch.zeros(1, 2, 4) - with pytest.raises(RuntimeError, match="shape_mismatch.*decoder.mask_token"): - validate_minimax_h3_video_vae_checkpoint(_parameter_module(_native_specs()), _write(tmp_path / "bad.safetensors", tensors)) - - -def test_unexpected_official_key_fails(tmp_path): - tensors = _official_tensors() - tensors["decoder.surprise"] = torch.zeros(1) - with pytest.raises(RuntimeError, match="unknown.*decoder.surprise"): - validate_minimax_h3_video_vae_checkpoint(_parameter_module(_native_specs()), _write(tmp_path / "bad.safetensors", tensors)) - - -def test_missing_native_target_fails(tmp_path): - specs = _native_specs() - specs["decoder.proj_out.weight"] = torch.empty(4, 4) - with pytest.raises(RuntimeError, match="missing.*decoder.proj_out.weight"): - validate_minimax_h3_video_vae_checkpoint(_parameter_module(specs), _write(tmp_path / "bad.safetensors")) - - -def test_duplicate_target_fails(tmp_path): - tensors = _official_tensors() - tensors["decoder.proj_in.weight"] = tensors["decoder.x_embedder.weight"] - with pytest.raises(RuntimeError, match="duplicate.*decoder.proj_in.weight"): - validate_minimax_h3_video_vae_checkpoint(_parameter_module(_native_specs()), _write(tmp_path / "bad.safetensors", tensors)) - - -def test_legacy_subset_loader_regression(tmp_path): - expected = {"layer.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3)} - module = _parameter_module(expected) - path = tmp_path / "native.safetensors" - save_file(expected, path) - report = load_safetensors_subset(module, path) - assert report.loaded_keys == ("layer.weight",) - assert torch.equal(module.state_dict()["layer.weight"], expected["layer.weight"]) - - -@pytest.mark.parametrize("is_weight", [False, True]) -def test_qkv_per_head_components_and_reinterleave(is_weight): - heads, head_dim = 3, 2 - rows = torch.tensor([100 * h + 10 * c + d for h in range(heads) for c in range(3) for d in range(head_dim)]) - source = rows.float() - if is_weight: - source = source[:, None] * 10 + torch.arange(5) - parts = _WEIGHTS._split_video_vae_qkv(source, heads, head_dim) - for c, part in enumerate(parts): - expected = torch.tensor([100 * h + 10 * c + d for h in range(heads) for d in range(head_dim)]).float() - if is_weight: - expected = expected[:, None] * 10 + torch.arange(5) - assert torch.equal(part, expected) - assert not torch.equal(part, source.chunk(3)[c]) - assert part.dtype == source.dtype and part.device == source.device and part.is_contiguous() - rebuilt = torch.stack([p.reshape(heads, head_dim, *source.shape[1:]) for p in parts], dim=1).reshape_as(source) - assert torch.equal(rebuilt, source) - - -@pytest.mark.parametrize("shape,heads,dim", [((11,), 2, 2), ((12,), 3, 2), ((12,), 2, 0), ((12, 2, 2), 2, 2)]) -def test_qkv_invalid_geometry_rejected(shape, heads, dim): - with pytest.raises(ValueError, match="Video VAE fused QKV"): - _WEIGHTS._split_video_vae_qkv(torch.empty(shape), heads, dim) - - -def test_loader_rejects_incompatible_attention_geometry(tmp_path): - module = _parameter_module(_native_specs()) - module.decoder.transformer_blocks[0].attn.heads = 3 - with pytest.raises(ValueError, match="target num_heads"): - load_minimax_h3_video_vae_checkpoint(module, _write(tmp_path / "model.safetensors")) diff --git a/tests/models/minimax_h3/test_video_vae_loader.py b/tests/models/minimax_h3/test_video_vae_loader.py deleted file mode 100644 index a3b6e80cd..000000000 --- a/tests/models/minimax_h3/test_video_vae_loader.py +++ /dev/null @@ -1,295 +0,0 @@ -import importlib.util -import sys -import types -from pathlib import Path - -import pytest - -REPO_ROOT = Path(__file__).parents[3] - - -def _install_module(monkeypatch, name, **attrs): - module = types.ModuleType(name) - for attr_name, value in attrs.items(): - setattr(module, attr_name, value) - monkeypatch.setitem(sys.modules, name, module) - return module - - -@pytest.fixture() -def video_vae_module(monkeypatch): - for package_name in [ - "lightx2v", - "lightx2v.models", - "lightx2v.models.video_encoders", - "lightx2v.models.video_encoders.hf", - "lightx2v.models.video_encoders.hf.minimax_h3", - "lightx2v.utils", - "lightx2v_platform", - "lightx2v_platform.base", - ]: - package = types.ModuleType(package_name) - package.__path__ = [] - monkeypatch.setitem(sys.modules, package_name, package) - - # Load main's FP8 helpers without importing the full runtime package. - for module_name in ( - "lightx2v.common.ops.mm.fp8_f16_accum", - "lightx2v.models.networks.minimax_h3.fp8_f16_accum_policy", - ): - spec = importlib.util.spec_from_file_location(module_name, REPO_ROOT / (module_name.replace(".", "/") + ".py")) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, module_name, module) - spec.loader.exec_module(module) - - _install_module( - monkeypatch, - "lightx2v.models.video_encoders.hf.minimax_h3.weights", - SafetensorsSubsetReport=object, - _is_official_video_vae_checkpoint=lambda *_args, **_kwargs: False, - load_minimax_h3_video_vae_checkpoint=lambda *_args, **_kwargs: None, - load_safetensors_subset=lambda *_args, **_kwargs: None, - ) - _install_module(monkeypatch, "lightx2v.utils.registry_factory", ATTN_WEIGHT_REGISTER={}) - _install_module(monkeypatch, "lightx2v_platform.base.global_var", AI_DEVICE="cpu") - - module_path = REPO_ROOT / "lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py" - spec = importlib.util.spec_from_file_location("minimax_h3_video_vae_under_test", module_path) - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -def _write_json(path, data): - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(__import__("json").dumps(data), encoding="utf-8") - - -def _wrapper_config(): - return { - "source_path": "source", - "source_safetensors_path": "model.safetensors", - "latent_channels": 24, - "latents_mean": [0.1] * 24, - "latents_std": [1.1] * 24, - "vae_clip_length": 17, - "vae_token_drop": 3, - } - - -def _source_config(): - return { - "in_channels": 3, - "out_ch": 3, - "z_channels": 24, - "ch": 8, - "ch_mult": [1, 2, 4], - "num_res_blocks": 2, - "space_down": [2, 2, 1], - "time_down": [1, 2, 1], - "padding_mode": "reflect", - "vit_decoder_kwargs": { - "num_layers": 12, - "heads": 4, - "dim_head": 16, - "rope_theta": 100.0, - "rope_dim_ratio": 0.75, - }, - } - - -def test_resolves_official_video_vae_layout_before_legacy(video_vae_module, tmp_path): - official = tmp_path / "video_vae" - legacy = tmp_path / "vae" - official.mkdir() - legacy.mkdir() - - assert video_vae_module._resolve_video_vae_dir(tmp_path) == official - - -def test_resolves_legacy_vae_layout(video_vae_module, tmp_path): - legacy = tmp_path / "vae" - legacy.mkdir() - - assert video_vae_module._resolve_video_vae_dir(tmp_path) == legacy - - -@pytest.mark.parametrize("name", ["video_vae", "vae"]) -def test_resolves_direct_component_path(video_vae_module, tmp_path, name): - component = tmp_path / name - component.mkdir() - - assert video_vae_module._resolve_video_vae_dir(component) == component - - -def test_missing_video_vae_and_vae_reports_tried_paths(video_vae_module, tmp_path): - with pytest.raises(FileNotFoundError, match="video_vae.*vae"): - video_vae_module._resolve_video_vae_dir(tmp_path) - - -def test_official_layout_uses_source_safetensors_by_default(video_vae_module, tmp_path): - vae_dir = tmp_path / "video_vae" - _write_json(vae_dir / "config.json", _wrapper_config()) - _write_json(vae_dir / "source/config.json", _source_config()) - (vae_dir / "source/model.safetensors").write_bytes(b"") - - config, weight_path = video_vae_module._load_video_vae_config_and_weight_path(vae_dir, None) - - assert weight_path == vae_dir / "source/model.safetensors" - assert config["block_out_channels"] == [8, 16, 32] - - -def test_explicit_checkpoint_path_wins_for_official_layout(video_vae_module, tmp_path): - vae_dir = tmp_path / "video_vae" - explicit = tmp_path / "quantized.safetensors" - _write_json(vae_dir / "config.json", _wrapper_config()) - _write_json(vae_dir / "source/config.json", _source_config()) - (vae_dir / "source/model.safetensors").write_bytes(b"") - - _config, weight_path = video_vae_module._load_video_vae_config_and_weight_path(vae_dir, explicit) - - assert weight_path == explicit - - -def test_official_config_normalization_maps_wrapper_and_source_fields(video_vae_module): - config = video_vae_module._normalize_official_video_vae_config(_wrapper_config(), _source_config()) - - assert config["in_channels"] == 3 - assert config["out_channels"] == 3 - assert config["latent_channels"] == 24 - assert config["block_out_channels"] == [8, 16, 32] - assert config["layers_per_block"] == 2 - assert config["spatial_downsample_factors"] == [2, 2, 1] - assert config["temporal_downsample_factors"] == [1, 2, 1] - assert config["spatial_padding_mode"] == "reflect" - assert config["decoder_num_layers"] == 12 - assert config["decoder_num_attention_heads"] == 4 - assert config["decoder_attention_head_dim"] == 16 - assert config["decoder_rope_theta"] == 100.0 - assert config["decoder_rope_dim_ratio"] == 0.75 - assert config["clip_length"] == 17 - assert config["token_drop"] == 3 - assert config["latents_mean"] == [0.1] * 24 - assert config["latents_std"] == [1.1] * 24 - - -def test_legacy_config_normalization_keeps_existing_fields(video_vae_module): - legacy = { - "latent_channels": 8, - "block_out_channels": [4, 8], - "clip_length": 9, - "token_drop": 1, - "custom": "kept", - } - - assert video_vae_module._normalize_official_video_vae_config(legacy, None) == legacy - - -def test_legacy_layout_uses_component_dir_as_weight_path(video_vae_module, tmp_path): - vae_dir = tmp_path / "vae" - legacy = {"latent_channels": 8, "block_out_channels": [4, 8]} - _write_json(vae_dir / "config.json", legacy) - (vae_dir / "model.safetensors").write_bytes(b"") - - config, weight_path = video_vae_module._load_video_vae_config_and_weight_path(vae_dir, None) - - assert config == legacy - assert weight_path == vae_dir - - -def test_from_pretrained_dispatches_official_and_legacy(video_vae_module, monkeypatch, tmp_path): - calls = [] - - class TinyVAE(video_vae_module.MiniMaxH3VideoVAE): - def __init__(self, _config, **_kwargs): - video_vae_module.nn.Module.__init__(self) - self.weight = video_vae_module.nn.Parameter(video_vae_module.torch.empty(1)) - self.execution_device = video_vae_module.torch.device("cpu") - - def _reset_runtime_buffers(self): - pass - - def _prepare_inference_dtypes(self): - pass - - monkeypatch.setattr(video_vae_module, "_resolve_video_vae_dir", lambda _path: tmp_path) - monkeypatch.setattr(video_vae_module, "_load_video_vae_config_and_weight_path", lambda *_args: ({}, tmp_path / "model.safetensors")) - monkeypatch.setattr(video_vae_module, "load_minimax_h3_video_vae_checkpoint", lambda *_args: calls.append("official") or "official-report") - monkeypatch.setattr(video_vae_module, "load_safetensors_subset", lambda *_args: calls.append("legacy") or "legacy-report") - - monkeypatch.setattr(video_vae_module, "_is_official_video_vae_checkpoint", lambda _path: True) - assert TinyVAE.from_pretrained(tmp_path, cpu_offload=True).load_report == "official-report" - monkeypatch.setattr(video_vae_module, "_is_official_video_vae_checkpoint", lambda _path: False) - assert TinyVAE.from_pretrained(tmp_path, cpu_offload=True).load_report == "legacy-report" - assert calls == ["official", "legacy"] - - -def test_from_pretrained_quantized_path_bypasses_official_adapter(video_vae_module, monkeypatch, tmp_path): - calls = [] - - class TinyQuantizedVAE(video_vae_module.MiniMaxH3VideoVAE): - def __init__(self, _config, **_kwargs): - video_vae_module.nn.Module.__init__(self) - self.weight = video_vae_module.nn.Parameter(video_vae_module.torch.empty(1)) - self.execution_device = video_vae_module.torch.device("cpu") - - def _reset_runtime_buffers(self): - pass - - def _pack_decoder_fp8_qkv(self): - calls.append("pack") - - def _prepare_inference_dtypes(self): - pass - - monkeypatch.setattr(video_vae_module, "_resolve_video_vae_dir", lambda _path: tmp_path) - monkeypatch.setattr(video_vae_module, "_load_video_vae_config_and_weight_path", lambda *_args: ({}, tmp_path / "quant.safetensors")) - monkeypatch.setattr(video_vae_module, "_is_official_video_vae_checkpoint", lambda _path: pytest.fail("official detection must not run for quantized checkpoints")) - monkeypatch.setattr(video_vae_module, "load_minimax_h3_video_vae_checkpoint", lambda *_args: pytest.fail("official adapter must not run for quantized checkpoints")) - monkeypatch.setattr(video_vae_module, "load_safetensors_subset", lambda *_args: calls.append("legacy") or "quant-report") - - model = TinyQuantizedVAE.from_pretrained(tmp_path, checkpoint_path=tmp_path / "quant.safetensors", quant_scheme="fp8-sgl", cpu_offload=True) - assert model.load_report == "quant-report" - assert calls == ["legacy", "pack"] - - -@pytest.mark.parametrize("device", ["cpu", "mps"]) -@pytest.mark.parametrize("dtype_name", ["float32", "float16"]) -@pytest.mark.parametrize("size", [64, 258]) -def test_causal_temporal_padding_preserves_interior(video_vae_module, device, dtype_name, size): - import torch - import torch.nn.functional as F - - if device == "mps" and not torch.backends.mps.is_available(): - pytest.skip("MPS is unavailable") - dtype = getattr(torch, dtype_name) - source = torch.randn((1, 3, 17, size, size), generator=torch.Generator().manual_seed(123)).to(dtype) - conv = video_vae_module.MiniMaxH3VideoCausalConv3d(3, 3, 1, temporal_padding=2) - actual_device = conv._pad_temporal(source.to(device)) - actual = actual_device.cpu() - expected = F.pad(source, (0, 0, 0, 0, 2, 0)) - - assert actual.shape == (1, 3, 19, size, size) - assert actual_device.device.type == device - assert actual.dtype == dtype - assert torch.count_nonzero(actual[:, :, :2]) == 0 - assert torch.equal(actual[:, :, 2:], source) - assert torch.equal(actual, expected) - assert torch.isfinite(actual).all() - - -@pytest.mark.parametrize("device", ["cpu", "mps"]) -def test_causal_conv_forward_uses_temporal_padding(video_vae_module, device): - import torch - import torch.nn.functional as F - - if device == "mps" and not torch.backends.mps.is_available(): - pytest.skip("MPS is unavailable") - conv = video_vae_module.MiniMaxH3VideoCausalConv3d(1, 1, 1, temporal_padding=2).to(device) - with torch.no_grad(): - conv.weight.fill_(1) - conv.bias.zero_() - source = torch.randn((1, 1, 3, 258, 258), generator=torch.Generator().manual_seed(123)) - actual = conv(source.to(device)).cpu() - assert torch.equal(actual, F.pad(source, (0, 0, 0, 0, 2, 0))) diff --git a/tools/cache_minimax_h3_adaln/builder.py b/tools/cache_minimax_h3_adaln/builder.py index 639e24528..5aef69535 100644 --- a/tools/cache_minimax_h3_adaln/builder.py +++ b/tools/cache_minimax_h3_adaln/builder.py @@ -30,7 +30,6 @@ _timesteps_from_bits, _validate_cache, ) -from lightx2v.models.networks.minimax_h3.checkpoint import MiniMaxH3ShardCheckpoint from lightx2v.models.networks.minimax_h3.infer.pre_infer import timestep_embedding from lightx2v_platform.base.global_var import AI_DEVICE @@ -48,13 +47,9 @@ def _checkpoint_files(config) -> list[Path]: class _CheckpointTensors: """Read individual tensors without materializing the whole checkpoint.""" - def __init__(self, files: list[Path], config=None): + def __init__(self, files: list[Path]): self.files = files - self.adapter = None - if (files[0].parent / "model.safetensors.index.json").is_file(): - checkpoint = MiniMaxH3ShardCheckpoint(files[0].parent, config=config) - self.adapter = checkpoint.selected_reader - self.locations = self._find_locations() if self.adapter is None else {} + self.locations = self._find_locations() def _find_locations(self) -> dict[str, Path]: directory = self.files[0].parent @@ -71,17 +66,6 @@ def _find_locations(self) -> dict[str, Path]: return locations def get(self, name: str) -> torch.Tensor: - if self.adapter is not None: - # Reuse the inference mapping and validation. The builder consumes - # logical (out, in) weights; _linear owns the eventual transpose. - try: - spec = self.adapter.plan.targets[name][1] - except KeyError as error: - raise KeyError(f"MiniMax-H3 checkpoint tensor is missing: {name}") from error - dtype = torch.float32 if spec.dtype == "F32" else torch.bfloat16 - tensor = torch.empty(spec.shape, dtype=dtype, device="cpu") - self.adapter.write_targets({name: (tensor, False)}) - return tensor path = self.locations.get(name) if path is None: raise KeyError(f"MiniMax-H3 checkpoint tensor is missing: {name}") @@ -110,10 +94,10 @@ def _empty_device_cache() -> None: torch_device_module.empty_cache() -def _build_cache(spec: dict, cache_path: Path, checkpoint_files: list[Path], config=None) -> None: +def _build_cache(spec: dict, cache_path: Path, checkpoint_files: list[Path]) -> None: stage_path = Path(tempfile.mkdtemp(prefix=".building-", dir=cache_path.parent)) try: - checkpoint = _CheckpointTensors(checkpoint_files, config=config) + checkpoint = _CheckpointTensors(checkpoint_files) with torch.inference_mode(): # ADALN CACHE SYNC: Keep activation placement and casts aligned with # the three online infer modules named in this file's contract. @@ -184,7 +168,7 @@ def build_persistent_adaln_cache(config) -> Path: raise FileExistsError(f"MiniMax-H3 AdaLN cache path already exists: {cache_path}") checkpoint_files = _checkpoint_files(config) logger.info("Building MiniMax-H3 AdaLN cache on {}: {}", AI_DEVICE, cache_path) - _build_cache(spec, cache_path, checkpoint_files, config=config) + _build_cache(spec, cache_path, checkpoint_files) if not _validate_cache(cache_path, spec): raise RuntimeError(f"MiniMax-H3 AdaLN cache validation failed: {cache_path}") diff --git a/tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh b/tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh index 8f0b5f950..e99723c62 100755 --- a/tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh +++ b/tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh @@ -1,14 +1,19 @@ #!/bin/bash +set -euo pipefail # set path firstly -lightx2v_path=/data/nvme1/yongyang/dan/LightX2V -model_path=/data/nvme1/models/MiniMaxAI/MiniMax-H3 +lightx2v_path=/Users/yongyang/Documents/x2v/LightX2V +model_path=/Users/yongyang/Documents/x2v/models/MiniMaxAI/diffusers/MiniMax-H3 +config_json="${CONFIG_JSON:-${lightx2v_path}/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json}" -# Select one platform. NVIDIA is enabled by default. +# Select one platform. Apple MPS is enabled by default. + +# Apple MPS +export PLATFORM=mps # NVIDIA -export PLATFORM=cuda -export CUDA_VISIBLE_DEVICES=0 +# export PLATFORM=cuda +# export CUDA_VISIBLE_DEVICES=0 # Intel XPU # export PLATFORM=intel_xpu @@ -51,10 +56,15 @@ export CUDA_VISIBLE_DEVICES=0 # export CUDA_VISIBLE_DEVICES=0 # set environment variables -source "${lightx2v_path}/scripts/base/base.sh" +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 +export TOKENIZERS_PARALLELISM=false +export PYTHONUNBUFFERED=1 +export PYTHONPATH="${lightx2v_path}:${PYTHONPATH:-}" -# Supported tasks: fl2av, ref2av -python "${lightx2v_path}/tools/cache_minimax_h3_adaln/cache_minimax_h3_adaln.py" \ +# fl2av also generates the cache used by t2av/i2av/l2av inference. +# Cache generation and inference must use the same model and JSON config. +/opt/miniconda3/envs/torch/bin/python "${lightx2v_path}/tools/cache_minimax_h3_adaln/cache_minimax_h3_adaln.py" \ --model_path "${model_path}" \ - --config_json "${lightx2v_path}/configs/minimax_h3/minimax_h3.json" \ + --config_json "${config_json}" \ --task fl2av From f502c4a61a62f5c1e38ca7489e48df7b22224015 Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Thu, 10 Sep 2026 23:02:02 +0800 Subject: [PATCH 31/31] =?UTF-8?q?perf(mps):=20=E4=B8=BA=20MiniMax-H3=20?= =?UTF-8?q?=E5=AE=9E=E7=8E=B0=E5=85=B1=E4=BA=AB=E5=86=85=E5=AD=98=E5=8F=8C?= =?UTF-8?q?=20buffer=20offload?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 复用现有 block offload 架构，重叠后台权重预读与 GPU 计算 - 通过 MPS 共享视图直接读取 safetensors，省去中间 CPU 权重副本 - 保持 diffusers 权重发现规则及原有计算逻辑 - 完善 buffer 同步、异常处理、释放及重新初始化流程 - 默认启用 dit_mps_shared_buffer，并补充使用说明 --- configs/platforms/mps/minimax_h3_t2av.json | 1 + .../mps/minimax_h3_t2av_4step_512_22.json | 1 + lightx2v/common/offload/manager.py | 14 ++- lightx2v/common/offload/mps_manager.py | 89 +++++++++++++++++++ .../models/networks/minimax_h3/checkpoint.py | 56 ++++++++++++ .../infer/offload/transformer_infer.py | 39 ++++---- lightx2v/models/networks/minimax_h3/model.py | 13 ++- .../minimax_h3/weights/transformer_weights.py | 67 +++++++++++++- .../runners/minimax_h3/minimax_h3_runner.py | 6 +- scripts/platforms/mps/README.md | 4 + 10 files changed, 265 insertions(+), 25 deletions(-) create mode 100644 lightx2v/common/offload/mps_manager.py diff --git a/configs/platforms/mps/minimax_h3_t2av.json b/configs/platforms/mps/minimax_h3_t2av.json index 8a8d488b6..75fff1ab9 100644 --- a/configs/platforms/mps/minimax_h3_t2av.json +++ b/configs/platforms/mps/minimax_h3_t2av.json @@ -11,6 +11,7 @@ "offload_granularity": "block", "dit_prepost_resident": false, "dit_disk_streaming": true, + "dit_mps_shared_buffer": true, "use_adaln_cache": true, "adaln_cache_dir": "~/.cache/lightx2v/adaln/diffusers", diff --git a/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json b/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json index f953292f1..f266be95b 100644 --- a/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json +++ b/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json @@ -11,6 +11,7 @@ "offload_granularity": "block", "dit_prepost_resident": false, "dit_disk_streaming": true, + "dit_mps_shared_buffer": true, "use_adaln_cache": true, "adaln_cache_dir": "~/.cache/lightx2v/adaln/diffusers", diff --git a/lightx2v/common/offload/manager.py b/lightx2v/common/offload/manager.py index 3fb476ad6..08d5383f9 100755 --- a/lightx2v/common/offload/manager.py +++ b/lightx2v/common/offload/manager.py @@ -1,4 +1,5 @@ from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext import torch from loguru import logger @@ -14,9 +15,12 @@ class WeightAsyncStreamManager(object): def __init__(self, offload_granularity): self.offload_granularity = offload_granularity - self.init_stream = torch_device_module.Stream(priority=0) self.need_init_first_buffer = True self.lazy_load = False + self._init_streams() + + def _init_streams(self): + self.init_stream = torch_device_module.Stream(priority=0) torch_version = parse(torch.__version__.split("+")[0]) # Legacy name: this is the active device backend's weight-loading stream, not a CUDA-only stream. if AI_DEVICE == "cuda" and torch_version >= parse("2.7"): @@ -26,6 +30,14 @@ def __init__(self, offload_granularity): self.cuda_load_stream = torch_device_module.Stream(priority=0) self.compute_stream = torch_device_module.Stream(priority=-1) + def prepare_compute(self): + self.compute_stream.wait_stream(torch_device_module.current_stream()) + + def compute_context(self): + if AI_DEVICE == "xpu": + return nullcontext() + return torch_device_module.stream(self.compute_stream) + def init_cpu_buffer(self, blocks_cpu_buffer=None, phases_cpu_buffer=None): self.need_init_first_buffer = True if self.offload_granularity == "block": diff --git a/lightx2v/common/offload/mps_manager.py b/lightx2v/common/offload/mps_manager.py new file mode 100644 index 000000000..a58d56299 --- /dev/null +++ b/lightx2v/common/offload/mps_manager.py @@ -0,0 +1,89 @@ +"""Double-buffered disk offload into CPU-visible MPS weight storage.""" + +from contextlib import nullcontext + +import torch +from loguru import logger + +from lightx2v.common.offload.manager import WeightAsyncStreamManager + + +def host_view(tensor): + """Alias MPS storage without copying; the caller orders CPU/GPU accesses.""" + if tensor.device.type != "mps": + raise ValueError("Shared weight buffers require an MPS tensor") + if not hasattr(torch.mps, "_host_alias_storage"): + raise RuntimeError("Shared MPS weight buffers require torch.mps._host_alias_storage (PyTorch >= 2.13)") + storage = torch.mps._host_alias_storage(tensor.untyped_storage()) + return torch.empty(0, dtype=tensor.dtype).set_(storage, tensor.storage_offset(), tensor.shape, tensor.stride()) + + +class MpsSharedWeightAsyncStreamManager(WeightAsyncStreamManager): + """Read into the idle slot while the default MPS stream uses the other. + + The source implements load_block_into(target_buffer, block_idx). This runs + on a CPU worker and must only write host views of the supplied idle buffer. + """ + + def __init__(self, offload_granularity="block"): + if offload_granularity != "block": + raise ValueError("Shared MPS weight offload only supports block granularity") + super().__init__(offload_granularity) + self.cuda_buffers = [] + self.init_lazy_load(num_workers=1) + logger.info("MPS shared weight offload: two device buffers, direct file reads, one prefetch worker") + + def _init_streams(self): + # Disk I/O runs on a CPU thread; GPU work stays on the default stream. + pass + + def prepare_compute(self): + pass + + def compute_context(self): + return nullcontext() + + def init_cuda_buffer(self, blocks_cuda_buffer=None, phases_cuda_buffer=None): + if blocks_cuda_buffer is None or len(blocks_cuda_buffer) != 2: + raise ValueError("Shared MPS offload requires exactly two device buffers") + if self.prefetch_futures: + raise RuntimeError("Cannot replace shared buffers while a prefetch is pending") + super().init_cuda_buffer(blocks_cuda_buffer, phases_cuda_buffer) + if self.executor is None: + self.init_lazy_load(num_workers=1) + + def init_first_buffer(self, blocks, adapter_block_idx=None): + torch.mps.synchronize() + blocks.load_block_into(self.cuda_buffers[0], 0) + torch.mps.synchronize() + self.need_init_first_buffer = False + + def prefetch_weights(self, block_idx, blocks, adapter_block_idx=None): + if self.prefetch_futures: + raise RuntimeError("Call swap_blocks before scheduling another prefetch") + self.prefetch_block_idx = block_idx + self.prefetch_futures = [self.executor.submit(blocks.load_block_into, self.cuda_buffers[1], block_idx)] + + def swap_blocks(self): + if not self.prefetch_futures: + raise RuntimeError("No shared-buffer prefetch to complete") + # This commits GPU work and releases the GIL while waiting, so the disk + # reader continues filling slot 1 while the GPU consumes slot 0. + torch.mps.synchronize() + for future in self.prefetch_futures: + future.result() + torch.mps.synchronize() + self.prefetch_futures.clear() + self.cuda_buffers.reverse() + + def close(self): + """Drain both users before releasing aliases; allow later reinitialization.""" + try: + if self.executor is not None: + self.executor.shutdown(wait=True) + torch.mps.synchronize() + finally: + self.executor = None + self.prefetch_futures.clear() + self.cuda_buffers = [] + self.need_init_first_buffer = True diff --git a/lightx2v/models/networks/minimax_h3/checkpoint.py b/lightx2v/models/networks/minimax_h3/checkpoint.py index 6e764d626..1d7d514e8 100644 --- a/lightx2v/models/networks/minimax_h3/checkpoint.py +++ b/lightx2v/models/networks/minimax_h3/checkpoint.py @@ -1,9 +1,14 @@ """Load MiniMax-H3 diffusers checkpoints by requested tensor or block.""" +import json +import math import re +import struct +import sys from collections import defaultdict from pathlib import Path +import torch from safetensors import safe_open _H3_BLOCK_KEY_RE = re.compile(r"^transformer_blocks\.(\d+)\.") @@ -19,12 +24,63 @@ def __init__(self, checkpoint_dir): raise FileNotFoundError(f"MiniMax-H3 safetensors checkpoint not found: {checkpoint}") self.checkpoint_dir = checkpoint if checkpoint.is_dir() else checkpoint.parent self.weight_map = {} + self._shard_headers = {} # Match the upstream model loader's directory/single-file discovery. # Read only headers here; tensor data is loaded when a block requests it. for path in files: with safe_open(path, framework="pt", device="cpu") as source: self.weight_map.update(dict.fromkeys(source.keys(), path.name)) + def tensor_metadata(self, name): + """Return dtype, shape and absolute file range without loading payloads.""" + shard_name = self.shard_for_tensor(name) + if shard_name not in self._shard_headers: + with (self.checkpoint_dir / shard_name).open("rb") as source: + header_size = struct.unpack("