diff --git a/configs/platforms/mps/minimax_h3_t2av.json b/configs/platforms/mps/minimax_h3_t2av.json new file mode 100644 index 000000000..395d821af --- /dev/null +++ b/configs/platforms/mps/minimax_h3_t2av.json @@ -0,0 +1,51 @@ +{ + "infer_steps": 29, + "num_frames": 124, + "size": [480, 480], + "fps": 24, + "enable_cfg": false, + + "cpu_offload": true, + "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", + + "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/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..2360748f9 --- /dev/null +++ b/configs/platforms/mps/minimax_h3_t2av_4step_512_22.json @@ -0,0 +1,51 @@ +{ + "infer_steps": 4, + "num_frames": 22, + "size": [512, 512], + "fps": 24, + "enable_cfg": false, + + "cpu_offload": true, + "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", + + "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/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/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/attn/torch_sdpa.py b/lightx2v/common/ops/attn/torch_sdpa.py index fd12abd80..a9cfd4f21 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,17 +71,21 @@ def apply( enable_mem_efficient=True, ) with sdpa_ctx: - # q/k/v are (B, H, S, D) here, so head count is dim 1. GQA models such as - # neopp (32 q heads, 8 kv heads) need SDPA to broadcast the kv groups. - x = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=attn_mask, - dropout_p=drop_rate, - is_causal=causal, - enable_gqa=q.shape[1] != k.shape[1], - ) + 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: + # q/k/v are (B, H, S, D) here, so head count is dim 1. GQA models such as + # neopp (32 q heads, 8 kv heads) need SDPA to broadcast the kv groups. + x = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=drop_rate, + is_causal=causal, + enable_gqa=q.shape[1] != k.shape[1], + ) x = x.transpose(1, 2) b, s, a, d = x.shape out = x.reshape(b, s, -1) diff --git a/lightx2v/common/ops/mm/fp8_f16_accum.py b/lightx2v/common/ops/mm/fp8_f16_accum.py index f04034f3b..53ff0a063 100644 --- a/lightx2v/common/ops/mm/fp8_f16_accum.py +++ b/lightx2v/common/ops/mm/fp8_f16_accum.py @@ -2,8 +2,6 @@ import torch -from lightx2v.common.ops.mm.triton_kernels import fp8_quantize_range_triton - try: from lightx2v_kernel.gemm import FP8_F16_ACCUM_MM_AVAILABLE, cutlass_scaled_fp8_mm_f16_accum except ImportError: @@ -35,6 +33,8 @@ def validate_fp8_f16_accum_qmax(activation_qmax): def fp8_f16_accum_linear(input_tensor, weight, weight_scale, bias, activation_qmax): + from lightx2v.common.ops.mm.triton_kernels import fp8_quantize_range_triton + input_shape = input_tensor.shape input_matrix = input_tensor.reshape(-1, input_shape[-1]) quantized, activation_scale = fp8_quantize_range_triton(input_matrix, activation_qmax) diff --git a/lightx2v/common/ops/mm/mm_weight.py b/lightx2v/common/ops/mm/mm_weight.py index 46a1e6f74..b0dc9e74e 100755 --- a/lightx2v/common/ops/mm/mm_weight.py +++ b/lightx2v/common/ops/mm/mm_weight.py @@ -6,6 +6,8 @@ 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: @@ -17,14 +19,24 @@ validate_fp8_f16_accum_qmax, ) 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 @@ -32,7 +44,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/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/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py b/lightx2v/models/input_encoders/hf/minimax_h3/qwen3vl.py index 0d2c03c2f..5c466ffa6 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("model_variant") != "fl2av": + raise ValueError("MiniMax-H3 Qwen3-VL text_encoder_disk_streaming requires model_variant='fl2av' and supports t2av requests 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 @@ -1236,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 @@ -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/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py index 8fd1436cc..c5078b3ca 100644 --- a/lightx2v/models/networks/base_model.py +++ b/lightx2v/models/networks/base_model.py @@ -337,7 +337,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): diff --git a/lightx2v/models/networks/minimax_h3/checkpoint.py b/lightx2v/models/networks/minimax_h3/checkpoint.py new file mode 100644 index 000000000..1d7d514e8 --- /dev/null +++ b/lightx2v/models/networks/minimax_h3/checkpoint.py @@ -0,0 +1,125 @@ +"""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+)\.") + + +class MiniMaxH3ShardCheckpoint: + """Index safetensors headers for selective loading, using upstream file discovery.""" + + 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 = {} + 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("= 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: diff --git a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py index c73a174ef..c14c9835f 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 @@ -112,6 +113,8 @@ def __init__(self, config): def get_supported_tasks(self): """Return tasks supported by the loaded transformer weights.""" + if self.config.get("text_encoder_disk_streaming", False): + return ("t2av",) if self.config["model_variant"] == "ref2av": return ("ref2av",) return ("t2av", "i2av", "l2av", "fl2av") @@ -204,10 +207,25 @@ 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("model_variant") == "fl2av" 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"], @@ -229,14 +247,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}") @@ -469,12 +487,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 = [], [] @@ -566,6 +603,10 @@ 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_mps_shared_buffer", False): + logger.info("MiniMax-H3 diffusers disk offload enabled; prefetching directly into two shared MPS block buffers") + 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() @@ -596,7 +637,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 buffers") + if not self.model.prepost_resident: + self.model.pre_weight.to_cpu() + self.model.post_weight.to_cpu() + self.model.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() @@ -613,6 +660,7 @@ def _offload_transformer(self): metrics_labels=["MiniMaxH3Runner"], ) def run_vae_decoder(self, video_rows, audio_rows): + 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( @@ -708,8 +756,9 @@ def run_main(self): with suppress(Exception): self._offload_transformer() try: - self.end_run() + self._release_low_memory_vae() 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/lightx2v/models/schedulers/minimax_h3/scheduler.py b/lightx2v/models/schedulers/minimax_h3/scheduler.py index 4b1a0b0df..c0cf34d5e 100644 --- a/lightx2v/models/schedulers/minimax_h3/scheduler.py +++ b/lightx2v/models/schedulers/minimax_h3/scheduler.py @@ -30,7 +30,7 @@ def _make_schedule(infer_steps: int, shift: float, device) -> tuple[torch.Tensor 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/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py index 0f121fcb6..79aa8668f 100644 --- a/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py +++ b/lightx2v/models/video_encoders/hf/minimax_h3/video_vae.py @@ -147,6 +147,15 @@ 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 _enable_fp8(self, policy: Fp8EncoderConvPolicy) -> None: # The model is still on meta here, so replacing the parameter only # changes the expected checkpoint schema; no materialized weight is lost. @@ -180,7 +189,7 @@ def forward(self, hidden_states): 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) if self.weight.dtype == torch.float8_e4m3fn: return self._run_fp8_conv3d(hidden_states) return F.conv3d(hidden_states, self.weight, self.bias, stride=self.stride, dilation=self.dilation) @@ -800,6 +809,8 @@ def from_pretrained( raise ValueError(f"Unsupported MiniMax-H3 VAE encoder_conv_mode {encoder_conv_mode!r}; expected one of {sorted(_SUPPORTED_ENCODER_CONV_MODES)}") if (checkpoint_path is None) != (quant_scheme is None): raise ValueError("MiniMax-H3 video VAE checkpoint_path and quant_scheme must be configured together") + 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 encoder_conv_checkpoint_info = inspect_fp8_encoder_conv_checkpoint(weight_path) if encoder_conv_mode in FP8_ENCODER_CONV_MODES: @@ -826,8 +837,6 @@ def from_pretrained( ) else: logger.warning("MiniMax-H3 Video VAE FP8-F16 accumulation requested but {}; falling back to FP8-SGL", fallback_reason) - with (vae_dir / "config.json").open("r", encoding="utf-8") as handle: - config = json.load(handle) # The released decoder is several GiB. Constructing it on meta avoids # allocating and then immediately overwriting random initialized weights. 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 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.") diff --git a/scripts/platforms/mps/README.md b/scripts/platforms/mps/README.md new file mode 100644 index 000000000..56aa34145 --- /dev/null +++ b/scripts/platforms/mps/README.md @@ -0,0 +1,46 @@ +# 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 步配置用于本机快速验证。 + +启动参数与上游一致:`--model-variant fl2av` 选择基础权重,`--task t2av` 指定本次请求;AdaLN 缓存脚本也使用 `--model-variant fl2av`。配置通过 `size: [height, width]`、`num_frames` 和 `fps` 设置输出尺寸、帧数与帧率。当前文本编码器磁盘加载只支持 `t2av` 请求。 + +DiT 默认开启 `dit_mps_shared_buffer=true`,复用现有 block offload 的 `init_first_buffer → prefetch_weights → run_block → swap_blocks` 流程。两套 MPS 权重 buffer 交替使用:后台线程通过 CPU 共享视图,将 safetensors 数据直接读入空闲 buffer;GPU 同时计算当前 block。交换前等待 GPU 计算和后台读取完成,并同步 CPU 写入,再复用上一套 buffer。无需中间 CPU 权重副本,也无需另建 GPU stream。 + +共享视图依赖 `torch.mps._host_alias_storage`(PyTorch 2.13 起提供的私有接口,本机使用 2.14.0 验证);接口缺失时会明确报错。此模式要求 `cpu_offload=true`、`offload_granularity="block"`、`dit_disk_streaming=true`,文件与推理 dtype 一致,并沿用现有的非量化、AdaLN 缓存配置。设置 `dit_mps_shared_buffer=false` 可使用原来的单 buffer 磁盘加载路径。当前 H3 配置的两套 block 权重合计约 1.44 GiB,进入 VAE 阶段前会等待预读结束并释放共享视图和设备 buffer;文本编码器和 VAE 的加载方式保持不变。 + +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 new file mode 100755 index 000000000..5159eca87 --- /dev/null +++ b/scripts/platforms/mps/run_minimax_h3_t2av.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +# 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 PYTHONPATH="${lightx2v_path}:$PYTHONPATH" + +prompt='A cinematic fox walking through a snowy forest' + +mkdir -p "${lightx2v_path}/save_results" + +/opt/miniconda3/envs/torch/bin/python -m lightx2v.infer \ + --model_cls minimax_h3 \ + --model-variant fl2av \ + --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/tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh b/tools/cache_minimax_h3_adaln/run_cache_minimax_h3_adaln.sh index 22a8b568e..129a75cce 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}" \ --model-variant fl2av