Forward pass to quantize decode params.
(self, state, rng: PRNGKeyType | None = None)
| 319 | lora_utils.unapply_lora_from_base_params(base_params, adapter_params, lora_scale_factor) |
| 320 | |
| 321 | def quantize_params(self, state, rng: PRNGKeyType | None = None): |
| 322 | """Forward pass to quantize decode params.""" |
| 323 | if rng is None: |
| 324 | rng = jax.random.PRNGKey(0) |
| 325 | |
| 326 | self.model.quant.quant_mode = quantizations.get_quant_mode("convert") |
| 327 | |
| 328 | @jax.jit |
| 329 | def model_apply(_p, _rng): |
| 330 | image_shape = multimodal_utils.get_dummy_image_shape_for_init( |
| 331 | self.config.model_name, batch_size=self.config.micro_batch_size_to_train_on |
| 332 | ) |
| 333 | return self.model.apply( |
| 334 | _p | {"aqt": {}}, |
| 335 | jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32), |
| 336 | jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32), |
| 337 | encoder_images=jnp.ones(image_shape, dtype=jnp.float32) if self.config.use_multimodal else None, |
| 338 | # encoder_image_masks indicates valid tiles if image tiling + padding is used in vision encoder input. |
| 339 | encoder_image_masks=jnp.ones(image_shape[:2], dtype=jnp.int32) |
| 340 | if self.config.use_multimodal and "llama4" in self.config.model_name |
| 341 | else None, |
| 342 | decoder_segment_ids=jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32), |
| 343 | enable_dropout=False, |
| 344 | model_mode=MODEL_MODE_PREFILL, |
| 345 | rngs={"params": _rng}, |
| 346 | mutable=True, |
| 347 | ) |
| 348 | |
| 349 | _, new_vars = model_apply(state.params, rng) |
| 350 | # Remove param values which have corresponding qtensors in aqt to save memory. |
| 351 | params = {} |
| 352 | params["aqt"] = new_vars["aqt"] |
| 353 | params["params"] = quantizations.remove_quantized_params(state.params["params"], new_vars["aqt"]) |
| 354 | self.abstract_params = jax.tree_util.tree_map( |
| 355 | lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding), |
| 356 | params, |
| 357 | ) |
| 358 | maxtext_utils.save_quantized_checkpoint_if_configured(self.config, params) |
| 359 | self.model.quant.quant_mode = quantizations.get_quant_mode("serve") |
| 360 | return params |
| 361 | |
| 362 | def _maybe_stack_prefill_result_cache(self, cache): |
| 363 | """Stack the caches across the layers.""" |