Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| import torch.cuda.amp as amp | |
| from .fuser import (get_sequence_parallel_rank, | |
| get_sequence_parallel_world_size, get_sp_group, | |
| init_distributed_environment, initialize_model_parallel, | |
| xFuserLongContextAttention) | |
| def pad_freqs(original_tensor, target_len): | |
| seq_len, s1, s2 = original_tensor.shape | |
| pad_size = target_len - seq_len | |
| padding_tensor = torch.ones( | |
| pad_size, | |
| s1, | |
| s2, | |
| dtype=original_tensor.dtype, | |
| device=original_tensor.device) | |
| padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0) | |
| return padded_tensor | |
| def rope_apply(x, grid_sizes, freqs): | |
| """ | |
| x: [B, L, N, C]. | |
| grid_sizes: [B, 3]. | |
| freqs: [M, C // 2]. | |
| """ | |
| s, n, c = x.size(1), x.size(2), x.size(3) // 2 | |
| # split freqs | |
| freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) | |
| # loop over samples | |
| output = [] | |
| for i, (f, h, w) in enumerate(grid_sizes.tolist()): | |
| seq_len = f * h * w | |
| # precompute multipliers | |
| x_i = torch.view_as_complex(x[i, :s].to(torch.float32).reshape( | |
| s, n, -1, 2)) | |
| freqs_i = torch.cat([ | |
| freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), | |
| freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), | |
| freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) | |
| ], | |
| dim=-1).reshape(seq_len, 1, -1) | |
| # apply rotary embedding | |
| sp_size = get_sequence_parallel_world_size() | |
| sp_rank = get_sequence_parallel_rank() | |
| freqs_i = pad_freqs(freqs_i, s * sp_size) | |
| s_per_rank = s | |
| freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) * | |
| s_per_rank), :, :] | |
| x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2) | |
| x_i = torch.cat([x_i, x[i, s:]]) | |
| # append to collection | |
| output.append(x_i) | |
| return torch.stack(output) | |
| def rope_apply_qk(q, k, grid_sizes, freqs): | |
| q = rope_apply(q, grid_sizes, freqs) | |
| k = rope_apply(k, grid_sizes, freqs) | |
| return q, k | |
| def usp_attn_forward(self, | |
| x, | |
| seq_lens, | |
| grid_sizes, | |
| freqs, | |
| dtype=torch.bfloat16, | |
| t=0): | |
| b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim | |
| half_dtypes = (torch.float16, torch.bfloat16) | |
| def half(x): | |
| return x if x.dtype in half_dtypes else x.to(dtype) | |
| # query, key, value function | |
| def qkv_fn(x): | |
| q = self.norm_q(self.q(x)).view(b, s, n, d) | |
| k = self.norm_k(self.k(x)).view(b, s, n, d) | |
| v = self.v(x).view(b, s, n, d) | |
| return q, k, v | |
| q, k, v = qkv_fn(x) | |
| q, k = rope_apply_qk(q, k, grid_sizes, freqs) | |
| # TODO: We should use unpaded q,k,v for attention. | |
| # k_lens = seq_lens // get_sequence_parallel_world_size() | |
| # if k_lens is not None: | |
| # q = torch.cat([u[:l] for u, l in zip(q, k_lens)]).unsqueeze(0) | |
| # k = torch.cat([u[:l] for u, l in zip(k, k_lens)]).unsqueeze(0) | |
| # v = torch.cat([u[:l] for u, l in zip(v, k_lens)]).unsqueeze(0) | |
| x = xFuserLongContextAttention()( | |
| None, | |
| query=half(q), | |
| key=half(k), | |
| value=half(v), | |
| window_size=self.window_size) | |
| # TODO: padding after attention. | |
| # x = torch.cat([x, x.new_zeros(b, s - x.size(1), n, d)], dim=1) | |
| # output | |
| x = x.flatten(2) | |
| x = self.o(x) | |
| return x |