tpu-inference 0.11.1rc1__py3-none-any.whl → 0.11.1rc3__py3-none-any.whl

tpu_inference/layers/vllm/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py

@@ -0,0 +1,136 @@
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+import torch
+from compressed_tensors.quantization import QuantizationStrategy
+from jax.sharding import NamedSharding, PartitionSpec
+from torchax.interop import jax_view, torch_view
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_int8 import \
+    CompressedTensorsW8A8Int8
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import \
+    convert_to_channelwise
+
+from tpu_inference.layers.vllm.linear_common import (
+    sharded_quantized_matmul, slice_sharded_tensor_for_concatenation,
+    torch_to_jax_param)
+from tpu_inference.layers.vllm.quantization.common import JaxCommonLinearConfig
+
+P = PartitionSpec
+logger = init_logger(__name__)
+
+
+class VllmCompressedTensorsW8A8Int8(CompressedTensorsW8A8Int8):
+
+    def __init__(self, strategy: str, is_static_input_scheme: bool,
+                 input_symmetric: bool, jax_config: JaxCommonLinearConfig):
+        super().__init__(strategy, is_static_input_scheme, input_symmetric)
+
+        self.jax_config = jax_config
+        self.is_channelwise = (self.strategy == QuantizationStrategy.CHANNEL),
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = torch_to_jax_param(
+            layer.weight,
+            NamedSharding(self.jax_config.mesh,
+                          self.jax_config.weight_sharding),
+            self.jax_config.output_sizes,
+            self.jax_config.n_shards,
+            self.jax_config.fuse_matmuls,
+        )
+        delattr(layer, "weight")
+        layer.weight = weight
+
+        weight_scale = layer.weight_scale
+        is_fused_module = len(layer.logical_widths) > 1
+        if is_fused_module and not self.is_channelwise:
+            weight_scale = convert_to_channelwise(weight_scale,
+                                                  layer.logical_widths)
+        weight_scale = weight_scale.squeeze(-1)
+
+        weight_scale = torch_to_jax_param(
+            weight_scale,
+            NamedSharding(self.jax_config.mesh, self.jax_config.bias_sharding),
+            self.jax_config.output_sizes,
+            self.jax_config.n_shards,
+            self.jax_config.fuse_matmuls,
+        )
+        delattr(layer, "weight_scale")
+        layer.weight_scale = weight_scale
+
+        if layer.bias is not None and not layer.skip_bias_add:
+            if layer.return_bias:
+                logger.warning_once("Bias might return incorrect value.")
+
+            bias = torch_to_jax_param(
+                layer.bias,
+                NamedSharding(self.jax_config.mesh,
+                              self.jax_config.bias_sharding),
+                self.jax_config.output_sizes,
+                self.jax_config.n_shards,
+                self.jax_config.fuse_matmuls,
+            )
+            delattr(layer, "bias")
+            layer.bias = bias
+
+        # TODO(kyuyeunk): Support static range input quantization.
+        assert getattr(layer, "input_scale", None) is None
+        assert getattr(layer, "input_zero_point", None) is None
+        assert getattr(layer, "azp_adj", None) is None
+
+    def apply_weights(self, layer: torch.nn.Module, x: torch.Tensor,
+                      bias: Optional[torch.Tensor]) -> torch.Tensor:
+        with jax.named_scope(layer._get_name()):
+            if self.jax_config.fuse_matmuls:
+                out = self._apply_fused(layer, x, bias)
+            else:
+                out = self._apply_split(layer, x, bias)
+
+        return out
+
+    def _apply_fused(self, layer: torch.nn.Module, x: torch.Tensor,
+                     bias: Optional[torch.Tensor]) -> torch.Tensor:
+        x_jax = jax_view(x)
+        weight_jax = jax_view(layer.weight)
+        weight_scale_jax = jax_view(layer.weight_scale)
+
+        outs = sharded_quantized_matmul(
+            x_jax,
+            weight_jax,
+            weight_scale_jax,
+            self.jax_config.mesh,
+            self.jax_config.weight_sharding,
+        )
+        if bias is not None and not layer.skip_bias_add:
+            outs += jax_view(bias)
+
+        outs = slice_sharded_tensor_for_concatenation(
+            outs, self.jax_config.output_sizes, self.jax_config.n_shards)
+        out = jnp.concatenate(outs, axis=-1)
+        return torch_view(out)
+
+    def _apply_split(self, layer: torch.nn.Module, x: torch.Tensor,
+                     bias: Optional[torch.Tensor]) -> torch.Tensor:
+        assert isinstance(layer.weight, torch.nn.ParameterList)
+
+        x_jax = jax_view(x)
+        outs = []
+        for i, (weight, weight_scale) in enumerate(
+                zip(layer.weight, layer.weight_scale)):
+            weight_jax = jax_view(weight)
+            weight_scale_jax = jax_view(weight_scale)
+
+            out = sharded_quantized_matmul(
+                x_jax,
+                weight_jax,
+                weight_scale_jax,
+                self.jax_config.mesh,
+                self.jax_config.weight_sharding,
+            )
+            if bias is not None and not layer.skip_bias_add:
+                out += jax_view(bias[i])
+
+            outs.append(out)
+        out = jnp.concatenate(outs, axis=-1)
+        return torch_view(out)

