sglang 0.4.3.post2__py3-none-any.whl → 0.4.3.post3__py3-none-any.whl

sglang/srt/models/mllama.py

@@ -36,6 +36,7 @@ from sglang.srt.managers.schedule_batch import ImageInputs
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.llama import LlamaDecoderLayer, LlamaMLP
+from sglang.srt.utils import add_prefix
 
 
 class ColumnParallelConv2dPatch(torch.nn.Module):
@@ -147,7 +148,12 @@ class MllamaPrecomputedPositionEmbedding(nn.Module):
 
 
 class MllamaVisionMLP(nn.Module):
-    def __init__(self, config, quant_config: Optional[QuantizationConfig] = None):
+    def __init__(
+        self,
+        config,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ):
         super().__init__()
         self.config = config
         self.activation_fn = get_act_fn(config.hidden_act)
@@ -156,12 +162,14 @@ class MllamaVisionMLP(nn.Module):
             config.intermediate_size,
             bias=True,
             quant_config=quant_config,
+            prefix=add_prefix("fc1", prefix),
         )
         self.fc2 = RowParallelLinear(
             config.intermediate_size,
             config.hidden_size,
             bias=True,
             quant_config=quant_config,
+            prefix=add_prefix("fc2", prefix),
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -174,7 +182,10 @@ class MllamaVisionMLP(nn.Module):
 
 class MllamaVisionEncoderLayer(nn.Module):
     def __init__(
-        self, config: config_mllama.MllamaVisionConfig, is_gated: bool = False
+        self,
+        config: config_mllama.MllamaVisionConfig,
+        is_gated: bool = False,
+        prefix: str = "",
     ):
         super().__init__()
 
@@ -193,8 +204,9 @@ class MllamaVisionEncoderLayer(nn.Module):
             use_context_forward=False,
             use_full_precision_softmax=False,
             flatten_batch=False,
+            prefix=add_prefix("self_attn", prefix),
         )
-        self.mlp = MllamaVisionMLP(config)
+        self.mlp = MllamaVisionMLP(config, prefix=add_prefix("mlp", prefix))
 
         self.input_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps)
         self.post_attention_layernorm = nn.LayerNorm(
@@ -235,11 +247,17 @@ class MllamaVisionEncoder(nn.Module):
         num_layers=32,
         is_gated=False,
         output_hidden_states=None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
         self.layers = nn.ModuleList(
-            [MllamaVisionEncoderLayer(config, is_gated) for _ in range(num_layers)]
+            [
+                MllamaVisionEncoderLayer(
+                    config, is_gated, prefix=add_prefix(f"layers.{i}", prefix)
+                )
+                for i in range(num_layers)
+            ]
         )
         self.output_hidden_states = output_hidden_states or []
 
@@ -265,7 +283,7 @@ class MllamaVisionEncoder(nn.Module):
 
 
 class MllamaVisionModel(nn.Module):
-    def __init__(self, config: config_mllama.MllamaVisionConfig):
+    def __init__(self, config: config_mllama.MllamaVisionConfig, prefix: str = ""):
         super().__init__()
         self.image_size = config.image_size
         self.patch_size = config.patch_size
@@ -305,9 +323,13 @@ class MllamaVisionModel(nn.Module):
             config.num_hidden_layers,
             is_gated=False,
             output_hidden_states=config.intermediate_layers_indices,
+            prefix=add_prefix("transformer", prefix),
         )
         self.global_transformer = MllamaVisionEncoder(
-            config, config.num_global_layers, is_gated=True
+            config,
+            config.num_global_layers,
+            is_gated=True,
+            prefix=add_prefix("global_transformer", prefix),
         )
 
     def apply_class_embedding(self, hidden_state: torch.Tensor) -> torch.Tensor:
@@ -464,6 +486,7 @@ class MllamaTextCrossAttention(nn.Module):
         config: Optional[config_mllama.MllamaTextConfig] = None,
         layer_id: Optional[int] = None,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
@@ -489,6 +512,7 @@ class MllamaTextCrossAttention(nn.Module):
             self.num_key_value_heads,
             bias=False,
             quant_config=quant_config,
+            prefix=add_prefix("qkv_proj", prefix),
         )
         self.o_proj = RowParallelLinear(
             self.num_heads * self.head_dim,
@@ -496,6 +520,7 @@ class MllamaTextCrossAttention(nn.Module):
             bias=False,
             input_is_parallel=True,
             quant_config=quant_config,
+            prefix=add_prefix("o_proj", prefix),
         )
         # vllm.model_executor.layers.layernorm.RMSNorm has precision issue,
         # use huggingface's instead
