| 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, |