Layerwise quantization for large models.
| 55 | |
| 56 | |
| 57 | class LayerwiseQuantization: |
| 58 | """ |
| 59 | Layerwise quantization for large models. |
| 60 | """ |
| 61 | |
| 62 | def __init__(self, config: Any): |
| 63 | self.config = config |
| 64 | |
| 65 | # TODO(ranlihao): Remove this assertion once the Layerwise quantization is supported for other decoder blocks. |
| 66 | assert ( |
| 67 | config.decoder_block == common_types.DecoderBlockType.DEEPSEEK |
| 68 | ), f"Layerwise quantization is only supported for {common_types.DecoderBlockType.DEEPSEEK}\ |
| 69 | , but got {config.decoder_block}." |
| 70 | |
| 71 | # Mesh definition |
| 72 | devices_array = maxtext_utils.create_device_mesh(config=config) |
| 73 | self._mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) |
| 74 | |
| 75 | # Model and quantization config |
| 76 | self.quant = quantizations.configure_quantization(config) |
| 77 | model = models.transformer_as_linen( |
| 78 | config, mesh=self._mesh, quant=self.quant, model_mode=common_types.MODEL_MODE_TRAIN |
| 79 | ) |
| 80 | rng = jax.random.PRNGKey(1234) |
| 81 | self.unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(model, None, self.config, rng, self._mesh, False) |
| 82 | |
| 83 | def load_and_quantize(self, rng: None | PRNGKeyType = None) -> None: |
| 84 | """ |
| 85 | Load parameters layer by layer and quantize them. |
| 86 | """ |
| 87 | |
| 88 | quantized_params = {} |
| 89 | quantized_params["params"] = {"decoder": {}} |
| 90 | quantized_params["aqt"] = {"decoder": {}} |
| 91 | |
| 92 | config = self.config |
| 93 | |
| 94 | self.quant.quant_mode = quantizations.get_quant_mode("convert") |
| 95 | |
| 96 | layers = [ |
| 97 | deepseek.DeepSeekDenseLayer(config, mesh=self._mesh, quant=self.quant), |
| 98 | deepseek.DeepSeekMoELayer(config, mesh=self._mesh, quant=self.quant), |
| 99 | ] |
| 100 | layer_prefixes = ["dense_layers", "moe_layers"] |
| 101 | num_moe_layers = config.num_decoder_layers - config.first_num_dense_layers |
| 102 | num_layers_list = [config.first_num_dense_layers, num_moe_layers] |
| 103 | |
| 104 | def model_apply(_p, _rng, layer): |
| 105 | return layer.apply( |
| 106 | _p | {"aqt": {}}, |
| 107 | jnp.ones((1, self.config.max_prefill_predict_length, self.config.base_emb_dim), dtype=jnp.int32), |
| 108 | None, |
| 109 | jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32), |
| 110 | True, |
| 111 | model_mode=common_types.MODEL_MODE_PREFILL, |
| 112 | rngs={"params": _rng}, |
| 113 | mutable=True, |
| 114 | ) |