@@ -510,6 +535,7 @@ class MllamaTextCrossAttention(nn.Module):
             self.num_local_key_value_heads,
             layer_id=layer_id,
             is_cross_attention=True,
+            prefix=add_prefix("attn", prefix),
         )
 
     def forward(
@@ -551,6 +577,7 @@ class MllamaCrossAttentionDecoderLayer(torch.nn.Module):
         config: config_mllama.MllamaTextConfig,
         layer_id: int,
         quant_config: Optional[QuantizationConfig],
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.layer_id = layer_id
@@ -558,6 +585,7 @@ class MllamaCrossAttentionDecoderLayer(torch.nn.Module):
             config=config,
             layer_id=layer_id,
             quant_config=quant_config,
+            prefix=add_prefix("cross_attn", prefix),
         )
 
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -568,6 +596,7 @@ class MllamaCrossAttentionDecoderLayer(torch.nn.Module):
             intermediate_size=config.intermediate_size,
             hidden_act=config.hidden_act,
             quant_config=quant_config,
+            prefix=add_prefix("mlp", prefix),
         )
         self.post_attention_layernorm = RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
@@ -610,12 +639,15 @@ class MllamaTextModel(nn.Module):
         self,
         config: config_mllama.MllamaTextConfig,
         quant_config: Optional[QuantizationConfig],
+        prefix: str = "",
     ):
         super().__init__()
         self.padding_id = config.pad_token_id
         self.vocab_size = config.vocab_size
         self.embed_tokens = VocabParallelEmbedding(
-            config.vocab_size + 8, config.hidden_size
+            config.vocab_size + 8,
+            config.hidden_size,
+            prefix=add_prefix("embed_tokens", prefix),
         )
         self.cross_attention_layers = config.cross_attention_layers
 
@@ -624,14 +656,20 @@ class MllamaTextModel(nn.Module):
             if layer_id in self.cross_attention_layers:
                 layers.append(
                     MllamaCrossAttentionDecoderLayer(
-                        config, layer_id, quant_config=quant_config
+                        config,
+                        layer_id,
+                        quant_config=quant_config,
+                        prefix=add_prefix(f"layers.{layer_id}", prefix),
                     )
                 )
             else:
                 # TODO: force LlamaDecoderLayer to config.attention_bias=False
                 layers.append(
                     LlamaDecoderLayer(
-                        config, quant_config=quant_config, layer_id=layer_id
+                        config,
+                        quant_config=quant_config,
+                        layer_id=layer_id,
+                        prefix=add_prefix(f"layers.{layer_id}", prefix),
                     )
                 )
 
@@ -687,16 +725,20 @@ class MllamaForCausalLM(nn.Module):
         self,
         config: config_mllama.MllamaTextConfig,
         quant_config: Optional[QuantizationConfig],
+        prefix: str = "",
     ):
         super().__init__()
         self.vocab_size = config.vocab_size
-        self.model = MllamaTextModel(config, quant_config)
+        self.model = MllamaTextModel(
+            config, quant_config, prefix=add_prefix("model", prefix)
+        )
         self.lm_head = ParallelLMHead(
             config.vocab_size,
             config.hidden_size,
             org_num_embeddings=config.vocab_size,
             padding_size=DEFAULT_VOCAB_PADDING_SIZE,
             quant_config=quant_config,
+            prefix=add_prefix("lm_head", prefix),
         )
 
     def forward(
@@ -726,6 +768,7 @@ class MllamaForConditionalGeneration(nn.Module):
         self,
         config: config_mllama.MllamaConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.vocab_size = config.text_config.vocab_size
@@ -737,10 +780,13 @@ class MllamaForConditionalGeneration(nn.Module):
         )
         self.image_size = config.vision_config.image_size
 
-        self.vision_model = MllamaVisionModel(config.vision_config)
+        self.vision_model = MllamaVisionModel(
+            config.vision_config, prefix=add_prefix("vision_model", prefix)
+        )
         self.language_model = MllamaForCausalLM(
             config.text_config,
             quant_config=quant_config,
+            prefix=add_prefix("language_model", prefix),
         )
         self.multi_modal_projector = nn.Linear(
             config.vision_config.vision_output_dim,

sglang/srt/models/olmo.py

@@ -38,7 +38,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
-from sglang.srt.utils import make_layers
+from sglang.srt.utils import add_prefix, make_layers
 
 
 class OlmoAttention(nn.Module):
@@ -53,6 +53,7 @@ class OlmoAttention(nn.Module):
         config: OlmoConfig,
         layer_id: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
@@ -75,6 +76,7 @@ class OlmoAttention(nn.Module):
             self.head_dim,
             self.total_num_heads,
             bias=config.attention_bias,
+            prefix=add_prefix("qkv_proj", prefix),
         )
 
         # Rotary embeddings.
