diff --git a/tpu_inference/envs.py b/tpu_inference/envs.py index 5cf956f7de..5968ce7334 100644 --- a/tpu_inference/envs.py +++ b/tpu_inference/envs.py @@ -83,6 +83,8 @@ VLLM_TPU_BUCKET_PADDING_GAP: int = 0 VLLM_INCREMENTAL_FP8_LOADING: bool = False TPU_MESH_SORT_BY_COORDS: bool = False + MOE_REQUANTIZE_WEIGHT_DTYPE: str = "" + MOE_REQUANTIZE_BLOCK_SIZE: str = "" def env_with_choices( @@ -494,6 +496,12 @@ def _get_int_list_env() -> list[int]: # when initializing large FP8 models on smaller RAM TPUs such as TPU8i. "VLLM_INCREMENTAL_FP8_LOADING": env_bool("VLLM_INCREMENTAL_FP8_LOADING", default=False), + # Controls the target quantization dtype for MoE requantization (e.g. 'fp8', 'int4'). + "MOE_REQUANTIZE_WEIGHT_DTYPE": + lambda: os.getenv("MOE_REQUANTIZE_WEIGHT_DTYPE", ""), + # Controls the block size for MoE requantization (e.g. '32', '64'). + "MOE_REQUANTIZE_BLOCK_SIZE": + lambda: os.getenv("MOE_REQUANTIZE_BLOCK_SIZE", ""), } diff --git a/tpu_inference/layers/vllm/quantization/mxfp4.py b/tpu_inference/layers/vllm/quantization/mxfp4.py index 308e1d2bd9..fbf510bda2 100644 --- a/tpu_inference/layers/vllm/quantization/mxfp4.py +++ b/tpu_inference/layers/vllm/quantization/mxfp4.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools from typing import Optional import jax @@ -38,6 +39,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import \ is_layer_skipped +from tpu_inference import envs from tpu_inference.layers.common.moe import \ FusedMoEMethodBase as TpuFusedMoEMethodBase from tpu_inference.layers.common.process_weights.moe_weights import ( @@ -53,7 +55,7 @@ from tpu_inference.layers.vllm.quantization.unquantized import \ VllmUnquantizedLinearMethod from tpu_inference.logger import init_logger -from tpu_inference.utils import get_mesh_shape_product, t2j +from tpu_inference.utils import get_mesh_shape_product, t2j, to_jax_dtype P = PartitionSpec @@ -91,6 +93,56 @@ def get_quant_method(self, layer: torch.nn.Module, return None +@functools.partial( + jax.jit, + static_argnames=( + "moe_backend", + "w13_reorder_size", + "w13_interleave", + "desired_quant_dtype", + "requant_block_size", + ), +) +def _process_mxfp4_moe_weights( + w13_weight: jax.Array, + w13_weight_scale: jax.Array, + w13_bias: jax.Array | None, + w2_weight: jax.Array, + w2_weight_scale: jax.Array, + w2_bias: jax.Array | None, + moe_backend: Any, + w13_reorder_size: int, + w13_interleave: bool, + desired_quant_dtype: Any, + requant_block_size: int, +) -> FusedMoEWeights: + # Dequantize fp4 weights into fp32. + w13_weight = dequantize_tensor_from_mxfp4_packed( + w13_weight, w13_weight_scale, 2, jnp.float32) + w2_weight = dequantize_tensor_from_mxfp4_packed( + w2_weight, w2_weight_scale, 2, jnp.float32) + + weights = quantize_moe_weights( + FusedMoEWeights( + w13_weight=w13_weight, + w13_weight_scale=None, + w13_bias=w13_bias, + w2_weight=w2_weight, + w2_weight_scale=None, + w2_bias=w2_bias, + ), + desired_quant_dtype, + requant_block_size, + w13_interleave=w13_interleave, + ) + return process_moe_weights( + weights, + moe_backend=moe_backend, + w13_reorder_size=w13_reorder_size, + w13_interleave=w13_interleave, + ) + + class VllmMxfp4MoEMethod(Mxfp4MoEMethod, FusedMoEMethodBase): def __init__( @@ -141,51 +193,25 @@ def process_weights_after_loading(self, layer: torch.nn.Module): w2_weight_scale = t2j(layer.w2_weight_scale, use_dlpack=False) w2_bias = t2j(layer.w2_bias, use_dlpack=False) if has_bias else None - @jax.jit - def process_mxfp4_moe_weights( - w13_weight: jax.Array, - w13_weight_scale: jax.Array, - w13_bias: jax.Array | None, - w2_weight: jax.Array, - w2_weight_scale: jax.Array, - w2_bias: jax.Array | None, - ) -> FusedMoEWeights: - # Dequantize fp4 weights into fp32. - w13_weight = dequantize_tensor_from_mxfp4_packed( - w13_weight, w13_weight_scale, 2, jnp.float32) - w2_weight = dequantize_tensor_from_mxfp4_packed( - w2_weight, w2_weight_scale, 2, jnp.float32) - w13_interleave = layer.activation == MoEActivation.SWIGLUOAI - w13_reorder_size = get_mesh_shape_product( - self.mesh, ShardingAxisName.MLP_TENSOR) - - weights = quantize_moe_weights( - FusedMoEWeights( - w13_weight=w13_weight, - w13_weight_scale=None, - w13_bias=w13_bias, - w2_weight=w2_weight, - w2_weight_scale=None, - w2_bias=w2_bias, - ), - jnp.float4_e2m1fn, - MXFP4_REQUANTIZED_BLOCK_SIZE, - w13_interleave=w13_interleave, - ) - return process_moe_weights( - weights, - moe_backend=self.moe_backend, - w13_reorder_size=w13_reorder_size, - w13_interleave=w13_interleave, - ) - - weights = process_mxfp4_moe_weights( + w13_interleave = layer.activation == MoEActivation.SWIGLUOAI + w13_reorder_size = get_mesh_shape_product( + self.mesh, ShardingAxisName.MLP_TENSOR) + + desired_quant_dtype = to_jax_dtype(envs.MOE_REQUANTIZE_WEIGHT_DTYPE) if envs.MOE_REQUANTIZE_WEIGHT_DTYPE else jnp.float4_e2m1fn + requant_block_size = int(envs.MOE_REQUANTIZE_BLOCK_SIZE) if envs.MOE_REQUANTIZE_BLOCK_SIZE else MXFP4_REQUANTIZED_BLOCK_SIZE + + weights = _process_mxfp4_moe_weights( w13_weight, w13_weight_scale, w13_bias, w2_weight, w2_weight_scale, w2_bias, + moe_backend=self.moe_backend, + w13_reorder_size=w13_reorder_size, + w13_interleave=w13_interleave, + desired_quant_dtype=desired_quant_dtype, + requant_block_size=requant_block_size, ) weights = torch_view( shard_moe_weights(weights, self.moe_backend, self.mesh))