| 662 | |
| 663 | |
| 664 | class FlaxLLaMABlock(nn.Module): |
| 665 | config: LLaMAConfig |
| 666 | dtype: jnp.dtype=jnp.float32 |
| 667 | param_dtype: jnp.dtype=jnp.float32 |
| 668 | precision: Optional[Union[jax.lax.Precision, str]]=None |
| 669 | |
| 670 | def setup(self) -> None: |
| 671 | attention_module = FlaxLLaMAAttention |
| 672 | mlp_module = FlaxLLaMAMLP |
| 673 | if self.config.scan_mlp: |
| 674 | mlp_module = remat( |
| 675 | mlp_module, static_argnums=(1,), |
| 676 | policy=jax.checkpoint_policies.nothing_saveable, |
| 677 | prevent_cse=not self.config.scan_layers, |
| 678 | ) |
| 679 | self.attention = attention_module( |
| 680 | self.config, |
| 681 | dtype=self.dtype, |
| 682 | param_dtype=self.param_dtype, |
| 683 | precision=self.precision, |
| 684 | ) |
| 685 | self.feed_forward = mlp_module( |
| 686 | self.config, |
| 687 | dtype=self.dtype, |
| 688 | param_dtype=self.param_dtype, |
| 689 | precision=self.precision, |
| 690 | ) |
| 691 | self.attention_norm = RMSNorm( |
| 692 | self.config.hidden_size, |
| 693 | eps=self.config.rms_norm_eps, |
| 694 | dtype=self.dtype, |
| 695 | param_dtype=self.param_dtype, |
| 696 | ) |
| 697 | self.ffn_norm = RMSNorm( |
| 698 | self.config.hidden_size, |
| 699 | eps=self.config.rms_norm_eps, |
| 700 | dtype=self.dtype, |
| 701 | param_dtype=self.param_dtype, |
| 702 | ) |
| 703 | |
| 704 | def __call__( |
| 705 | self, |
| 706 | hidden_states, |
| 707 | attention_mask=None, |
| 708 | segment_ids=None, |
| 709 | position_ids=None, |
| 710 | deterministic: bool = True, |
| 711 | init_cache: bool = False, |
| 712 | output_attentions: bool = False, |
| 713 | ): |
| 714 | attn_outputs = self.attention( |
| 715 | self.attention_norm(hidden_states), |
| 716 | attention_mask, |
| 717 | segment_ids, |
| 718 | position_ids, |
| 719 | deterministic, |
| 720 | init_cache, |
| 721 | output_attentions, |
nothing calls this directly
no outgoing calls
no test coverage detected