address reviews, minor fix

Author: Edwardf0t1

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -649,16 +649,12 @@ def _load_w2(
                 shard_size,
                 not self.use_presharded_weights,
             )
-        if not self.use_presharded_weights:
-            if shard_size * tp_rank + shard_size > loaded_weight.shape[shard_dim]:
-                raise ValueError(
-                    f"Shard size {shard_size} at rank {tp_rank} exceeds loaded_weight dimension {loaded_weight.shape[shard_dim]}"
-                )
-            loaded_weight = loaded_weight.narrow(
-                shard_dim, shard_size * tp_rank, shard_size
-            )
         else:
             if not self.use_presharded_weights:
+                if shard_size * tp_rank + shard_size > loaded_weight.shape[shard_dim]:
+                    raise ValueError(
+                        f"Shard size {shard_size} at rank {tp_rank} exceeds loaded_weight dimension {loaded_weight.shape[shard_dim]}"
+                    )
                 loaded_weight = loaded_weight.narrow(
                     shard_dim, shard_size * tp_rank, shard_size
                 )

python/sglang/srt/models/mllama4.py

@@ -1,4 +1,5 @@
 import json as json_lib
+import logging
 import os
 from collections.abc import Iterable
 from typing import List, Optional, Set, Tuple
@@ -27,6 +28,8 @@
 )
 from sglang.srt.utils import add_prefix
 
+logger = logging.getLogger(__name__)
+
 
 class Llama4ForConditionalGeneration(nn.Module):
     packed_modules_mapping = {
@@ -46,6 +49,11 @@ def __init__(
 
         # Check if this is a text-only model (modelopt fp8 llama4 has no vision components)
         self.has_vision = self._has_vision_weights(config)
+        if not self.has_vision:
+            logger.warning(
+                "No vision weights found in checkpoint. Model will run in text-only mode. "
+                "Multimodal capabilities (image processing) will be unavailable."
+            )
 
         if self.has_vision:
             self.vision_model = Llama4VisionModel(config.vision_config)
@@ -225,12 +233,10 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
         )
 
         for name, loaded_weight in weights:
-            if not self._should_load_weight(name):
+            if self._should_skip_weight(name):
                 continue
 
             name = self._transform_weight_name(name)
-            if name is None:
-                continue
 
             if "vision" not in name:
                 name, loaded_weight = self.permute_qk_weight_for_rotary(
@@ -252,9 +258,9 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
 
             self._handle_default_weight(name, loaded_weight, params_dict)
 
-    def _should_load_weight(self, name: str) -> bool:
-        """Check if we should load this weight."""
-        return not ("vision" in name and not self.has_vision)
+    def _should_skip_weight(self, name: str) -> bool:
+        """Check if we should skip loading this weight."""
+        return "vision" in name and not self.has_vision
 
     def _transform_weight_name(self, name: str) -> str:
         """Transform weight name by adding language_model prefix if needed."""