FL2VA/video_vae/attention.py
5.6 KB · 164 lines · python Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # Attention module for the MiniMax H3 visual VAE (inference-only bundle).
3 import os
4 import torch
5 import torch.nn as nn
6 import torch.distributed as dist
7 from typing import Optional
8 from diffusers.utils import logging
9
10 from .parallel import all_to_all_4D, get_parallel_state
11 from .func import apply_rotary_pos_emb
12 from .flash import flash_attn
13
14 logger = logging.get_logger(__name__) # pylint: disable=invalid-name
15
16
17 def _env_flag(name, default="0"):
18 value = os.environ.get(name, default)
19 return str(value).strip().lower() in ("1", "true", "yes", "on")
20
21
22 def _vit_norm_input(module, hidden_states):
23 if _env_flag("MINIMAX_H3_VAE_DECODER_VIT_FP32_NORM", "1"):
24 return hidden_states.float()
25 weight = getattr(module, "weight", None)
26 return hidden_states.to(getattr(weight, "dtype", hidden_states.dtype))
27
28
29 def maybe_checkpoint(owner, function, *args):
30 if owner.training and getattr(owner, "gradient_checkpointing", False):
31 raise NotImplementedError(
32 "gradient checkpointing is not supported in this inference-only bundle"
33 )
34 return function(*args)
35
36
37 class Attention(nn.Module):
38 def __init__(
39 self,
40 heads,
41 dim_head,
42 embed_dim: Optional[int] = None,
43 qk_norm_type: Optional[str] = None,
44 qk_norm_affine: bool = False,
45 bias: bool = True,
46 out_bias: Optional[bool] = None,
47 eps: float = 1e-5,
48 **kwargs,
49 ):
50 super().__init__()
51 self.dim_head = dim_head
52 self.heads = heads
53 self.attn_inner_dim = dim_head * heads
54 self.embed_dim = embed_dim if embed_dim is not None else self.attn_inner_dim
55
56 out_bias = out_bias if out_bias is not None else bias
57
58 if qk_norm_type is None:
59 self.norm_q = None
60 self.norm_k = None
61 elif qk_norm_type == "layer_norm":
62 self.norm_q = nn.LayerNorm(
63 dim_head, eps=eps, elementwise_affine=qk_norm_affine
64 )
65 self.norm_k = nn.LayerNorm(
66 dim_head, eps=eps, elementwise_affine=qk_norm_affine
67 )
68 elif qk_norm_type == "rms_norm":
69 self.norm_q = nn.RMSNorm(
70 dim_head, eps=eps, elementwise_affine=qk_norm_affine
71 )
72 self.norm_k = nn.RMSNorm(
73 dim_head, eps=eps, elementwise_affine=qk_norm_affine
74 )
75 else:
76 raise ValueError(
77 f"unknown qk_norm_type: {qk_norm_type}. Should be None,'layer_norm','rms_norm'"
78 )
79
80 self.to_qkv = nn.Linear(self.embed_dim, self.attn_inner_dim * 3, bias=bias)
81
82 self.to_out = nn.Linear(self.attn_inner_dim, self.embed_dim, bias=out_bias)
83
84 self.spatial_parallel = get_parallel_state().get("sp_enabled", False)
85
86 state = get_parallel_state()
87 sp_size = state.get("sp_size", 1)
88 tp_size = state.get("tp_size", 1)
89 parallel_size = sp_size * tp_size
90 if parallel_size > 1 and self.heads % parallel_size != 0:
91 raise ValueError(
92 f"num_heads {self.heads} must be divisible by sp_size * tp_size ({sp_size} * {tp_size} = {parallel_size})"
93 )
94
95 if len(kwargs) > 0 and (not dist.is_initialized() or dist.get_rank() == 0):
96 logger.warning(f"Unused kwargs: {kwargs}")
97
98 def _perform_attention(self, query, key, value, pack_info):
99 cu_seqlens = pack_info.get("cu_seqlens", None)
100 mask_mod = pack_info.get("mask_mod", None)
101 block_sparse = pack_info.get("block_sparse", None)
102
103 if cu_seqlens is not None:
104 raise NotImplementedError(
105 "varlen attention is not supported in this inference-only bundle"
106 )
107
108 if mask_mod is not None:
109 hidden_states = flash_attn(
110 query,
111 key,
112 value,
113 mask_mod=mask_mod,
114 block_sparse=block_sparse,
115 )
116 else:
117 hidden_states = flash_attn(
118 query,
119 key,
120 value,
121 )
122
123 return hidden_states
124
125 def perform_attention(self, query, key, value, pack_info={}):
126 return self._perform_attention(query, key, value, pack_info)
127
128 def forward(
129 self,
130 hidden_states: torch.Tensor,
131 rotary_pos_emb: Optional[torch.Tensor] = None,
132 pack_info: dict = {},
133 ) -> torch.Tensor:
134 batch_size, seq_len, _ = hidden_states.shape
135
136 qkv = self.to_qkv(hidden_states)
137 qkv = qkv.view(batch_size, seq_len, -1, 3 * self.dim_head)
138 query, key, value = torch.chunk(qkv, 3, dim=-1)
139
140 if self.spatial_parallel:
141 local_process_group = get_parallel_state()["sp_process_group"]
142 query = all_to_all_4D(query, 2, 1, group=local_process_group)
143 key = all_to_all_4D(key, 2, 1, group=local_process_group)
144 value = all_to_all_4D(value, 2, 1, group=local_process_group)
145
146 if self.norm_q is not None:
147 query = self.norm_q(_vit_norm_input(self.norm_q, query)).to(query.dtype)
148 if self.norm_k is not None:
149 key = self.norm_k(_vit_norm_input(self.norm_k, key)).to(key.dtype)
150
151 if rotary_pos_emb is not None:
152 query = apply_rotary_pos_emb(query, rotary_pos_emb)
153 key = apply_rotary_pos_emb(key, rotary_pos_emb)
154
155 hidden_states = self.perform_attention(query, key, value, pack_info)
156
157 if self.spatial_parallel:
158 hidden_states = all_to_all_4D(hidden_states, 1, 2, group=local_process_group)
159
160 hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
161 hidden_states = self.to_out(hidden_states)
162
163 return hidden_states
164