(x, freqs, num_heads)
| 25 | return padded_tensor |
| 26 | |
| 27 | def rope_apply(x, freqs, num_heads): |
| 28 | x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) |
| 29 | s_per_rank = x.shape[1] |
| 30 | |
| 31 | x_out = torch.view_as_complex(x.to(torch.float64).reshape( |
| 32 | x.shape[0], x.shape[1], x.shape[2], -1, 2)) |
| 33 | |
| 34 | sp_size = get_sequence_parallel_world_size() |
| 35 | sp_rank = get_sequence_parallel_rank() |
| 36 | freqs = pad_freqs(freqs, s_per_rank * sp_size) |
| 37 | freqs_rank = freqs[(sp_rank * s_per_rank):((sp_rank + 1) * s_per_rank), :, :] |
| 38 | |
| 39 | x_out = torch.view_as_real(x_out * freqs_rank).flatten(2) |
| 40 | return x_out.to(x.dtype) |
| 41 | |
| 42 | def usp_dit_forward(self, |
| 43 | x: torch.Tensor, |
no test coverage detected