tpu-inference 0.0.1rc1__py3-none-any.whl → 0.11.1.dev202511130813__py3-none-any.whl

tpu_inference/layers/vllm/quantization/mxfp4.py

@@ -1,331 +0,0 @@
-from typing import 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.logger import init_logger
-from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEConfig, FusedMoEQuantConfig, biased_moe_quant_config)
-from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
-                                                        FusedMoEMethodBase)
-from vllm.model_executor.layers.linear import LinearBase
-from vllm.model_executor.layers.quantization import \
-    register_quantization_config
-from vllm.model_executor.layers.quantization.base_config import \
-    QuantizeMethodBase
-from vllm.model_executor.layers.quantization.mxfp4 import (Mxfp4Backend,
-                                                           Mxfp4Config,
-                                                           Mxfp4MoEMethod)
-from vllm.model_executor.layers.quantization.utils.quant_utils import \
-    is_layer_skipped
-
-from tpu_inference import envs
-from tpu_inference.kernels.fused_moe.v1.kernel import fused_ep_moe
-from tpu_inference.layers.common.quant_methods import (MXFP4,
-                                                       get_tpu_quant_method)
-from tpu_inference.layers.vllm.fused_moe import fused_moe_func
-from tpu_inference.layers.vllm.linear_common import \
-    reorder_concatenated_tensor_for_sharding
-from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
-from tpu_inference.layers.vllm.quantization.unquantized import \
-    VllmUnquantizedLinearMethod
-
-MXFP4_BLOCK_SIZE = 32
-
-P = PartitionSpec
-logger = init_logger(__name__)
-
-
-# TODO(kyuyeunk): Move these functions into a common utility file.
-def u8_unpack_e2m1(u8_packed_e2m1: jax.Array) -> jax.Array:
-    assert u8_packed_e2m1.dtype == jnp.uint8
-    e2m1 = jax.lax.bitcast_convert_type(u8_packed_e2m1, jnp.float4_e2m1fn)
-    # bitcast creates one more dimension that splits 8 bits into two e2m1.
-    # we flatten them with the last dim.
-    return jnp.reshape(e2m1, e2m1.shape[:-2] + (-1, ))
-
-
-def e8m0_to_fp32(u8: jax.Array) -> jax.Array:
-    e8_finfo = jnp.finfo(jnp.float8_e8m0fnu)
-    exponents = u8.astype(jnp.int32) + e8_finfo.minexp
-    ones = jnp.ones_like(u8, dtype=jnp.float32)
-    return jnp.ldexp(ones, exponents)
-
-
-def dequantize_block_weight(weight: jax.Array,
-                            scale: jax.Array,
-                            block_size: int,
-                            out_dtype: jnp.dtype = jnp.bfloat16) -> jax.Array:
-    orig_shape = weight.shape
-    weight_block = weight.reshape(orig_shape[:-1] + (-1, block_size))
-    weight_dequantized = weight_block.astype(jnp.float32) * jnp.expand_dims(
-        scale, -1)
-    return weight_dequantized.reshape(orig_shape).astype(out_dtype)
-
-
-@register_quantization_config(get_tpu_quant_method(MXFP4))
-class VllmMxfp4Config(Mxfp4Config, JaxCommonConfig):
-
-    @classmethod
-    def get_name(cls):
-        return MXFP4
-
-    def get_quant_method(self, layer: torch.nn.Module,
-                         prefix: str) -> Optional["QuantizeMethodBase"]:
-        from vllm.attention.layer import Attention  # Avoid circular import
-
-        if isinstance(layer, LinearBase):
-            linear_config = self.get_linear_config(layer)
-            if self.ignored_layers and is_layer_skipped(
-                    prefix=prefix,
-                    ignored_layers=self.ignored_layers,
-                    fused_mapping=self.packed_modules_mapping,
-            ):
-                return VllmUnquantizedLinearMethod(linear_config)
-            logger.warning_once(
-                "MXFP4 linear layer is not implemented - falling back to "
-                "UnquantizedLinearMethod.")
-            return VllmUnquantizedLinearMethod(linear_config)
-        elif isinstance(layer, FusedMoE):
-            moe_config = self.get_moe_config(layer)
-            return VllmMxfp4MoEMethod(moe_config, self.mesh)
-        elif isinstance(layer, Attention):
-            logger.warning_once("MXFP4 attention layer is not implemented. "
-                                "Skipping quantization for this layer.")
-        return None
-
-
-class VllmMxfp4MoEMethod(Mxfp4MoEMethod):
-
-    def __init__(self,
-                 moe: FusedMoEConfig,
-                 mesh: Mesh,
-                 ep_axis_name: str = 'model'):
-        FusedMoEMethodBase.__init__(self, moe)
-
-        # We piggyback on triton implementation as it applies minimal hardware
-        # specific post processing to the weights.
-        self.mxfp4_backend = Mxfp4Backend.TRITON
-
-        self.mesh = mesh
-        self.use_kernel = envs.USE_MOE_EP_KERNEL and moe.use_ep
-        self.ep_axis_name = ep_axis_name
-        # TODO: Use autotune table once we have it.
-        self.block_size = {
-            "bt": 64,
-            "bf": 1024,
-            "bd1": 1536,
-            "bd2": 1536,
-            "btc": 64,
-            "bfc": 1024,
-            "bd1c": 1536,
-            "bd2c": 1536,
-        }
-
-    def get_fused_moe_quant_config(
-            self, layer: torch.nn.Module) -> FusedMoEQuantConfig | None:
-        # Because we have dequantized weights, we only need biased moe config.
-        # TODO(kyuyeunk): Add native support for MXFP4.
-        return biased_moe_quant_config(
-            layer.w13_bias,
-            layer.w2_bias,
-        )
-
-    def process_weights_after_loading(self, layer: torch.nn.Module):
-        assert isinstance(layer, FusedMoE)
-        assert layer.moe_config.has_bias, "mxfp4 quantization alwyas use bias."
-
-        w13_weight = u8_unpack_e2m1(t2j(layer.w13_weight, use_dlpack=False))
-        w13_weight_scale = e8m0_to_fp32(
-            t2j(layer.w13_weight_scale, use_dlpack=False))
-        w13_bias = t2j(layer.w13_bias, use_dlpack=False)
-
-        w2_weight = u8_unpack_e2m1(t2j(layer.w2_weight, use_dlpack=False))
-        w2_weight_scale = e8m0_to_fp32(
-            t2j(layer.w2_weight_scale, use_dlpack=False))
-        w2_bias = t2j(layer.w2_bias, use_dlpack=False)
-
-        # We dequantize fp4 weights into bf16.
-        # TODO(kyuyeunk): Add native support for MXFP4.
-        w13_weight = dequantize_block_weight(w13_weight, w13_weight_scale,
-                                             MXFP4_BLOCK_SIZE, jnp.bfloat16)
-        w2_weight = dequantize_block_weight(w2_weight, w2_weight_scale,
-                                            MXFP4_BLOCK_SIZE, jnp.bfloat16)
-
-        num_experts, hidden_size, intermediate_size = w2_weight.shape
-
-        # Because we have dequantized weights, scales are not used anymore.
-        delattr(layer, "w13_weight_scale")
-        delattr(layer, "w2_weight_scale")
-
-        if layer.activation == "swigluoai":
-            # When using swigluoai, vLLM splits gmm output in a interleaved way.
-            # However, interleaved split is not performant on TPU. Therefore,
-            # we preprocess the weight so that splitting gmm output by middle
-            # can still get the same result.
-            w1_weight = w13_weight[:, ::2, :]
-            w3_weight = w13_weight[:, 1::2, :]
-            w13_weight = jnp.concat([w1_weight, w3_weight], axis=1)
-
-            w1_bias = w13_bias[:, ::2]
-            w3_bias = w13_bias[:, 1::2]
-            w13_bias = jnp.concat([w1_bias, w3_bias], axis=1)
-
-        if self.use_kernel:
-            # Kernel expects:
-            # w13: (num_experts, 2, hidden_size, intermediate_size)
-            # w2: (num_experts, intermediate_size, hidden_size)
-            # Current format:
-            # w13_weight: (num_experts, 2*intermediate_size, hidden_size)
-            # w2_weight: (num_experts, hidden_size, intermediate_size)
-
-            w13_reshaped = w13_weight.reshape(num_experts, 2,
-                                              intermediate_size, hidden_size)
-
-            # Transpose non-constracting dim to right most dim
-            w13_weight_transposed = jnp.swapaxes(w13_reshaped, 2, 3)
-            w2_weight_transposed = jnp.swapaxes(w2_weight, 1, 2)
-
-            # Apply EP sharding
-            ep_sharding = NamedSharding(self.mesh, P("model"))
-
-            w13_weight = jax.device_put(
-                w13_weight_transposed, Format(Layout((0, 1, 2, 3)),
-                                              ep_sharding))
-            w2_weight = jax.device_put(w2_weight_transposed,
-                                       Format(Layout((0, 1, 2)), ep_sharding))
-
-            w13_bias = w13_bias.reshape(num_experts, 2, intermediate_size)
-            w13_bias = jax.device_put(w13_bias,
-                                      Format(Layout((0, 1, 2)), ep_sharding))
-            w2_bias = jax.device_put(w2_bias,
-                                     Format(Layout((0, 1)), ep_sharding))
-
-        else:
-            if layer.use_ep:
-                ep_sharding = NamedSharding(self.mesh, P("model"))
-                w13_weight = jax.device_put(
-                    w13_weight, Format(Layout((0, 1, 2)), ep_sharding))
-                w2_weight = jax.device_put(
-                    w2_weight, Format(Layout((0, 1, 2)), ep_sharding))
-
-                w13_bias = jax.device_put(w13_bias,
-                                          Format(Layout((0, 1)), ep_sharding))
-                w2_bias = jax.device_put(w2_bias,
-                                         Format(Layout((0, 1)), ep_sharding))
-
-            else:
-                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_bias = reorder_concatenated_tensor_for_sharding(
-                    w13_bias,
-                    output_sizes,
-                    n_shards,
-                    dim=1,
-                )
-                w13_bias = jax.device_put(
-                    w13_bias,
-                    Format(Layout((0, 1)),
-                           NamedSharding(self.mesh, P(None, "model"))))
-                w2_bias = jax.device_put(
-                    w2_bias,
-                    Format(Layout((0, 1)),
-                           NamedSharding(self.mesh, P(None, None))))
-
-        layer.w13_weight = Parameter(torch_view(w13_weight),
-                                     requires_grad=False)
-        layer.w2_weight = Parameter(torch_view(w2_weight), requires_grad=False)
-
-        layer.w13_bias = Parameter(torch_view(w13_bias), requires_grad=False)
-        layer.w2_bias = Parameter(torch_view(w2_bias), requires_grad=False)
-
-        pass
-
-    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 scoring_func != "softmax":
-            raise NotImplementedError(
-                "Only softmax is supported for scoring_func")
-
-        x = jax_view(x)
-        w13_weight = jax_view(layer.w13_weight)
-        w2_weight = jax_view(layer.w2_weight)
-        w13_bias = jax_view(layer.w13_bias)
-        w2_bias = jax_view(layer.w2_bias)
-        gating_output = jax_view(router_logits)
-
-        if self.use_kernel:
-            output = fused_ep_moe(
-                mesh=self.mesh,
-                tokens=x,
-                w1=w13_weight,
-                w2=w2_weight,
-                b1=w13_bias,
-                b2=w2_bias,
-                gating_output=gating_output,
-                top_k=top_k,
-                ep_axis_name=self.ep_axis_name,
-                renormalize_topk_logits=renormalize,
-                act_fn=activation,
-                **self.block_size,
-            )
-        else:
-            output = fused_moe_func(
-                hidden_states=x,
-                w1=w13_weight,
-                w2=w2_weight,
-                w1_bias=w13_bias,
-                w2_bias=w2_bias,
-                gating_output=gating_output,
-                topk=top_k,
-                renormalize=renormalize,
-                mesh=self.mesh,
-                use_ep=layer.use_ep,
-                activation=activation,
-            )
-
-        return torch_view(output)

