sglang 0.5.4__py3-none-any.whl → 0.5.4.post2__py3-none-any.whl

sglang/srt/models/nvila.py

@@ -0,0 +1,355 @@
+import itertools
+import math
+from collections.abc import Iterable
+from typing import Any
+
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+from transformers.models.siglip import SiglipVisionConfig, SiglipVisionModel
+
+import sglang.srt.managers.mm_utils as mm_utils
+import sglang.srt.model_loader.weight_utils as weight_utils
+import sglang.srt.utils as utils
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
+from sglang.srt.managers.schedule_batch import (
+    Modality,
+    MultimodalDataItem,
+    MultimodalInputs,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.models.qwen2 import Qwen2ForCausalLM
+
+MM_HIDDEN_SIZE = 3456
+
+
+class NVILAConfig(PretrainedConfig):
+    model_type = "nvila"
+    sub_configs = {
+        "text_config": Qwen2Config,
+        "vision_config": SiglipVisionConfig,
+    }
+    _auto_class = "AutoConfig"
+
+    def __init__(
+        self,
+        *,
+        text_config: dict[str, Any] | None = None,
+        vision_config: dict[str, Any] | None = None,
+        image_token_id: int | None = None,
+        video_token_id: int | None = None,
+        **kwargs,
+    ):
+        self.text_config = (
+            Qwen2Config(**text_config) if text_config is not None else Qwen2Config()
+        )
+        self.vision_config = (
+            SiglipVisionConfig(**vision_config)
+            if vision_config is not None
+            else SiglipVisionConfig()
+        )
+
+        self.image_token_id = image_token_id if image_token_id is not None else -1
+        self.video_token_id = video_token_id if video_token_id is not None else -1
+
+        super().__init__(**kwargs)
+
+
+class NVILAMultiModalProjectorDownsampleBlock(nn.Module):
+    def forward(self, x: Tensor) -> Tensor:
+        batch_size, sequence_length, hidden_size = x.shape
+
+        feat_size = math.isqrt(sequence_length)
+
+        features = x.reshape(batch_size, feat_size, feat_size, hidden_size)
+
+        pad_after = feat_size % 2
+        if pad_after > 0:
+            features = F.pad(features, (0, 0, 0, pad_after, 0, pad_after))
+            feat_size = feat_size + pad_after
+
+        features = features.reshape(
+            batch_size, feat_size // 2, 2, feat_size // 2, 2, hidden_size
+        )
+        features = features.permute(0, 1, 3, 2, 4, 5).contiguous()
+        features = features.reshape(batch_size, -1, 4 * hidden_size)
+
+        return features
+
+
+class NVILAMultiModalProjector(nn.Module):
+    def __init__(self, config: NVILAConfig):
+        super().__init__()
+
+        self.layers = nn.Sequential(
+            NVILAMultiModalProjectorDownsampleBlock(),
+            nn.LayerNorm(MM_HIDDEN_SIZE * 4),
+            nn.Linear(MM_HIDDEN_SIZE * 4, config.text_config.hidden_size),
+            nn.GELU(),
+            nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.layers(x)
+
+
+class NVILAForConditionalGeneration(nn.Module):
+    def __init__(
+        self,
+        config: NVILAConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+
+        self.vision_tower = SiglipVisionModel(config.vision_config)
+        self.mm_projector = NVILAMultiModalProjector(config)
+        self.llm = Qwen2ForCausalLM(
+            config=config.text_config,
+            quant_config=quant_config,
+            prefix=utils.add_prefix("llm", prefix),
+        )
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        positions: Tensor,
+        forward_batch: ForwardBatch,
+        get_embedding: bool = False,
+    ) -> LogitsProcessorOutput:
+        output = mm_utils.general_mm_embed_routine(
+            input_ids=input_ids,
+            forward_batch=forward_batch,
+            language_model=self.llm,
+            data_embedding_funcs={
+                Modality.IMAGE: self.get_image_feature,
+                Modality.VIDEO: self.get_image_feature,
+            },
+            get_embedding=get_embedding,
+            positions=positions,
+        )
+
+        assert isinstance(output, LogitsProcessorOutput)
+
+        return output
+
+    def get_image_feature(self, mm_input: list[MultimodalDataItem]) -> Tensor:
+        block_sizes = (
+            list(
+                itertools.chain.from_iterable(
+                    x.block_sizes for x in mm_input if hasattr(x, "block_sizes")
+                )
+            )
+            or None
+        )
+        pixel_values = torch.cat([torch.tensor(x.feature) for x in mm_input], dim=0)
+
+        vision_tower_output: BaseModelOutputWithPooling = self.vision_tower(
+            pixel_values.to(
+                device=self.vision_tower.device, dtype=self.vision_tower.dtype
+            ),
+            output_hidden_states=True,
+        )
+        assert vision_tower_output.hidden_states is not None
+
+        vision_features: Tensor = vision_tower_output.hidden_states[-2]
+
+        vision_features_list, block_sizes = merge_features_for_dynamic_s2(
+            vision_features,
+            block_sizes=(
+                block_sizes
+                if block_sizes is not None
+                else [None] * vision_features.shape[0]
+            ),
+            resize_output_to_scale_idx=-1,
+            scales=[448, 896, 1344],
+        )
+
+        vision_features_list = [
+            split_chessboard(x, block_size[0], block_size[1])
+            for x, block_size in zip(vision_features_list, block_sizes)
+        ]
+
+        vision_features = torch.cat(
+            [einops.rearrange(x, "b c h w -> b (h w) c") for x in vision_features_list]
+        )
+
+        vision_features = self.mm_projector(vision_features)
+
+        vision_features_list = list(
+            vision_features.split(
+                [block_size[0] * block_size[1] for block_size in block_sizes], dim=0
+            )
+        )
+        vision_features_list = [
+            merge_chessboard(x, block_size[0], block_size[1])
+            for x, block_size in zip(vision_features_list, block_sizes)
+        ]
+
+        vision_features = torch.stack(
+            [einops.rearrange(x, "1 c h w -> (h w) c") for x in vision_features_list]
+        )
+
+        vision_features = einops.rearrange(vision_features, "n p d -> (n p) d")
+
+        return vision_features
+
+    def load_weights(self, weights: Iterable[tuple[str, Tensor]]) -> None:
+        params_dict = dict(self.named_parameters())
+
+        for name, loaded_weight in weights:
+            if name.startswith("llm."):
+                self.llm.load_weights([(name[len("llm.") :], loaded_weight)])
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(
+                    param, "weight_loader", weight_utils.default_weight_loader
+                )
+                weight_loader(param, loaded_weight)
+
+    def pad_input_ids(
+        self, input_ids: list[int], mm_inputs: MultimodalInputs
+    ) -> list[int]:
+        pattern = MultiModalityDataPaddingPatternMultimodalTokens()
+        return pattern.pad_input_tokens(input_ids, mm_inputs)
+
+
+def merge_chessboard(x, num_split_h, num_split_w):
+    """
+    x: b * n * c or b * h * w * c
+    out: b * c * h * w
+    Assuming x contains num_split**2 sub-squares concatenated along batch dimension, merge the sub-squares back to the original whole square.
+    """
+    B = x.shape[0]
+    if x.dim() == 3:
+        N = x.shape[1]
+        x = einops.rearrange(
+            x, "b (h w) c -> b c h w", h=math.isqrt(N), w=math.isqrt(N)
+        )
+
+    assert B % (num_split_h * num_split_w) == 0
+    b = B // (num_split_h * num_split_w)
+
+    x_merge = torch.cat(
+        [
+            torch.cat(
+                [
+                    x[(i * num_split_w + j) * b : (i * num_split_w + j + 1) * b]
+                    for j in range(num_split_w)
+                ],
+                dim=-1,
+            )
+            for i in range(num_split_h)
+        ],
+        dim=-2,
+    )
+
+    return x_merge
+
+
+def merge_features_for_dynamic_s2(
+    image_features, block_sizes, *, scales, resize_output_to_scale_idx
+):
+    image_features_each_image = []
+    new_block_sizes = []
+    block_cnt = 0
+    for block_size_each_image in block_sizes:
+        if block_size_each_image is None:
+            cur_features = image_features[block_cnt : block_cnt + 1]
+            cur_features = einops.rearrange(
+                cur_features,
+                "1 (h w) c -> 1 c h w",
+                h=math.isqrt(cur_features.shape[1]),
+            )
+            cur_features = cur_features.repeat(1, len(scales), 1, 1)
+            image_features_each_image.append(cur_features)
+            new_block_sizes.append((1, 1))
+            block_cnt += 1
+        else:
+            cur_features_each_scale = []
+            for scale in scales[:-1]:
+                num_blocks_this_scale = (scale // scales[0]) ** 2
+                cur_features_each_scale.append(
+                    merge_chessboard(
+                        image_features[block_cnt : block_cnt + num_blocks_this_scale],
+                        num_split_h=scale // scales[0],
+                        num_split_w=scale // scales[0],
+                    )
+                )  # 1 * C * H * W
+                block_cnt += num_blocks_this_scale
+            num_blocks_last_scale = block_size_each_image[0] * block_size_each_image[1]
+            cur_features_each_scale.append(
+                merge_chessboard(
+                    image_features[block_cnt : block_cnt + num_blocks_last_scale],
+                    num_split_h=block_size_each_image[0],
+                    num_split_w=block_size_each_image[1],
+                )
+            )  # 1 * C * H * W
+            block_cnt += num_blocks_last_scale
+
+            # resize and concat features from different scales
+            output_size = cur_features_each_scale[resize_output_to_scale_idx].shape[-2:]
+            cur_features = torch.cat(
+                [
+                    F.interpolate(
+                        cur_features_each_scale[i].to(torch.float32),
+                        size=output_size,
+                        mode="area",
+                    ).to(cur_features_each_scale[i].dtype)
+                    for i in range(len(cur_features_each_scale))
+                ],
+                dim=1,
+            )
+
+            image_features_each_image.append(cur_features)
+
+            if (
+                resize_output_to_scale_idx == len(scales) - 1
+                or resize_output_to_scale_idx == -1
+            ):
+                new_block_sizes.append(block_size_each_image)
+            else:
+                new_block_sizes.append(
+                    (
+                        scales[resize_output_to_scale_idx] // scales[0],
+                        scales[resize_output_to_scale_idx] // scales[0],
+                    )
+                )
+
+    assert block_cnt == len(
+        image_features
+    ), f"The number of blocks ({block_cnt}) does not match length of image_features ({len(image_features)})!"
+
+    return image_features_each_image, new_block_sizes
+
+
+def split_chessboard(x, num_split_h, num_split_w):
+    """
+    x: b * c * h * w
+    out: b * c * h * w
+    Deividing x into num_split**2 sub-squares, and concatenate all the sub-squares on the batch dimension
+    """
+    B, C, H, W = x.shape
+    assert H % num_split_h == 0 and W % num_split_w == 0
+    h, w = H // num_split_h, W // num_split_w
+    x_split = torch.cat(
+        [
+            x[:, :, i * h : (i + 1) * h, j * w : (j + 1) * w]
+            for i in range(num_split_h)
+            for j in range(num_split_w)
+        ],
+        dim=0,
+    )
+    return x_split
+
+
+EntryClass = [NVILAForConditionalGeneration]

sglang/srt/models/nvila_lite.py

@@ -0,0 +1,184 @@
+import math
+from collections.abc import Iterable
+from typing import Any
+
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+from transformers.models.siglip import SiglipVisionConfig, SiglipVisionModel
+
+import sglang.srt.managers.mm_utils as mm_utils
+import sglang.srt.model_loader.weight_utils as weight_utils
+import sglang.srt.utils as utils
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
+from sglang.srt.managers.schedule_batch import (
+    Modality,
+    MultimodalDataItem,
+    MultimodalInputs,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.models.qwen2 import Qwen2ForCausalLM
+
+MM_HIDDEN_SIZE = 1152
+
+
+class NVILALiteConfig(PretrainedConfig):
+    model_type = "nvila_lite"
+    sub_configs = {
+        "text_config": Qwen2Config,
+        "vision_config": SiglipVisionConfig,
+    }
+    _auto_class = "AutoConfig"
+
+    def __init__(
+        self,
+        *,
+        text_config: dict[str, Any] | None = None,
+        vision_config: dict[str, Any] | None = None,
+        image_token_id: int | None = None,
+        video_token_id: int | None = None,
+        **kwargs,
+    ):
+        self.text_config = (
+            Qwen2Config(**text_config) if text_config is not None else Qwen2Config()
+        )
+        self.vision_config = (
+            SiglipVisionConfig(**vision_config)
+            if vision_config is not None
+            else SiglipVisionConfig()
+        )
+
+        self.image_token_id = image_token_id if image_token_id is not None else -1
+        self.video_token_id = video_token_id if video_token_id is not None else -1
+
+        super().__init__(**kwargs)
+
+
+class NVILALiteMultiModalProjectorDownsampleBlock(nn.Module):
+    def forward(self, x: Tensor) -> Tensor:
+        batch_size, sequence_length, hidden_size = x.shape
+
+        feat_size = math.isqrt(sequence_length)
+
+        features = x.reshape(batch_size, feat_size, feat_size, hidden_size)
+
+        pad_after = (3 - feat_size % 3) % 3
+        if pad_after > 0:
+            features = F.pad(features, (0, 0, 0, pad_after, 0, pad_after))
+            feat_size = feat_size + pad_after
+
+        features = features.reshape(
+            batch_size, feat_size // 3, 3, feat_size // 3, 3, hidden_size
+        )
+        features = features.permute(0, 1, 3, 2, 4, 5).contiguous()
+        features = features.reshape(batch_size, -1, 9 * hidden_size)
+
+        return features
+
+
+class NVILALiteMultiModalProjector(nn.Module):
+    def __init__(self, config: NVILALiteConfig):
+        super().__init__()
+
+        self.layers = nn.Sequential(
+            NVILALiteMultiModalProjectorDownsampleBlock(),
+            nn.LayerNorm(MM_HIDDEN_SIZE * 9),
+            nn.Linear(MM_HIDDEN_SIZE * 9, MM_HIDDEN_SIZE * 3),
+            nn.GELU(),
+            nn.LayerNorm(MM_HIDDEN_SIZE * 3),
+            nn.Linear(MM_HIDDEN_SIZE * 3, config.text_config.hidden_size),
+            nn.GELU(),
+            nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.layers(x)
+
+
+class NVILALiteForConditionalGeneration(nn.Module):
+    def __init__(
+        self,
+        config: NVILALiteConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+
+        self.vision_tower = SiglipVisionModel(config.vision_config)
+        self.mm_projector = NVILALiteMultiModalProjector(config)
+        self.llm = Qwen2ForCausalLM(
+            config=config.text_config,
+            quant_config=quant_config,
+            prefix=utils.add_prefix("llm", prefix),
+        )
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        positions: Tensor,
+        forward_batch: ForwardBatch,
+        get_embedding: bool = False,
+    ) -> LogitsProcessorOutput:
+        output = mm_utils.general_mm_embed_routine(
+            input_ids=input_ids,
+            forward_batch=forward_batch,
+            language_model=self.llm,
+            data_embedding_funcs={
+                Modality.IMAGE: self.get_image_feature,
+                Modality.VIDEO: self.get_image_feature,
+            },
+            get_embedding=get_embedding,
+            positions=positions,
+        )
+
+        assert isinstance(output, LogitsProcessorOutput)
+
+        return output
+
+    def get_image_feature(self, mm_input: list[MultimodalDataItem]) -> Tensor:
+        pixel_values = torch.cat([torch.tensor(x.feature) for x in mm_input], dim=0)
+
+        vision_tower_output: BaseModelOutputWithPooling = self.vision_tower(
+            pixel_values,
+            output_hidden_states=True,
+        )
+        assert vision_tower_output.hidden_states is not None
+
+        vision_features = vision_tower_output.hidden_states[-2]
+
+        vision_features = self.mm_projector(vision_features)
+
+        vision_features = einops.rearrange(vision_features, "n p d -> (n p) d")
+
+        return vision_features
+
+    def load_weights(self, weights: Iterable[tuple[str, Tensor]]) -> None:
+        params_dict = dict(self.named_parameters())
+
+        for name, loaded_weight in weights:
+            if name.startswith("llm."):
+                self.llm.load_weights([(name[len("llm.") :], loaded_weight)])
+            else:
+                param = params_dict[name]
+                weight_loader = getattr(
+                    param, "weight_loader", weight_utils.default_weight_loader
+                )
+                weight_loader(param, loaded_weight)
+
+    def pad_input_ids(
+        self, input_ids: list[int], mm_inputs: MultimodalInputs
+    ) -> list[int]:
+        pattern = MultiModalityDataPaddingPatternMultimodalTokens()
+        return pattern.pad_input_tokens(input_ids, mm_inputs)
+
+
+EntryClass = [NVILALiteForConditionalGeneration]

sglang/srt/models/qwen2.py

@@ -49,6 +49,7 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     kv_cache_scales_loader,
 )
+from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix, make_layers
 
 Qwen2Config = None
@@ -89,6 +90,9 @@ class Qwen2MLP(nn.Module):
         self.act_fn = SiluAndMul()
 
     def forward(self, x):
+        if get_global_server_args().rl_on_policy_target == "fsdp":
+            x = x.bfloat16()
+
         gate_up, _ = self.gate_up_proj(x)
         x = self.act_fn(gate_up)
         x, _ = self.down_proj(x)
@@ -275,6 +279,11 @@ class Qwen2Model(nn.Module):
                 quant_config=quant_config,
                 enable_tp=not is_dp_attention_enabled(),
                 prefix=add_prefix("embed_tokens", prefix),
+                params_dtype=(
+                    torch.float32
+                    if get_global_server_args().rl_on_policy_target == "fsdp"
+                    else None
+                ),
             )
         else:
             self.embed_tokens = PPMissingLayer()
@@ -295,7 +304,19 @@ class Qwen2Model(nn.Module):
             prefix=add_prefix("layers", prefix),
         )
         if self.pp_group.is_last_rank:
-            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+            norm_kwargs = (
+                dict(
+                    weight_dtype=torch.float32,
+                    cast_x_before_out_mul=True,
+                    override_orig_dtype=torch.float32,
+                    fp32_residual=True,
+                )
+                if get_global_server_args().rl_on_policy_target == "fsdp"
+                else {}
+            )
+            self.norm = RMSNorm(
+                config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
+            )
         else:
             self.norm = PPMissingLayer(return_tuple=True)
 
@@ -441,7 +462,7 @@ class Qwen2ForCausalLM(nn.Module):
                 self.pp_group.send(
                     self.model.embed_tokens.weight, dst=self.pp_group.last_rank
                 )
-            else:
+            elif self.pp_group.is_last_rank:
                 emb_token_weight = self.pp_group.recv(
                     size=(config.vocab_size, config.hidden_size),
                     dtype=next(self.model.parameters()).dtype,

sglang/srt/models/qwen2_moe.py

@@ -473,10 +473,16 @@ class Qwen2MoeDecoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
         residual: Optional[torch.Tensor],
+        captured_last_layer_outputs: Optional[List[torch.Tensor]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
 
-        hidden_states, residual = self.layer_communicator.prepare_attn(
-            hidden_states, residual, forward_batch
+        hidden_states, residual = (
+            self.layer_communicator.prepare_attn_and_capture_last_layer_outputs(
+                hidden_states,
+                residual,
+                forward_batch,
+                captured_last_layer_outputs=captured_last_layer_outputs,
+            )
         )
 
         if hidden_states.shape[0] != 0:
@@ -553,6 +559,11 @@ class Qwen2MoeModel(nn.Module):
         # For EAGLE3 support
         self.layers_to_capture = []
 
+    def set_eagle3_layers_to_capture(self, layers_to_capture: List[int]):
+        self.layers_to_capture = layers_to_capture
+        for layer_id in self.layers_to_capture:
+            setattr(self.layers[layer_id], "_is_layer_to_capture", True)
+
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -585,12 +596,6 @@ class Qwen2MoeModel(nn.Module):
             )
         else:
             for i in range(self.start_layer, self.end_layer):
-                if i in self.layers_to_capture:
-                    aux_hidden_states.append(
-                        hidden_states + residual
-                        if residual is not None
-                        else hidden_states
-                    )
                 ctx = (
                     nullcontext()
                     if get_global_server_args().enable_piecewise_cuda_graph
@@ -599,7 +604,15 @@ class Qwen2MoeModel(nn.Module):
                 with ctx:
                     layer = self.layers[i]
                     hidden_states, residual = layer(
-                        positions, hidden_states, forward_batch, residual
+                        positions,
+                        hidden_states,
+                        forward_batch,
+                        residual,
+                        captured_last_layer_outputs=(
+                            aux_hidden_states
+                            if getattr(layer, "_is_layer_to_capture", False)
+                            else None
+                        ),
                     )
         if not self.pp_group.is_last_rank:
             return PPProxyTensors(
@@ -830,13 +843,15 @@ class Qwen2MoeForCausalLM(nn.Module):
         self.capture_aux_hidden_states = True
         if layer_ids is None:
             num_layers = self.config.num_hidden_layers
-            self.model.layers_to_capture = [
-                2,
-                num_layers // 2,
-                num_layers - 3,
-            ]  # Specific layers for EAGLE3 support
+            self.model.set_eagle3_layers_to_capture(
+                [
+                    2,
+                    num_layers // 2,
+                    num_layers - 3,
+                ]
+            )  # Specific layers for EAGLE3 support
         else:
-            self.model.layers_to_capture = [val + 1 for val in layer_ids]
+            self.model.set_eagle3_layers_to_capture([val + 1 for val in layer_ids])
 
 
 EntryClass = Qwen2MoeForCausalLM