(
self,
*,
mode: ForwardMode,
input_batch: Nested[Tensor],
self_attention_logit_biases: Optional[Tensor] = None,
cross_attention_data: Optional[Tensor] = None,
cross_attention_logit_biases: Optional[Tensor] = None,
cached_states: Optional[NestedTensor] = None,
page_pool: Optional[Nested[Tensor]] = None,
)
| 546 | self, |
| 547 | *, |
| 548 | mode: ForwardMode, |
| 549 | input_batch: Nested[Tensor], |
| 550 | self_attention_logit_biases: Optional[Tensor] = None, |
| 551 | cross_attention_data: Optional[Tensor] = None, |
| 552 | cross_attention_logit_biases: Optional[Tensor] = None, |
| 553 | cached_states: Optional[NestedTensor] = None, |
| 554 | page_pool: Optional[Nested[Tensor]] = None, |
| 555 | ) -> tuple[Optional[NestedTensor], Tensor]: |
| 556 | validate_contains_paths(input_batch, paths=["input_ids"]) |
| 557 | input_segment_ids = input_batch.get("input_segment_ids", None) |
| 558 | positions = input_batch.get("positions", None) |
| 559 | |
| 560 | emb_batch = {**input_batch} |
| 561 | emb_batch["inputs"] = emb_batch["input_ids"] |
| 562 | |
| 563 | if mode == ForwardMode.FORWARD: |
| 564 | x = self.emb(input_batch=emb_batch) |
| 565 | x = self.transformer( |
| 566 | x, |
| 567 | self_attention_logit_biases=self_attention_logit_biases, |
| 568 | target_segment_ids=input_segment_ids, |
| 569 | target_positions=positions, |
| 570 | cross_attention_data=cross_attention_data, |
| 571 | cross_attention_logit_biases=cross_attention_logit_biases, |
| 572 | ) |
| 573 | cached_states = None |
| 574 | elif mode in (ForwardMode.PREFILL, ForwardMode.EXTEND_STEP): |
| 575 | assert cached_states is not None |
| 576 | cached_states["emb"], x = self.emb.extend_step( |
| 577 | cached_states=cached_states["emb"], |
| 578 | input_batch=emb_batch, |
| 579 | is_prefill=(mode == ForwardMode.PREFILL), |
| 580 | ) |
| 581 | cached_states["transformer_state"], x = self.transformer.extend_step( |
| 582 | cached_states=cached_states["transformer_state"], |
| 583 | data=x, |
| 584 | is_prefill=(mode == ForwardMode.PREFILL), |
| 585 | target_segment_ids=input_segment_ids, |
| 586 | self_attention_logit_biases=self_attention_logit_biases, |
| 587 | cross_attention_data=cross_attention_data, |
| 588 | cross_attention_logit_biases=cross_attention_logit_biases, |
| 589 | page_pool=page_pool, |
| 590 | ) |
| 591 | else: |
| 592 | raise ValueError(f"Unrecognized mode {mode}.") |
| 593 | x = x.data |
| 594 | self._add_tensor_stats("outputs", x) |
| 595 | |
| 596 | if "output_norm" in self.children: |
| 597 | x = self.output_norm(x) |
| 598 | self._add_tensor_stats("norm_outputs", x) |
| 599 | x = self.output_dropout(x) |
| 600 | |
| 601 | return cached_states, dict(hidden_states=x) |
| 602 | |
| 603 | def compute_logits(self, forward_outputs: Nested[Tensor]) -> Tensor: |
| 604 | """Computes logits from decoder forward outputs. |
| 605 |
no test coverage detected