| 980 | |
| 981 | |
| 982 | class FlaxLLaMAModule(nn.Module): |
| 983 | config: LLaMAConfig |
| 984 | dtype: jnp.dtype = jnp.float32 |
| 985 | param_dtype: jnp.dtype=jnp.float32 |
| 986 | precision: Optional[Union[jax.lax.Precision, str]]=None |
| 987 | |
| 988 | def setup(self): |
| 989 | self.embed_dim = self.config.hidden_size |
| 990 | |
| 991 | self.wte = nn.Embed( |
| 992 | self.config.vocab_size, |
| 993 | self.config.hidden_size, |
| 994 | embedding_init=jax.nn.initializers.normal(stddev=self.config.initializer_range), |
| 995 | dtype=self.dtype, |
| 996 | param_dtype=self.param_dtype, |
| 997 | ) |
| 998 | self.dropout = nn.Dropout(rate=self.config.embd_pdrop) |
| 999 | self.h = FlaxLLaMABlockCollection(self.config, dtype=self.dtype, param_dtype=self.param_dtype, precision=self.precision) |
| 1000 | self.ln_f = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps, dtype=self.dtype, param_dtype=self.param_dtype) |
| 1001 | |
| 1002 | def __call__( |
| 1003 | self, |
| 1004 | input_ids, |
| 1005 | attention_mask, |
| 1006 | segment_ids, |
| 1007 | position_ids, |
| 1008 | deterministic=True, |
| 1009 | init_cache: bool = False, |
| 1010 | output_attentions: bool = False, |
| 1011 | output_hidden_states: bool = False, |
| 1012 | return_dict: bool = True, |
| 1013 | ): |
| 1014 | input_embeds = self.wte(input_ids.astype("i4")) |
| 1015 | assert input_embeds.shape[1] <= self.config.max_sequence_length, f"Input sequence length {input_embeds.shape[1]} larger than max supported sequence length {self.config.max_sequence_length}" |
| 1016 | |
| 1017 | hidden_states = self.dropout(input_embeds, deterministic=deterministic) |
| 1018 | |
| 1019 | outputs = self.h( |
| 1020 | hidden_states, |
| 1021 | attention_mask, |
| 1022 | segment_ids=segment_ids, |
| 1023 | position_ids=position_ids, |
| 1024 | deterministic=deterministic, |
| 1025 | init_cache=init_cache, |
| 1026 | output_attentions=output_attentions, |
| 1027 | output_hidden_states=output_hidden_states, |
| 1028 | return_dict=return_dict, |
| 1029 | ) |
| 1030 | |
| 1031 | hidden_states = outputs[0] |
| 1032 | hidden_states = self.ln_f(hidden_states) |
| 1033 | |
| 1034 | if output_hidden_states: |
| 1035 | all_hidden_states = outputs[1] + (hidden_states,) |
| 1036 | outputs = (hidden_states, all_hidden_states) + outputs[2:] |
| 1037 | else: |
| 1038 | outputs = (hidden_states,) + outputs[1:] |
| 1039 | |