@@ -91,6 +93,7 @@ class OlmoAttention(nn.Module):
             self.scaling,
             num_kv_heads=self.num_heads,
             layer_id=layer_id,
+            prefix=add_prefix("attn", prefix),
         )
 
         # Attention output projection.
@@ -98,6 +101,7 @@ class OlmoAttention(nn.Module):
             self.hidden_size,
             self.hidden_size,
             bias=config.attention_bias,
+            prefix=add_prefix("o_proj", prefix),
         )
 
     def forward(
@@ -127,6 +131,7 @@ class OlmoMLP(nn.Module):
         self,
         config: OlmoConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
@@ -139,6 +144,7 @@ class OlmoMLP(nn.Module):
             [self.intermediate_size] * 2,
             bias=False,
             quant_config=quant_config,
+            prefix=add_prefix("gate_up_proj", prefix),
         )
 
         # Activation function.
@@ -150,6 +156,7 @@ class OlmoMLP(nn.Module):
             self.hidden_size,
             bias=False,
             quant_config=quant_config,
+            prefix=add_prefix("down_proj", prefix),
         )
 
     def forward(
@@ -174,13 +181,23 @@ class OlmoDecoderLayer(nn.Module):
         config: OlmoConfig,
         layer_id: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         # Attention block.
-        self.self_attn = OlmoAttention(config, layer_id, quant_config)
+        self.self_attn = OlmoAttention(
+            config,
+            layer_id,
+            quant_config,
+            prefix=add_prefix("self_attn", prefix),
+        )
 
         # MLP block.
-        self.mlp = OlmoMLP(config, quant_config)
+        self.mlp = OlmoMLP(
+            config,
+            quant_config,
+            prefix=add_prefix("mlp", prefix),
+        )
 
         # LayerNorm
         self.input_layernorm = nn.LayerNorm(
@@ -213,13 +230,18 @@ class OlmoDecoderLayer(nn.Module):
 class OlmoModel(nn.Module):
 
     def __init__(
-        self, config: OlmoConfig, quant_config: Optional[QuantizationConfig] = None
+        self,
+        config: OlmoConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
 
         self.embed_tokens = VocabParallelEmbedding(
-            config.vocab_size, config.hidden_size
+            config.vocab_size,
+            config.hidden_size,
+            prefix=add_prefix("embed_tokens", prefix),
         )
         self.layers = make_layers(
             config.num_hidden_layers,
@@ -227,7 +249,9 @@ class OlmoModel(nn.Module):
                 layer_id=idx,
                 config=config,
                 quant_config=quant_config,
+                prefix=prefix,
             ),
+            prefix=add_prefix("layers", prefix),
         )
         self.norm = nn.LayerNorm(
             config.hidden_size, elementwise_affine=False, bias=False
@@ -275,10 +299,11 @@ class OlmoForCausalLM(nn.Module):
         self,
         config: OlmoConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
-        self.model = OlmoModel(config, quant_config)
+        self.model = OlmoModel(config, quant_config, prefix=add_prefix("model", prefix))
         if config.tie_word_embeddings:
             self.lm_head = self.model.embed_tokens
         else:
@@ -288,6 +313,7 @@ class OlmoForCausalLM(nn.Module):
                 config.hidden_size,
                 org_num_embeddings=config.vocab_size,
                 quant_config=quant_config,
+                prefix=add_prefix("lm_head", prefix),
             )
         self.logits_processor = LogitsProcessor(config)
 
@@ -325,6 +351,8 @@ class OlmoForCausalLM(nn.Module):
                 # Models trained using ColossalAI may include these tensors in
                 # the checkpoint. Skip them.
                 continue
+            if self.config.tie_word_embeddings and "lm_head.weight" in name:
+                continue
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue

sglang/srt/models/olmo2.py

@@ -45,7 +45,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
-from sglang.srt.utils import make_layers
+from sglang.srt.utils import add_prefix, make_layers
 
 
 class Olmo2Attention(nn.Module):
@@ -60,28 +60,29 @@ class Olmo2Attention(nn.Module):
         config: PretrainedConfig,
         layer_id: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        self.tp_size = get_tensor_model_parallel_world_size()
         self.total_num_heads = config.num_attention_heads
 
         assert self.hidden_size % self.total_num_heads == 0
-        assert self.total_num_heads % tp_size == 0
+        assert self.total_num_heads % self.tp_size == 0
 
-        self.num_heads = self.total_num_heads // tp_size
+        self.num_heads = self.total_num_heads // self.tp_size
         self.total_num_kv_heads = self.config.num_key_value_heads
 
-        if self.total_num_kv_heads >= tp_size:
+        if self.total_num_kv_heads >= self.tp_size:
             # Number of KV heads is greater than TP size, so we partition
             # the KV heads across multiple tensor parallel GPUs.
-            assert self.total_num_kv_heads % tp_size == 0
+            assert self.total_num_kv_heads % self.tp_size == 0
         else:
             # Number of KV heads is less than TP size, so we replicate
             # the KV heads across multiple tensor parallel GPUs.
-            assert tp_size % self.total_num_kv_heads == 0
-        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+            assert self.tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size)
 
         self.head_dim = self.hidden_size // self.total_num_heads
         self.max_position_embeddings = config.max_position_embeddings
@@ -93,6 +94,8 @@ class Olmo2Attention(nn.Module):
             self.head_dim,
             self.total_num_heads,
             bias=config.attention_bias,
+            quant_config=quant_config,
+            prefix=add_prefix("qkv_proj", prefix),
         )
         self.tp_rank = get_tensor_model_parallel_rank()
 
