From 7161ecbc51e2996ca80655ec3e1ba8078516e95e Mon Sep 17 00:00:00 2001 From: zhangbilang <2045955563@qq.com> Date: Fri, 7 Aug 2026 07:01:44 +0000 Subject: [PATCH] minimax_h3 warmup and compile --- .../minimax_h3/minimax_h3_fl2av_compile.json | 27 ++++++++ .../minimax_h3/minimax_h3_i2av_compile.json | 27 ++++++++ .../minimax_h3/minimax_h3_l2av_compile.json | 27 ++++++++ .../minimax_h3/minimax_h3_ref2av_compile.json | 27 ++++++++ .../minimax_h3/minimax_h3_t2av_compile.json | 27 ++++++++ .../infer/offload/transformer_infer.py | 7 ++ .../minimax_h3/infer/transformer_infer.py | 5 +- .../runners/minimax_h3/minimax_h3_runner.py | 68 ++++++++++++++++++- .../minimax_h3/run_minimax_h3_t2av_warmup.sh | 24 +++++++ 9 files changed, 236 insertions(+), 3 deletions(-) create mode 100644 configs/minimax_h3/minimax_h3_fl2av_compile.json create mode 100644 configs/minimax_h3/minimax_h3_i2av_compile.json create mode 100644 configs/minimax_h3/minimax_h3_l2av_compile.json create mode 100644 configs/minimax_h3/minimax_h3_ref2av_compile.json create mode 100644 configs/minimax_h3/minimax_h3_t2av_compile.json create mode 100755 scripts/minimax_h3/run_minimax_h3_t2av_warmup.sh diff --git a/configs/minimax_h3/minimax_h3_fl2av_compile.json b/configs/minimax_h3/minimax_h3_fl2av_compile.json new file mode 100644 index 000000000..2c76bc35f --- /dev/null +++ b/configs/minimax_h3/minimax_h3_fl2av_compile.json @@ -0,0 +1,27 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 544, + "target_width": 960, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "sage_attn2", + "rms_type": "sgl-kernel", + "rope_type": "torch_real_rope", + "feature_caching": "NoCaching", + "use_compile": true, + "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/minimax_h3/minimax_h3_i2av_compile.json b/configs/minimax_h3/minimax_h3_i2av_compile.json new file mode 100644 index 000000000..2c76bc35f --- /dev/null +++ b/configs/minimax_h3/minimax_h3_i2av_compile.json @@ -0,0 +1,27 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 544, + "target_width": 960, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "sage_attn2", + "rms_type": "sgl-kernel", + "rope_type": "torch_real_rope", + "feature_caching": "NoCaching", + "use_compile": true, + "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/minimax_h3/minimax_h3_l2av_compile.json b/configs/minimax_h3/minimax_h3_l2av_compile.json new file mode 100644 index 000000000..2c76bc35f --- /dev/null +++ b/configs/minimax_h3/minimax_h3_l2av_compile.json @@ -0,0 +1,27 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 544, + "target_width": 960, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "sage_attn2", + "rms_type": "sgl-kernel", + "rope_type": "torch_real_rope", + "feature_caching": "NoCaching", + "use_compile": true, + "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/minimax_h3/minimax_h3_ref2av_compile.json b/configs/minimax_h3/minimax_h3_ref2av_compile.json new file mode 100644 index 000000000..2c76bc35f --- /dev/null +++ b/configs/minimax_h3/minimax_h3_ref2av_compile.json @@ -0,0 +1,27 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 544, + "target_width": 960, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "sage_attn2", + "rms_type": "sgl-kernel", + "rope_type": "torch_real_rope", + "feature_caching": "NoCaching", + "use_compile": true, + "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/minimax_h3/minimax_h3_t2av_compile.json b/configs/minimax_h3/minimax_h3_t2av_compile.json new file mode 100644 index 000000000..2c76bc35f --- /dev/null +++ b/configs/minimax_h3/minimax_h3_t2av_compile.json @@ -0,0 +1,27 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 544, + "target_width": 960, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": true, + "offload_granularity": "model", + "text_encoder_cpu_offload": true, + "vae_cpu_offload": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "sage_attn2", + "rms_type": "sgl-kernel", + "rope_type": "torch_real_rope", + "feature_caching": "NoCaching", + "use_compile": true, + "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/infer/offload/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/offload/transformer_infer.py index 6f81670ca..fe8b77d15 100644 --- a/lightx2v/models/networks/minimax_h3/infer/offload/transformer_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/offload/transformer_infer.py @@ -19,6 +19,13 @@ def __init__(self, config): elif offload_granularity != "model": raise NotImplementedError(f"MiniMax-H3 does not support offload_granularity={offload_granularity!r}") + def get_compile_block_key(self, block_idx, block): + # block offload + if hasattr(self, "offload_manager"): + return id(block) + # model offload + return super().get_compile_block_key(block_idx, block) + def infer_with_blocks_offload(self, blocks, hidden_states, pre_infer_out): num_blocks = len(blocks) current_stream = torch_device_module.current_stream() diff --git a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py index 9577299e9..2b3c74267 100644 --- a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py @@ -21,6 +21,7 @@ def __init__(self, config): self.tp_rank = 0 self.num_heads = self.global_num_heads // self.tp_size self.head_dim = int(config.get("attention_head_dim", 128)) + self.infer_dtype = GET_DTYPE() if config.get("seq_parallel", False): self.seq_p_group = config["device_mesh"].get_group(mesh_dim="seq_p") parallel = config.get("parallel", {}) @@ -90,7 +91,7 @@ def _attention(self, weights, hidden_states, pre_infer_out): ) if aux_out is not None: out = torch.cat((aux_out, out), dim=0) - return weights.to_out.apply(out.to(GET_DTYPE())) + return weights.to_out.apply(out.to(self.infer_dtype)) @staticmethod def _ff(weights, hidden_states): @@ -100,7 +101,7 @@ def _ff(weights, hidden_states): def infer_block(self, weights, hidden_states, pre_infer_out): # Activation is evaluated in fp32, then cast to the inference dtype # immediately before the (possibly quantized) AdaLN projection. - modulation = weights.adaln.apply(F.silu(pre_infer_out.temb).to(GET_DTYPE())) + modulation = weights.adaln.apply(F.silu(pre_infer_out.temb).to(self.infer_dtype)) modulation = self._gather_tp_last_dim(modulation) modulation = modulation.view(-1, 6 * self.hidden_size) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation.chunk(6, dim=-1) diff --git a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py index 0ab310b0f..b4b969336 100644 --- a/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py +++ b/lightx2v/models/runners/minimax_h3/minimax_h3_runner.py @@ -61,6 +61,9 @@ class MiniMaxH3Runner(DefaultRunner): graph. """ + _WARMUP_RESOLUTIONS = ((480, 480), (544, 960)) + _WARMUP_TASKS = ("t2av", "fl2av", "i2av", "l2av", "ref2av") + def __init__(self, config): if config.get("task") not in {"t2av", "i2av", "l2av", "fl2av", "ref2av"}: raise ValueError("MiniMax-H3 supports t2av/i2av/l2av/fl2av/ref2av") @@ -73,8 +76,71 @@ def init_modules(self): super().init_modules() self.run_input_encoder = self._run_input_encoder_local_h3 + @ProfilingContext4DebugL1("Warmup") def run_warmup(self): - raise NotImplementedError("MiniMax-H3 warmup is not implemented") + task = self.config["task"] + if task not in self._WARMUP_TASKS: + raise NotImplementedError(f"MiniMax-H3 warmup does not support task: {task}") + + for height, width in self._WARMUP_RESOLUTIONS: + logger.info(f"Warmup: {height}x{width}") + transformer_offloaded = not self.config.get("cpu_offload", False) + try: + self.scheduler.generator = None + self._prepare_warmup_inputs(height, width) + self.inputs = self._run_input_encoder_local_h3() + self.init_run() + + self.scheduler.step_pre(0) + self.model.infer(self.inputs) + self.scheduler.step_post() + video_rows = self.scheduler.video_latents + audio_rows = self.scheduler.audio_latents + + if self.config.get("cpu_offload", False): + self._offload_transformer() + transformer_offloaded = True + + self.run_vae_decoder(video_rows, audio_rows) + torch_device_module.synchronize() + finally: + if self.config.get("cpu_offload", False) and not transformer_offloaded: + with suppress(Exception): + self._offload_transformer() + self.clear_warmup_state() + + logger.info("[Warmup] Warmup completed") + self._maybe_freeze_gc() + + def _prepare_warmup_inputs(self, height, width): + task = self.config["task"] + common = { + "seed": 0, + "prompt": "warmup" if (height, width) == self._WARMUP_RESOLUTIONS[0] else "A cinematic fox walking through a snowy forest.", + "target_shape": [height, width], + "target_video_length": int(self.config.get("target_video_length", 124)), + "return_result_tensor": True, + } + image = Image.new("RGB", (width, height), color=0) + if task == "t2av": + self.input_info = T2AVInputInfo(**common) + elif task == "i2av": + self.input_info = I2AVInputInfo(**common, image_path=image) + elif task == "l2av": + self.input_info = L2AVInputInfo(**common, last_frame_path=image) + elif task == "fl2av": + self.input_info = FL2AVInputInfo(**common, image_path=image, last_frame_path=image.copy()) + else: + self.input_info = Ref2AVInputInfo(**common, image_path=image) + + def clear_warmup_state(self): + self.scheduler.clear() + self.condition_video_latents = [] + self.condition_audio_latents = [] + self.keyframe_anchors = () + self.prepared_references = None + self.input_info = None + self.__dict__.pop("inputs", None) def init_scheduler(self): self.scheduler = MiniMaxH3Scheduler(self.config) diff --git a/scripts/minimax_h3/run_minimax_h3_t2av_warmup.sh b/scripts/minimax_h3/run_minimax_h3_t2av_warmup.sh new file mode 100755 index 000000000..e993d3150 --- /dev/null +++ b/scripts/minimax_h3/run_minimax_h3_t2av_warmup.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +# set path firstly +lightx2v_path= +model_path= + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 + +prompt='A cinematic fox walking through a snowy forest.' + +python -m lightx2v.infer \ +--model_cls minimax_h3 \ +--task t2av \ +--model_path $model_path \ +--config_json ${lightx2v_path}/configs/minimax_h3/minimax_h3_t2av_compile.json \ +--warmup \ +--prompt "$prompt" \ +--save_result_path ${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av10.mp4 \ +--seed 0