From 0d92a0c74273c7eb5c5d4ad15d72aa4d7e3ac28e Mon Sep 17 00:00:00 2001 From: root Date: Tue, 4 Aug 2026 02:38:57 +0000 Subject: [PATCH 1/2] fix multi talk oom --- .../4090/infinitetalk_mutli_distilled.json | 12 ++- .../infer/infinitetalk/transformer_infer.py | 82 +++++++++++++------ 2 files changed, 67 insertions(+), 27 deletions(-) diff --git a/configs/infinitetalk/4090/infinitetalk_mutli_distilled.json b/configs/infinitetalk/4090/infinitetalk_mutli_distilled.json index 35a3c17d1..9f05b5939 100755 --- a/configs/infinitetalk/4090/infinitetalk_mutli_distilled.json +++ b/configs/infinitetalk/4090/infinitetalk_mutli_distilled.json @@ -7,13 +7,13 @@ "infinitetalk_mode": "multi", "audio_type": "para", "infinitetalk_size": "infinitetalk-720", - "dit_quantized_ckpt": "/path/to/InfiniteTalk/seko/InfiniteTalk-4StepDistill-Mean-w-ITAudioAdaptorV6.1-fp8.safetensors", + "dit_quantized_ckpt": "/models/infinitetalk_models_0703/seko/InfiniteTalk-4StepDistill-Mean-w-ITAudioAdaptorV6.1-fp8.safetensors", "dit_quantized": true, "dit_quant_scheme": "fp8-q8f", - "adapter_model_path": "/path/to/InfiniteTalk/multi/multi/infinitetalk-fp8.safetensors", + "adapter_model_path": "/models/infinitetalk_models_0703/multi/multi/infinitetalk-fp8.safetensors", "adapter_quantized": true, "adapter_quant_scheme": "fp8-q8f", - "audio_encoder_path": "/path/to/InfiniteTalk/TencentGameMate/chinese-wav2vec2-base", + "audio_encoder_path": "/models/infinitetalk_models_0703/TencentGameMate/chinese-wav2vec2-base", "clip_quantized": true, "clip_quant_scheme": "fp8-q8f", "t5_quantized": true, @@ -44,5 +44,9 @@ "infinitetalk_audio_output_dim": 768, "norm_output_audio": true, "mxfp8_fuse_enable": false, - "use_timestep_transform": true + "use_timestep_transform": true, + "parallel": { + "seq_p_size": 8, + "seq_p_attn_type": "ulysses" + } } diff --git a/lightx2v/models/networks/wan/infer/infinitetalk/transformer_infer.py b/lightx2v/models/networks/wan/infer/infinitetalk/transformer_infer.py index a39b313eb..c76b34e16 100755 --- a/lightx2v/models/networks/wan/infer/infinitetalk/transformer_infer.py +++ b/lightx2v/models/networks/wan/infer/infinitetalk/transformer_infer.py @@ -48,6 +48,10 @@ def _seq_parallel_token_count(self, pre_infer_out): return int(grid_t * grid_h * grid_w) def _seq_parallel_gather_tokens(self, x): + # Reference attention is evaluated before Ulysses exchanges sequence + # shards for head shards. Gather only the token dimension here and + # keep every attention head resident on every SP rank. + x = x.contiguous() gathered = [torch.empty_like(x) for _ in range(dist.get_world_size(self.seq_p_group))] dist.all_gather(gathered, x, group=self.seq_p_group) return torch.cat(gathered, dim=0) @@ -115,6 +119,8 @@ def infer_self_attn(self, phase, x, shift_msa, scale_msa, pre_infer_out): map_k = self._seq_parallel_gather_tokens(k)[:token_count] else: map_q, map_k = q, k + if map_q.shape[1] != self.num_heads or map_k.shape[1] != self.num_heads: + raise RuntimeError(f"InfiniteTalk reference attention requires all heads on every SP rank; expected {self.num_heads}, got q={map_q.shape[1]} and k={map_k.shape[1]}.") x_ref_attn_map = self._get_attn_map_with_target(map_q.unsqueeze(0), map_k.unsqueeze(0), pre_infer_out.grid_sizes.tuple, ref_target_masks) img_qkv_len = q.shape[0] @@ -151,38 +157,68 @@ def infer_self_attn(self, phase, x, shift_msa, scale_msa, pre_infer_out): y = phase.self_attn_o.apply(attn_out) return y, x_ref_attn_map - def _get_attn_map_with_target(self, visual_q, ref_k, shape, ref_target_masks, split_num=2): + def _get_attn_map_with_target(self, visual_q, ref_k, shape, ref_target_masks): _, grid_h, grid_w = shape ref_seqlen = grid_h * grid_w - ref_k = ref_k[:, :ref_seqlen] + device = visual_q.device + visual_q = visual_q.to(dtype=GET_DTYPE()) + ref_k = ref_k[:, :ref_seqlen].to(device=device, dtype=GET_DTYPE()) + ref_target_masks = ref_target_masks.to(device=device, dtype=GET_DTYPE()) _, seq_lens, heads, head_dim = visual_q.shape class_num, _ = ref_target_masks.shape - x_ref_attn_maps = torch.zeros(class_num, seq_lens, device=visual_q.device, dtype=visual_q.dtype) - split_chunk = max(1, heads // split_num) - split_count = 0 - for start in range(0, heads, split_chunk): - end = min(start + split_chunk, heads) + x_ref_attn_maps = torch.zeros(class_num, seq_lens, device=device, dtype=GET_DTYPE()) + # InfiniteTalk's reference implementation always uses two equal head + # groups. This is a local, sequential memory split and is unrelated to + # the Ulysses SP world size; never shard these heads across SP ranks. + split_num = 2 + if heads % split_num != 0: + raise ValueError(f"Reference attention heads ({heads}) must be divisible by the fixed split_num={split_num}.") + split_chunk = heads // split_num + for split_idx in range(split_num): + start = split_idx * split_chunk + end = (split_idx + 1) * split_chunk maps = self._calculate_x_ref_attn_map(visual_q[:, :, start:end], ref_k[:, :, start:end], ref_target_masks, head_dim) x_ref_attn_maps += maps - split_count += 1 - return x_ref_attn_maps / max(1, split_count) + return x_ref_attn_maps / split_num @staticmethod def _calculate_x_ref_attn_map(visual_q, ref_k, ref_target_masks, head_dim): - ref_k = ref_k.to(visual_q.dtype).to(visual_q.device) - visual_q = (visual_q * (head_dim**-0.5)).transpose(1, 2) - ref_k = ref_k.transpose(1, 2) - attn = visual_q @ ref_k.transpose(-2, -1) - attn = attn.softmax(-1) - ref_target_masks = ref_target_masks.to(visual_q.dtype).to(visual_q.device) - - x_ref_attn_maps = [] - for ref_target_mask in ref_target_masks: - mask = ref_target_mask[None, None, None, :] - x_ref_attnmap = (attn * mask).sum(-1) / mask.sum().clamp_min(1.0) - x_ref_attnmap = x_ref_attnmap.permute(0, 2, 1).mean(-1) - x_ref_attn_maps.append(x_ref_attnmap) - return torch.concat(x_ref_attn_maps, dim=0) + scale = visual_q.new_tensor(head_dim**-0.5) + visual_q = (visual_q * scale).transpose(1, 2).contiguous() + ref_k = ref_k.transpose(1, 2).contiguous() + + batch_size, heads, ref_seqlen, value_dim = ref_k.shape + class_num, mask_seqlen = ref_target_masks.shape + if mask_seqlen != ref_seqlen: + raise ValueError(f"Reference mask length ({mask_seqlen}) does not match reference K length ({ref_seqlen}).") + if class_num > value_dim: + raise ValueError(f"Reference mask count ({class_num}) exceeds attention head dimension ({value_dim}).") + + # The required map is softmax(QK^T) @ mask. Use each target mask as + # a value channel so fused SDPA can compute it without materializing + # the B x H x query_len x ref_len attention matrix. Pad V to the Q/K + # head dimension for compatibility with fused CUDA kernels. + mask_values = ref_k.new_zeros((batch_size, heads, ref_seqlen, value_dim)) + mask_values[..., :class_num] = ref_target_masks.transpose(0, 1)[None, None, :, :] + + # Q was scaled above to match InfiniteTalk's operation order, so SDPA + # must not apply the default head_dim**-0.5 scale a second time. Keep + # the math backend disabled: it may materialize the full attention map. + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=True): + mask_attn = F.scaled_dot_product_attention( + visual_q, + ref_k, + mask_values, + dropout_p=0.0, + is_causal=False, + scale=1.0, + ) + + mask_attn = mask_attn[..., :class_num] + mask_sums = ref_target_masks.sum(-1).clamp_min(1.0) + mask_attn = mask_attn / mask_sums[None, None, None, :] + mask_attn = mask_attn.mean(dim=1) # B, query_len, class_num + return mask_attn.permute(2, 0, 1).reshape(class_num * batch_size, -1) def infer_audio_cross_attn(self, phase, x, pre_infer_out, x_ref_attn_map): audio_embedding = pre_infer_out.adapter_args["audio_embedding"].to(device=x.device, dtype=GET_DTYPE()) From 40acc4f4056a54d4fc3105475ff04fdbb81c6ddd Mon Sep 17 00:00:00 2001 From: llmc-reviewer Date: Tue, 4 Aug 2026 10:41:00 +0800 Subject: [PATCH 2/2] Update infinitetalk_mutli_distilled.json --- configs/infinitetalk/4090/infinitetalk_mutli_distilled.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/configs/infinitetalk/4090/infinitetalk_mutli_distilled.json b/configs/infinitetalk/4090/infinitetalk_mutli_distilled.json index 9f05b5939..0fd5457fe 100755 --- a/configs/infinitetalk/4090/infinitetalk_mutli_distilled.json +++ b/configs/infinitetalk/4090/infinitetalk_mutli_distilled.json @@ -7,13 +7,13 @@ "infinitetalk_mode": "multi", "audio_type": "para", "infinitetalk_size": "infinitetalk-720", - "dit_quantized_ckpt": "/models/infinitetalk_models_0703/seko/InfiniteTalk-4StepDistill-Mean-w-ITAudioAdaptorV6.1-fp8.safetensors", + "dit_quantized_ckpt": "/path/to/InfiniteTalk/seko/InfiniteTalk-4StepDistill-Mean-w-ITAudioAdaptorV6.1-fp8.safetensors", "dit_quantized": true, "dit_quant_scheme": "fp8-q8f", - "adapter_model_path": "/models/infinitetalk_models_0703/multi/multi/infinitetalk-fp8.safetensors", + "adapter_model_path": "/path/to/InfiniteTalk/multi/multi/infinitetalk-fp8.safetensors", "adapter_quantized": true, "adapter_quant_scheme": "fp8-q8f", - "audio_encoder_path": "/models/infinitetalk_models_0703/TencentGameMate/chinese-wav2vec2-base", + "audio_encoder_path": "/path/to/InfiniteTalk/TencentGameMate/chinese-wav2vec2-base", "clip_quantized": true, "clip_quant_scheme": "fp8-q8f", "t5_quantized": true,