@@ -115,6 +118,7 @@ class Olmo2Attention(nn.Module):
             self.scaling,
             num_kv_heads=self.num_kv_heads,
             layer_id=layer_id,
+            prefix=add_prefix("attn", prefix),
         )
 
         # Attention output projection.
@@ -122,6 +126,8 @@ class Olmo2Attention(nn.Module):
             self.head_dim * self.total_num_heads,
             self.hidden_size,
             bias=config.attention_bias,
+            quant_config=quant_config,
+            prefix=add_prefix("o_proj", prefix),
         )
 
     def _apply_qk_norm(
@@ -164,6 +170,7 @@ class Olmo2MLP(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
@@ -176,6 +183,7 @@ class Olmo2MLP(nn.Module):
             [self.intermediate_size] * 2,
             bias=False,
             quant_config=quant_config,
+            prefix=add_prefix("gate_up_proj", prefix),
         )
 
         # Activation function.
@@ -187,6 +195,7 @@ class Olmo2MLP(nn.Module):
             self.hidden_size,
             bias=False,
             quant_config=quant_config,
+            prefix=add_prefix("down_proj", prefix),
         )
 
     def forward(
@@ -211,13 +220,16 @@ class Olmo2DecoderLayer(nn.Module):
         config: PretrainedConfig,
         layer_id: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         # Attention block.
-        self.self_attn = Olmo2Attention(config, layer_id, quant_config)
+        self.self_attn = Olmo2Attention(
+            config, layer_id, quant_config, prefix=add_prefix("self_attn", prefix)
+        )
 
         # MLP block.
-        self.mlp = Olmo2MLP(config, quant_config)
+        self.mlp = Olmo2MLP(config, quant_config, prefix=add_prefix("mlp", prefix))
 
         # RMSNorm
         self.post_attention_layernorm = RMSNorm(
@@ -254,12 +266,15 @@ class Olmo2Model(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
 
         self.embed_tokens = VocabParallelEmbedding(
-            config.vocab_size, config.hidden_size
+            config.vocab_size,
+            config.hidden_size,
+            prefix=add_prefix("embed_tokens", prefix),
         )
         self.layers = make_layers(
             config.num_hidden_layers,
@@ -267,7 +282,9 @@ class Olmo2Model(nn.Module):
                 layer_id=idx,
                 config=config,
                 quant_config=quant_config,
+                prefix=prefix,
             ),
+            prefix=add_prefix("layers", prefix),
         )
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
@@ -313,10 +330,13 @@ class Olmo2ForCausalLM(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.config = config
-        self.model = Olmo2Model(config, quant_config)
+        self.model = Olmo2Model(
+            config, quant_config, prefix=add_prefix("model", prefix)
+        )
         if config.tie_word_embeddings:
             self.lm_head = self.model.embed_tokens
         else:
@@ -326,6 +346,7 @@ class Olmo2ForCausalLM(nn.Module):
                 config.hidden_size,
                 org_num_embeddings=config.vocab_size,
                 quant_config=quant_config,
+                prefix=add_prefix("lm_head", prefix),
             )
         self.logits_processor = LogitsProcessor(config)
 
@@ -343,7 +364,7 @@ class Olmo2ForCausalLM(nn.Module):
             input_embeds=input_embeds,
         )
         return self.logits_processor(
-            input_ids, hidden_states, self.lm_head.weight, forward_batch
+            input_ids, hidden_states, self.lm_head, forward_batch
         )
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):

