(
self,
x,
e,
seq_lens,
freqs,
context,
context_lens,
)
| 419 | self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) |
| 420 | |
| 421 | def forward( |
| 422 | self, |
| 423 | x, |
| 424 | e, |
| 425 | seq_lens, |
| 426 | freqs, |
| 427 | context, |
| 428 | context_lens, |
| 429 | ): |
| 430 | assert e.dtype == torch.float32 |
| 431 | with amp.autocast(dtype=torch.float32, device_type="cuda"): |
| 432 | e = (self.modulation.to(dtype=e.dtype, device=e.device) + e).chunk(6, dim=1) |
| 433 | assert e[0].dtype == torch.float32 |
| 434 | |
| 435 | # self-attention |
| 436 | y = self.self_attn( |
| 437 | self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, |
| 438 | freqs) |
| 439 | with amp.autocast(dtype=torch.float32, device_type="cuda"): |
| 440 | x = x + y * e[2] |
| 441 | |
| 442 | # cross-attention & ffn function |
| 443 | def cross_attn_ffn(x, context, context_lens, e): |
| 444 | x = x + self.cross_attn(self.norm3(x), context, context_lens) |
| 445 | y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3]) |
| 446 | with amp.autocast(dtype=torch.float32, device_type="cuda"): |
| 447 | x = x + y * e[5] |
| 448 | return x |
| 449 | |
| 450 | x = cross_attn_ffn(x, context, context_lens, e) |
| 451 | return x |
| 452 | |
| 453 | |
| 454 | class Head(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected