Load Parameters from GCS or reshard given Parameters
(self, *args, params=None, rng: PRNGKeyType | None = None, **kwargs)
| 217 | ) |
| 218 | |
| 219 | def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwargs) -> Params: |
| 220 | """Load Parameters from GCS or reshard given Parameters""" |
| 221 | # pylint: disable=unused-argument |
| 222 | |
| 223 | if rng is None: |
| 224 | rng = jax.random.PRNGKey(0) |
| 225 | |
| 226 | if self.model.quant and self.config.checkpoint_is_quantized: |
| 227 | print("Loading from the quantized checkpoint...") |
| 228 | self.model.quant.quant_mode = quantizations.get_quant_mode("serve") |
| 229 | |
| 230 | rng1, rng2, rng3 = jax.random.split(rng, 3) |
| 231 | if params: |
| 232 | print("Resharding given params") |
| 233 | _, self.state_mesh_annotations, state_mesh_shardings = maxtext_utils.get_abstract_state( |
| 234 | self.model, None, self.config, rng, self._mesh, False |
| 235 | ) |
| 236 | # reshard given params based on shardings from config in MaxEngine |
| 237 | params = jax.device_put(params, state_mesh_shardings.params) |
| 238 | state = maxtext_utils.init_decode_state(None, params) |
| 239 | state = max_utils.unbox_logicallypartioned(state) |
| 240 | else: |
| 241 | state, self.state_mesh_annotations = maxtext_utils.setup_decode_state( |
| 242 | self.model, self.config, rng1, self._mesh, None |
| 243 | ) |
| 244 | # pylint: disable=isinstance-second-argument-not-valid-type |
| 245 | self.abstract_params = jax.tree_util.tree_map( |
| 246 | lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding) |
| 247 | if isinstance(x, jax.Array) |
| 248 | else None, |
| 249 | state.params, |
| 250 | ) |
| 251 | |
| 252 | self.prefill_kv_cache_annotations = maxtext_utils.get_prefill_kv_cache_annotations( |
| 253 | self.model, self.config, rng2, self._mesh, self.page_state |
| 254 | ) |
| 255 | self.prefill_kv_cache_shardings = jax.tree_util.tree_map( |
| 256 | lambda x: jax.sharding.NamedSharding(self._mesh, x), |
| 257 | self.prefill_kv_cache_annotations, |
| 258 | ) |
| 259 | |
| 260 | if self.config.stack_prefill_result_cache: |
| 261 | # Add extra axis for the axis generated by the stack. |
| 262 | self.prefill_kv_cache_shardings = jax.tree_util.tree_map( |
| 263 | lambda x: jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec(None, *x.spec)), |
| 264 | self.prefill_kv_cache_shardings, |
| 265 | ) |
| 266 | self.prefill_kv_cache_shardings = self.prefill_kv_cache_shardings["decoder"]["layers_0"] |
| 267 | |
| 268 | self.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations( |
| 269 | self.model, self.config, rng2, self._mesh, self.page_state |
| 270 | ) |
| 271 | self.kv_cache_shardings = jax.tree_util.tree_map( |
| 272 | lambda x: jax.sharding.NamedSharding(self._mesh, x), |
| 273 | self.kv_cache_annotations, |
| 274 | ) |
| 275 | |
| 276 | if self.model.quant and not self.config.checkpoint_is_quantized: |