sglang/srt/models/olmoe.py

@@ -41,7 +41,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
-from sglang.srt.utils import make_layers, print_warning_once
+from sglang.srt.utils import add_prefix, make_layers, print_warning_once
 
 
 class OlmoeMoE(nn.Module):
@@ -69,7 +69,11 @@ class OlmoeMoE(nn.Module):
 
         # Gate always runs at half / full precision for now.
         self.gate = ReplicatedLinear(
-            hidden_size, num_experts, bias=False, quant_config=None
+            hidden_size,
+            num_experts,
+            bias=False,
+            quant_config=None,
+            prefix=add_prefix("gate", prefix),
         )
 
         self.experts = FusedMoE(
@@ -81,6 +85,7 @@ class OlmoeMoE(nn.Module):
             renormalize=False,
             quant_config=quant_config,
             tp_size=tp_size,
+            prefix=add_prefix("experts", prefix),
         )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -107,6 +112,7 @@ class OlmoeAttention(nn.Module):
         rope_scaling: Optional[Dict[str, Any]] = None,
         max_position_embeddings: int = 4096,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
@@ -138,6 +144,7 @@ class OlmoeAttention(nn.Module):
             self.total_num_kv_heads,
             bias=False,
             quant_config=quant_config,
+            prefix=add_prefix("qkv_proj", prefix),
         )
         self.q_norm = RMSNorm(hidden_size, eps=1e-5)
         self.k_norm = RMSNorm(hidden_size, eps=1e-5)
@@ -146,6 +153,7 @@ class OlmoeAttention(nn.Module):
             hidden_size,
             bias=False,
             quant_config=quant_config,
+            prefix=add_prefix("o_proj", prefix),
         )
 
         self.rotary_emb = get_rope(
@@ -162,6 +170,7 @@ class OlmoeAttention(nn.Module):
             self.scaling,
             layer_id=layer_id,
             num_kv_heads=self.num_kv_heads,
+            prefix=add_prefix("attn", prefix),
         )
 
     def forward(
@@ -186,6 +195,7 @@ class OlmoeDecoderLayer(nn.Module):
         config: PretrainedConfig,
         layer_id: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -202,6 +212,7 @@ class OlmoeDecoderLayer(nn.Module):
             rope_scaling=rope_scaling,
             max_position_embeddings=max_position_embeddings,
             quant_config=quant_config,
+            prefix=add_prefix("self_attn", prefix),
         )
 
         self.mlp = OlmoeMoE(
@@ -210,6 +221,7 @@ class OlmoeDecoderLayer(nn.Module):
             hidden_size=config.hidden_size,
             intermediate_size=config.intermediate_size,
             quant_config=quant_config,
+            prefix=add_prefix("mlp", prefix),
         )
         self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
@@ -246,6 +258,7 @@ class OlmoeModel(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.padding_idx = config.pad_token_id
@@ -254,6 +267,7 @@ class OlmoeModel(nn.Module):
         self.embed_tokens = VocabParallelEmbedding(
             config.vocab_size,
             config.hidden_size,
+            prefix=add_prefix("embed_tokens", prefix),
         )
         self.layers = make_layers(
             config.num_hidden_layers,
@@ -261,7 +275,9 @@ class OlmoeModel(nn.Module):
                 config=config,
                 quant_config=quant_config,
                 layer_id=idx,
+                prefix=prefix,
             ),
+            prefix=add_prefix("layers", prefix),
         )
         self.norm = RMSNorm(config.hidden_size, eps=1e-5)
 
@@ -294,13 +310,19 @@ class OlmoeForCausalLM(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.config = config
         self.quant_config = quant_config
-        self.model = OlmoeModel(config, quant_config)
+        self.model = OlmoeModel(
+            config, quant_config, prefix=add_prefix("model", prefix)
+        )
         self.lm_head = ParallelLMHead(
-            config.vocab_size, config.hidden_size, quant_config=quant_config
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=add_prefix("lm_head", prefix),
         )
         self.logits_processor = LogitsProcessor(config)