tpu_inference/layers/vllm/quantization/unquantized.py

@@ -0,0 +1,263 @@
+import functools
+from typing import Any, Callable, Optional, Union
+
+import jax
+import jax.numpy as jnp
+import torch
+from jax.experimental.layout import Format, Layout
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from torch.nn.parameter import Parameter
+from torchax.interop import jax_view, torch_view
+from torchax.ops.mappings import t2j
+from vllm.attention.layer import Attention
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.layer import (
+    FusedMoE, FusedMoEConfig, UnquantizedFusedMoEMethod)
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
+    FusedMoEPermuteExpertsUnpermute, FusedMoEPrepareAndFinalize)
+from vllm.model_executor.layers.linear import (LinearBase,
+                                               UnquantizedLinearMethod)
+from vllm.model_executor.layers.quantization import \
+    register_quantization_config
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig, QuantizeMethodBase)
+
+from tpu_inference.layers.vllm.fused_moe import jax_fused_moe_func_padded
+from tpu_inference.layers.vllm.linear_common import (
+    reorder_concatenated_tensor_for_sharding,
+    slice_sharded_tensor_for_concatenation, torch_to_jax_param)
+from tpu_inference.layers.vllm.quantization.common import (
+    JaxCommonConfig, JaxCommonLinearConfig)
+
+P = PartitionSpec
+logger = init_logger(__name__)
+
+
+@register_quantization_config("jax-unquantized")
+class VllmUnquantizedConfig(QuantizationConfig, JaxCommonConfig):
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "jax-unquantized"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.float32, torch.float16, torch.bfloat16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 0  # Always supported
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []  # No extra configs required.
+
+    @classmethod
+    def from_config(cls, _: dict[str, Any]) -> "VllmUnquantizedConfig":
+        return cls()
+
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional[QuantizeMethodBase]:
+        if isinstance(layer, LinearBase):
+            linear_config = self.get_linear_config(layer)
+            return VllmUnquantizedLinearMethod(linear_config)
+        if isinstance(layer, FusedMoE):
+            moe_config = self.get_moe_config(layer)
+            return VllmUnquantizedFusedMoEMethod(moe_config, self.mesh)
+        if isinstance(layer, Attention):
+            return None
+        return None
+
+
+class VllmUnquantizedLinearMethod(UnquantizedLinearMethod):
+
+    def __init__(self, jax_config: JaxCommonLinearConfig):
+        self.jax_config = jax_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = torch_to_jax_param(
+            layer.weight,
+            NamedSharding(self.jax_config.mesh,
+                          self.jax_config.weight_sharding),
+            self.jax_config.output_sizes,
+            self.jax_config.n_shards,
+            self.jax_config.fuse_matmuls,
+        )
+        delattr(layer, "weight")
+        layer.weight = weight
+
+        if layer.bias is not None and not layer.skip_bias_add:
+            if layer.return_bias:
+                logger.warning_once("Bias might return incorrect value.")
+
+            bias = torch_to_jax_param(
+                layer.bias,
+                NamedSharding(self.jax_config.mesh,
+                              self.jax_config.bias_sharding),
+                self.jax_config.output_sizes,
+                self.jax_config.n_shards,
+                self.jax_config.fuse_matmuls,
+            )
+            delattr(layer, "bias")
+            layer.bias = bias
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+        with jax.named_scope(layer._get_name()):
+            if in_sharding := self.jax_config.get_input_sharding(x):
+                x.shard_(NamedSharding(self.jax_config.mesh, in_sharding))
+
+            if self.jax_config.fuse_matmuls:
+                out = self._apply_fused(layer, x, bias)
+            else:
+                out = self._apply_split(layer, x, bias)
+
+            if out_sharding := self.jax_config.get_output_sharding(out):
+                out.shard_(NamedSharding(self.jax_config.mesh, out_sharding))
+
+        return out
+
+    def _apply_fused(self,
+                     layer: torch.nn.Module,
+                     x: torch.Tensor,
+                     bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+        x_jax = jax_view(x)
+        weight_jax = jax_view(layer.weight)
+
+        outs = jnp.einsum("mn,pn->mp", x_jax, weight_jax)
+        if bias is not None and not layer.skip_bias_add:
+            outs += bias.jax()
+
+        outs = slice_sharded_tensor_for_concatenation(
+            outs, self.jax_config.output_sizes, self.jax_config.n_shards)
+        out = jnp.concatenate(outs, axis=-1)
+        return torch_view(out)
+
+    def _apply_split(self,
+                     layer: torch.nn.Module,
+                     x: torch.Tensor,
+                     bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+        assert isinstance(layer.weight, torch.nn.ParameterList)
+
+        x_jax = x.jax()
+        outs = []
+        for i, weight in enumerate(layer.weight):
+            weight_jax = jax_view(weight)
+
+            out = jnp.einsum("mn,pn->mp", x_jax, weight_jax)
+            if bias is not None and not layer.skip_bias_add:
+                out += jax_view(bias[i])
+
+            outs.append(out)
+        out = jnp.concatenate(outs, axis=-1)
+        return torch_view(out)
+
+
+class VllmUnquantizedFusedMoEMethod(UnquantizedFusedMoEMethod):
+
+    def __init__(self, moe: FusedMoEConfig, mesh: Mesh):
+        super().__init__(moe)
+        self.mesh = mesh
+
+    def select_gemm_impl(
+        self,
+        prepare_finalize: FusedMoEPrepareAndFinalize,
+        moe: FusedMoEConfig,
+        layer: torch.nn.Module,
+    ) -> FusedMoEPermuteExpertsUnpermute:
+        raise NotImplementedError(
+            "Selecting gemm implementation is currently not supported.")
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        assert isinstance(layer, FusedMoE)
+
+        w2_weight = t2j(layer.w2_weight, use_dlpack=False)
+        w13_weight = t2j(layer.w13_weight, use_dlpack=False)
+
+        if layer.use_ep:
+            w13_weight = jax.device_put(
+                w13_weight,
+                Format(Layout((0, 1, 2)),
+                       NamedSharding(self.mesh, P("model", None, None))))
+            w2_weight = jax.device_put(
+                w2_weight,
+                Format(Layout((0, 1, 2)),
+                       NamedSharding(self.mesh, P("model", None, None))))
+        else:
+            intermediate_size = w13_weight.shape[1] // 2
+            assert intermediate_size == w2_weight.shape[-1]
+            output_sizes = [intermediate_size, intermediate_size]
+            n_shards = self.mesh.shape["model"]
+            assert intermediate_size % n_shards == 0
+            w13_weight = reorder_concatenated_tensor_for_sharding(w13_weight,
+                                                                  output_sizes,
+                                                                  n_shards,
+                                                                  dim=1)
+            w13_weight = jax.device_put(
+                w13_weight,
+                Format(Layout((0, 1, 2)),
+                       NamedSharding(self.mesh, P(None, "model", None))))
+            w2_weight = jax.device_put(
+                w2_weight,
+                Format(Layout((0, 1, 2)),
+                       NamedSharding(self.mesh, P(None, None, "model"))))
+        w13_weight = Parameter(torch_view(w13_weight), requires_grad=False)
+        w2_weight = Parameter(torch_view(w2_weight), requires_grad=False)
+
+        layer.w13_weight = w13_weight
+        layer.w2_weight = w2_weight
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+        top_k: int,
+        renormalize: bool,
+        use_grouped_topk: bool = False,
+        topk_group: Optional[int] = None,
+        num_expert_group: Optional[int] = None,
+        global_num_experts: int = -1,
+        expert_map: Optional[torch.Tensor] = None,
+        custom_routing_function: Optional[Callable] = None,
+        scoring_func: str = "softmax",
+        routed_scaling_factor: float = 1.0,
+        e_score_correction_bias: Optional[torch.Tensor] = None,
+        apply_router_weight_on_input: bool = False,
+        activation: str = "silu",
+        enable_eplb: bool = False,
+        expert_load_view: Optional[torch.Tensor] = None,
+        logical_to_physical_map: Optional[torch.Tensor] = None,
+        logical_replica_count: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        assert isinstance(layer, FusedMoE)
+        if activation != "silu":
+            raise NotImplementedError(
+                "Only silu is supported for activation function.")
+        if scoring_func != "softmax":
+            raise NotImplementedError(
+                "Only softmax is supported for scoring_func")
+
+        _fused_moe_func = functools.partial(
+            jax.jit(jax_fused_moe_func_padded,
+                    static_argnames=[
+                        "topk", "global_num_experts", "renormalize",
+                        "reduce_results", "mesh", "use_ep"
+                    ]),
+            topk=top_k,
+            global_num_experts=global_num_experts,
+            renormalize=renormalize,
+            reduce_results=layer.reduce_results,
+            mesh=self.mesh,
+            use_ep=layer.use_ep)
+
+        output = _fused_moe_func(
+            jax_view(x),
+            jax_view(layer.w13_weight),
+            jax_view(layer.w2_weight),
+            jax_view(router_logits),
+        )
+
+        return torch_view(output)

tpu_inference/layers/vllm/sharding.py

@@ -0,0 +1,151 @@
+import jax
+import torch
+import torchax
+from jax.sharding import Mesh, NamedSharding, PartitionSpec
+from torch.nn import Parameter
+from torch.utils import _pytree as pytree
+from torchax.interop import jax_view, torch_view
+from torchax.ops.mappings import t2j
+from vllm.lora.layers import (MergedColumnParallelLinearWithLoRA,
+                              MergedQKVParallelLinearWithLoRA,
+                              RowParallelLinearWithLoRA)
+from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead, VocabParallelEmbedding)
+
+from tpu_inference.logger import init_logger
+
+P = PartitionSpec
+
+logger = init_logger(__name__)
+
+
+def shard_model_to_tpu(model: torch.nn.Module,
+                       mesh: Mesh) -> dict[str, torchax.torch.Tensor]:
+    """
+    Shard the model weights and move them to TPU.
+
+    At the same time, also turn the weight tensors into torchax tensors so that
+    jax code can interop with it and the overall program can be traced and
+    compiled in XLA.
+
+    Args:
+        model: A PyTorch model whose weights are on CPU main memory.
+        mesh: JAX mesh object for sharding.
+    Returns:
+        Dictionary of parameters and buffers that will be used as arguments of
+        torch.func.functional_call
+    """
+
+    with jax.default_device(jax.devices("cpu")[0]):
+        _shard_module_to_tpu(model, mesh)
+
+        params, buffers = _extract_all_params_buffers(model)
+
+        # For other weight tensors, repliate them on all the TPU chips.
+        params, buffers = pytree.tree_map_only(
+            _tensor_is_in_cpu,
+            lambda tensor: _shard_tensor_to_tpu_replicated(tensor, mesh),
+            (params, buffers))
+
+        return {**params, **buffers}
+
+
+def _extract_all_params_buffers(model: torch.nn.Module):
+    return dict(model.named_parameters()), dict(model.named_buffers())
+
+
+def _tensor_is_in_cpu(tensor: torch.tensor) -> bool:
+    # Check if a tensor haven't been converted to torchax tensor.
+    if not isinstance(tensor, torchax.tensor.Tensor):
+        return True
+    # Check if torchax tensor is still in CPU.
+    return tensor.jax_device == jax.devices('cpu')[0]
+
+
+def _convert_to_torchax_and_shard(tensor: torch.Tensor,
+                                  sharding: NamedSharding) -> torch.Tensor:
+    if isinstance(tensor, torchax.tensor.Tensor):
+        tensor = jax_view(tensor)
+    else:
+        tensor = t2j(tensor)
+    return torch_view(jax.device_put(tensor, sharding))
+
+
+def _shard_tensor_to_tpu_replicated(tensor: torch.Tensor,
+                                    mesh: Mesh) -> torchax.tensor.Tensor:
+    return _convert_to_torchax_and_shard(tensor, NamedSharding(mesh, P()))
+
+
+def _shard_vocab_parallel_embedding(layer: VocabParallelEmbedding,
+                                    mesh: Mesh) -> None:
+    weight = _convert_to_torchax_and_shard(
+        layer.weight, NamedSharding(mesh, P('model', None)))
+    layer.weight = Parameter(weight, requires_grad=False)
+
+
+def _shard_lm_head(layer: ParallelLMHead, mesh: Mesh):
+    # TODO(qihqi): currently this is not handling case of tie_word_weights=True.
+    # if that config is set, then we should not create new weights but reuse the
+    # weight from VocabParallelEmbedding
+    weight = _convert_to_torchax_and_shard(
+        layer.weight, NamedSharding(mesh, P('model', None)))
+    layer.weight = Parameter(weight, requires_grad=False)
+    if layer.bias is not None:
+        bias = _convert_to_torchax_and_shard(layer.bias,
+                                             NamedSharding(mesh, P('model')))
+        layer.bias = Parameter(bias, requires_grad=False)
+
+
+def _shard_base_linear_lora(layer: BaseLinearLayerWithLoRA,
+                            mesh: Mesh) -> None:
+    # NOTE: lora_a_stacked[i] has shape [max_loras, 1, num_out, num_in]
+    sharded_lora_a_tpu = torch.nn.ParameterList()
+    sharded_lora_b_tpu = torch.nn.ParameterList()
+
+    for i in range(layer.n_slices):
+        sharded_lora_a_tpu.append(
+            _shard_tensor_to_tpu_replicated(layer.lora_a_stacked[i], mesh))
+        sharded_lora_b_tpu.append(
+            _shard_tensor_to_tpu_replicated(layer.lora_b_stacked[i], mesh))
+
+    layer.lora_a_stacked = sharded_lora_a_tpu
+    layer.lora_b_stacked = sharded_lora_b_tpu
+
+
+# TODO: Add custom sharding logic for following lora layers
+def _shard_column_parallel_linear_lora(
+        layer: MergedColumnParallelLinearWithLoRA, mesh: Mesh) -> None:
+    _shard_base_linear_lora(layer, mesh)
+
+
+def _shard_qkv_parallel_linear_lora(layer: MergedQKVParallelLinearWithLoRA,
+                                    mesh: Mesh) -> None:
+    _shard_base_linear_lora(layer, mesh)
+
+
+def _shard_row_parallel_linear_lora(layer: RowParallelLinearWithLoRA,
+                                    mesh: Mesh) -> None:
+    _shard_base_linear_lora(layer, mesh)
+
+
+# NOTE: Ordering is important as it calls first matched type of a given module
+MODULE_TYPE_TO_SHARDING_FUNC = [
+    # Shard embedding layers
+    (ParallelLMHead, _shard_lm_head),
+    (VocabParallelEmbedding, _shard_vocab_parallel_embedding),
+    # Shard LoRA layers
+    (MergedColumnParallelLinearWithLoRA, _shard_column_parallel_linear_lora),
+    (MergedQKVParallelLinearWithLoRA, _shard_qkv_parallel_linear_lora),
+    (RowParallelLinearWithLoRA, _shard_row_parallel_linear_lora),
+    (BaseLinearLayerWithLoRA, _shard_base_linear_lora),
+]
+
+
+def _shard_module_to_tpu(model: torch.nn.Module, mesh: Mesh) -> None:
+    for path, module in model.named_modules():
+        for module_type, sharding_func in MODULE_TYPE_TO_SHARDING_FUNC:
+            if isinstance(module, module_type):
+                logger.debug("shard %s with %s", path, sharding_func)
+                sharding_func(module, mesh)
+                break