tpu_inference/models/jax/llama_guard_4.py

@@ -1,361 +0,0 @@
-import re
-from typing import Any, List, Optional, Tuple
-
-import jax
-import jax.numpy as jnp
-import torch
-from flax import nnx
-from flax.typing import PRNGKey
-from jax.sharding import Mesh
-from jax.sharding import PartitionSpec as P
-from vllm.config import VllmConfig
-
-from tpu_inference.layers.jax.attention.attention import AttentionMetadata
-from tpu_inference.layers.jax.attention.llama4_attention import Llama4Attention
-from tpu_inference.layers.jax.constants import KVCacheType
-from tpu_inference.layers.jax.layers import DenseFFW, Embedder, LMhead, RMSNorm
-from tpu_inference.layers.jax.misc import shard_put
-from tpu_inference.layers.jax.transformer_block import TransformerBlock
-from tpu_inference.logger import init_logger
-from tpu_inference.models.jax.utils.weight_utils import (
-    get_param, model_weights_generator, print_param_info, reshape_params,
-    transpose_params)
-
-logger = init_logger(__name__)
-
-
-class LlamaGuard4ForCausalLM(nnx.Module):
-
-    def __init__(self,
-                 vllm_config: VllmConfig,
-                 rng: PRNGKey,
-                 mesh: Mesh,
-                 force_random_weights: bool = False):
-        logger.warning(
-            "🚨🚨🚨WARNING🚨🚨🚨 🚨🚨🚨WARNING🚨🚨🚨 🚨🚨🚨WARNING🚨🚨🚨\n"
-            "Llama Guard 4 (JAX) is WIP: Only the text modality is currently implemented.  "
-            "Multimodal inputs will fail.\n"
-            "🚨🚨🚨WARNING🚨🚨🚨 🚨🚨🚨WARNING🚨🚨🚨 🚨🚨🚨WARNING🚨🚨🚨")
-        assert mesh is not None
-
-        self.vllm_config = vllm_config
-        self.vllm_config.model_config.dtype = torch.bfloat16
-        model_config = vllm_config.model_config
-        text_config = model_config.hf_config.text_config
-
-        self.mesh = mesh
-        self.is_verbose = getattr(self.vllm_config.additional_config,
-                                  "is_verbose", False)
-
-        self.use_qk_norm = getattr(text_config, "use_qk_norm", True)
-
-        vocab_size = model_config.get_vocab_size()
-        self.hidden_size = model_config.get_hidden_size()
-
-        self.dtype: jnp.dtype = jnp.bfloat16
-
-        self.num_layers: int = getattr(text_config, "num_layers", 48)
-        hidden_act: str = getattr(text_config, "hidden_act", "silu")
-
-        rms_norm_eps = getattr(text_config, "rms_norm_eps", 1e-5)
-        self.num_attention_heads = getattr(text_config, "num_attention_heads",
-                                           40)
-        self.num_key_value_heads = getattr(text_config, "num_key_value_heads",
-                                           8)
-        self.head_dim = getattr(text_config, "head_dim", 128)
-
-        intermediate_size = getattr(text_config, "intermediate_size", 8192)
-
-        self.rope_theta_text = getattr(text_config, "rope_theta", 500000.0)
-        self.rope_scaling = getattr(text_config, "rope_scaling")
-
-        self.rng = nnx.Rngs(rng)
-
-        self.embedder = Embedder(
-            vocab_size=vocab_size,
-            hidden_size=self.hidden_size,
-            dtype=self.dtype,
-            vd_sharding=(('data', 'model'), None),
-            rngs=self.rng,
-            random_init=force_random_weights,
-        )
-
-        self.layers = []
-
-        for i in range(self.num_layers):
-            use_attention_rope = True
-
-            custom_module = DenseFFW(dtype=self.dtype,
-                                     hidden_act=hidden_act,
-                                     hidden_size=self.hidden_size,
-                                     intermediate_size=intermediate_size,
-                                     random_init=force_random_weights,
-                                     rngs=self.rng,
-                                     df_sharding=P(None, 'model'),
-                                     fd_sharding=P('model', None),
-                                     activation_ffw_td=P('data', None))
-
-            attn = Llama4Attention(
-                hidden_size=self.hidden_size,
-                dtype=self.dtype,
-                num_attention_heads=self.num_attention_heads,
-                num_key_value_heads=self.num_key_value_heads,
-                head_dim=self.head_dim,
-                rope_theta=self.rope_theta_text,
-                rope_scaling={
-                    "scale_factor":
-                    self.rope_scaling["factor"],
-                    "low_freq_factor":
-                    self.rope_scaling["low_freq_factor"],
-                    "high_freq_factor":
-                    self.rope_scaling["high_freq_factor"],
-                    "original_max_position_embeddings":
-                    self.rope_scaling["original_max_position_embeddings"]
-                },
-                rngs=self.rng,
-                rope_input_ordering="interleaved",
-                # TODO (jacobplatin): we should refactor this to pass a dtype (or config) directly
-                kv_cache_dtype=vllm_config.cache_config.cache_dtype,
-                temperature_tuning=True,
-                temperature_tuning_scale=0.1,
-                temperature_tuning_floor_scale=8192,
-                use_qk_norm=self.use_qk_norm,
-                attention_chunk_size=None if use_attention_rope else 8192,
-                mesh=self.mesh,
-                random_init=force_random_weights,
-                activation_attention_td=('data', 'model'),
-                activation_q_td=('data', 'model'),
-                query_tnh=P('data', 'model', None),
-                keyvalue_skh=P('data', 'model', None),
-                activation_attention_out_td=('data', 'model'),
-                attn_o_tnh=P('data', 'model', None),
-                dnh_sharding=(None, 'model', None),
-                dkh_sharding=(None, 'model', None),
-                nhd_sharding=('model', None, None),
-            )
-
-            pre_attention_norm = RMSNorm(
-                dims=self.hidden_size,
-                random_init=force_random_weights,
-                epsilon=rms_norm_eps,
-                rngs=self.rng,
-                activation_ffw_td=('data', None),
-                with_scale=True,
-                dtype=self.dtype,
-            )
-
-            pre_mlp_norm = RMSNorm(
-                dims=self.hidden_size,
-                activation_ffw_td=('data', None),
-                epsilon=rms_norm_eps,
-                rngs=self.rng,
-                with_scale=True,
-                dtype=self.dtype,
-                random_init=force_random_weights,
-            )
-
-            block = TransformerBlock(custom_module=custom_module,
-                                     attn=attn,
-                                     pre_attention_norm=pre_attention_norm,
-                                     pre_mlp_norm=pre_mlp_norm,
-                                     use_attention_rope=use_attention_rope)
-            self.layers.append(block)
-
-        self.final_norm = RMSNorm(
-            dims=self.hidden_size,
-            activation_ffw_td=P(),
-            epsilon=rms_norm_eps,
-            rngs=self.rng,
-            with_scale=True,
-            dtype=self.dtype,
-            random_init=force_random_weights,
-        )
-
-        self.lm_head = LMhead(vocab_size=vocab_size,
-                              hidden_size=self.hidden_size,
-                              dtype=self.dtype,
-                              rngs=self.rng,
-                              vd_sharding=(('data', 'model'), None),
-                              dv_sharding=(None, ('data', 'model')),
-                              random_init=force_random_weights)
-        if self.is_verbose:
-            self._print_model_architecture()
-
-    def _print_model_architecture(self):
-
-        logger.info("### Embedding ###")
-        nnx.display(self.embedder)
-
-        logger.info("\n### Layers ###")
-        for i, layer in enumerate(self.layers):
-            logger.info(f"\n--- Layer {i} ---")
-            nnx.display(layer)
-
-        logger.info("\n### LM Head ###")
-        nnx.display(self.lm_head)
-
-    def load_weights(self, rng: jax.Array, cache_dir: Optional[str] = None):
-        self.rng = nnx.Rngs(rng)
-
-        weight_loader = LlamaGuard4WeightLoader(
-            vllm_config=self.vllm_config,
-            hidden_size=self.hidden_size,
-            attn_heads=self.num_attention_heads,
-            num_key_value_heads=self.num_key_value_heads,
-            attn_head_dim=self.head_dim)
-        weight_loader.load_weights(self)
-
-    def __call__(
-        self,
-        kv_caches: List[jax.Array],
-        input_ids: jax.Array,
-        attention_metadata: AttentionMetadata,
-        inputs_embeds: Optional[jax.Array] = None,
-        layer_metadata_tuple: Optional[Tuple] = None,
-        lora_metadata: Optional[Any] = None,
-        *args,
-    ) -> Tuple[List[KVCacheType], jax.Array]:
-        is_prefill = False
-
-        if inputs_embeds is not None:
-            x_TD = inputs_embeds
-        elif input_ids is not None:
-            x_TD = self.embedder.encode(input_ids)
-        else:
-            raise ValueError(
-                "Cannot run forward pass: Both input_ids and inputs_embeds are None."
-            )
-
-        for (i, block) in enumerate(self.layers):
-            kv_cache = kv_caches[i]
-            new_kv_cache, x_TD = block(x_TD, is_prefill, kv_cache,
-                                       attention_metadata)
-            jax.block_until_ready(x_TD)
-            kv_caches[i] = new_kv_cache
-
-        final_activation_TD = self.final_norm(x_TD)
-
-        return kv_caches, final_activation_TD, []
-
-    def compute_logits(self, hidden_states: jax.Array) -> jax.Array:
-        logits_TV = jnp.dot(hidden_states,
-                            self.lm_head.input_embedding_table_DV.value)
-        return logits_TV
-
-    def get_input_embeddings(
-            self,
-            input_ids: jax.Array,
-            multimodal_embeddings: Optional[List[jax.Array]] = None
-    ) -> jax.Array:
-        """
-        Computes the embeddings for text input (used for input to fusion).
-        """
-        return self.embedder.encode(input_ids)
-
-
-class LlamaGuard4WeightLoader:
-
-    def __init__(self, vllm_config: VllmConfig, hidden_size, attn_heads,
-                 num_key_value_heads, attn_head_dim):
-        self.names_and_weights_generator = model_weights_generator(
-            model_name_or_path=vllm_config.model_config.model,
-            framework="flax",
-            filter_regex="language_model",
-            download_dir=vllm_config.load_config.download_dir)
-        self.is_verbose = getattr(vllm_config.additional_config, "is_verbose",
-                                  False)
-        self._transpose_map = {
-            "q_proj": (2, 0, 1),
-            "k_proj": (2, 0, 1),
-            "v_proj": (2, 0, 1),
-            "o_proj": (1, 2, 0),
-            "lm_head": (1, 0),
-            "feed_forward.down_proj": (1, 0),
-            "feed_forward.gate_proj": (1, 0),
-            "feed_forward.up_proj": (1, 0),
-            "mlp.down_proj": (1, 0),
-            "mlp.gate_proj": (1, 0),
-            "mlp.up_proj": (1, 0),
-        }
-        self._weight_shape_map = {
-            "q_proj": (attn_heads, attn_head_dim, hidden_size),
-            "k_proj": (num_key_value_heads, attn_head_dim, hidden_size),
-            "v_proj": (num_key_value_heads, attn_head_dim, hidden_size),
-            "o_proj": (hidden_size, attn_heads, attn_head_dim),
-        }
-
-        self._loaded_to_standardized_keys = {
-            "language_model.model.embed_tokens.weight":
-            "embedder.input_embedding_table_VD",
-            "language_model.lm_head.weight":
-            "lm_head.input_embedding_table_DV",
-            "language_model.model.norm.weight":
-            "final_norm.scale",
-            "language_model.model.layers.*.input_layernorm.weight":
-            "layers.*.pre_attention_norm.scale",
-            "language_model.model.layers.*.post_attention_layernorm.weight":
-            "layers.*.pre_mlp_norm.scale",
-            "language_model.model.layers.*.self_attn.q_proj.weight":
-            "layers.*.attn.kernel_q_proj_DNH",
-            "language_model.model.layers.*.self_attn.k_proj.weight":
-            "layers.*.attn.kernel_k_proj_DKH",
-            "language_model.model.layers.*.self_attn.v_proj.weight":
-            "layers.*.attn.kernel_v_proj_DKH",
-            "language_model.model.layers.*.self_attn.o_proj.weight":
-            "layers.*.attn.kernel_o_proj_NHD",
-            "language_model.model.layers.*.feed_forward.gate_proj.weight":
-            "layers.*.custom_module.kernel_gating_DF",
-            "language_model.model.layers.*.feed_forward.up_proj.weight":
-            "layers.*.custom_module.kernel_up_proj_DF",
-            "language_model.model.layers.*.feed_forward.down_proj.weight":
-            "layers.*.custom_module.kernel_down_proj_FD",
-        }
-
-    def map_loaded_to_standardized_name(self, loaded_key: str) -> str:
-        if "layer" in loaded_key:
-            layer_num = re.search(r"layers\.(\d+)", loaded_key).group(1)
-            layer_key = re.sub(r"layers\.\d+", "layers.*", loaded_key)
-            mapped_key = self._loaded_to_standardized_keys.get(
-                layer_key, loaded_key)
-            mapped_key = re.sub(r"layers\.\*", f"layers.{layer_num}",
-                                mapped_key)
-        else:
-            mapped_key = self._loaded_to_standardized_keys.get(
-                loaded_key, loaded_key)
-        return mapped_key
-
-    def load_weights(self, model_for_loading: nnx.Module):
-        model_params = nnx.state(model_for_loading)
-        with jax.default_device(jax.devices("cpu")[0]):
-            for loaded_name, loaded_weight in self.names_and_weights_generator:
-                if loaded_name.endswith(".bias"):
-                    continue
-                if "vision_model" in loaded_name or "multi_modal_projector" in loaded_name:
-                    continue
-
-                mapped_name = self.map_loaded_to_standardized_name(loaded_name)
-                model_weight = get_param(model_params, mapped_name)
-
-                if not loaded_name.endswith(".bias"):
-                    # For other layers, continue to use the transpose_params helper.
-                    loaded_weight = reshape_params(loaded_name, loaded_weight,
-                                                   self._weight_shape_map)
-                    loaded_weight = transpose_params(loaded_name,
-                                                     loaded_weight,
-                                                     self._transpose_map)
-                if model_weight.value.shape != loaded_weight.shape:
-                    raise ValueError(
-                        f"Loaded shape for {loaded_name}: {loaded_weight.shape} "
-                        f"does not match model shape for {mapped_name}: {model_weight.value.shape}!"
-                    )
-                logger.debug(
-                    f"Transformed parameter {loaded_name} to {mapped_name}: {loaded_weight.shape} --> {model_weight.value.shape}"
-                )
-
-                model_weight.value = shard_put(loaded_weight,
-                                               model_weight.sharding,
-                                               mesh=model_for_loading.mesh)
-                if self.is_verbose:
-                    print_param_info(model_weight, loaded_name)
-
-        nnx.update(model_for_loading, model_params)