| 408 | |
| 409 | |
| 410 | class RoPE3D(RoPE1D): |
| 411 | def __init__(self, freq=1e4, F0=1.0, scaling_factor=1.0): |
| 412 | super(RoPE3D, self).__init__(freq, F0, scaling_factor) |
| 413 | self.position_cache = {} |
| 414 | |
| 415 | def get_mesh_3d(self, rope_positions, bsz): |
| 416 | f, h, w = rope_positions |
| 417 | |
| 418 | if f"{f}-{h}-{w}" not in self.position_cache: |
| 419 | x = torch.arange(f, device='cpu') |
| 420 | y = torch.arange(h, device='cpu') |
| 421 | z = torch.arange(w, device='cpu') |
| 422 | self.position_cache[f"{f}-{h}-{w}"] = torch.cartesian_prod(x, y, z).view(1, f*h*w, 3).expand(bsz, -1, 3) |
| 423 | return self.position_cache[f"{f}-{h}-{w}"] |
| 424 | |
| 425 | def __call__(self, tokens, rope_positions, ch_split, parallel=False): |
| 426 | """ |
| 427 | input: |
| 428 | * tokens: batch_size x ntokens x nheads x dim |
| 429 | * rope_positions: list of (f, h, w) |
| 430 | output: |
| 431 | * tokens after applying RoPE2D (batch_size x ntokens x nheads x dim) |
| 432 | """ |
| 433 | assert sum(ch_split) == tokens.size(-1); |
| 434 | |
| 435 | mesh_grid = self.get_mesh_3d(rope_positions, bsz=tokens.shape[0]) |
| 436 | out = [] |
| 437 | for i, (D, x) in enumerate(zip(ch_split, torch.split(tokens, ch_split, dim=-1))): |
| 438 | cos, sin = self.get_cos_sin(D, int(mesh_grid.max()) + 1, tokens.device, tokens.dtype) |
| 439 | |
| 440 | if parallel: |
| 441 | pass |
| 442 | else: |
| 443 | mesh = mesh_grid[:, :, i].clone() |
| 444 | x = self.apply_rope1d(x, mesh.to(tokens.device), cos, sin) |
| 445 | out.append(x) |
| 446 | |
| 447 | tokens = torch.cat(out, dim=-1) |
| 448 | return tokens |
| 449 | |
| 450 | |
| 451 | class SelfAttention(Attention): |