sglang 0.5.2rc2__py3-none-any.whl → 0.5.3rc0__py3-none-any.whl

sglang/srt/models/qwen3_next_mtp.py

@@ -0,0 +1,112 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Inference-only Qwen3Next MTP Speculative Decoding."""
+import logging
+from typing import Iterable, Optional, Tuple
+
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+
+from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from sglang.srt.layers.layernorm import GemmaRMSNorm, RMSNorm
+from sglang.srt.layers.logits_processor import LogitsProcessor
+from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.models.qwen3_moe import Qwen3MoeModel
+from sglang.srt.models.qwen3_next import Qwen3NextForCausalLM, Qwen3NextModel
+from sglang.srt.utils import add_prefix
+
+logger = logging.getLogger(__name__)
+
+
+class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+    ) -> None:
+        nn.Module.__init__(self)
+        self.config = config
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.quant_config = quant_config
+        # if not set, model load will be broken in Qwen3NextForCausalLM load_weights()
+        self.pp_group = get_pp_group()
+        # self.determine_num_fused_shared_experts("Qwen3NextForCausalLMMTP")
+
+        # currently based on the provided ckpt, we:
+        # (1) do not use_dedicated_mtp_embeddings provided in ckpt since not provided and directly use the target model embeddings
+        # (2) hardcode bias=False since not provided
+        self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
+        RMSNorm_cls = GemmaRMSNorm
+        self.pre_fc_norm_embedding = RMSNorm_cls(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.pre_fc_norm_hidden = RMSNorm_cls(config.hidden_size, config.rms_norm_eps)
+        config.num_hidden_layers = 1
+        config.full_attention_interval = 1
+        self.model = Qwen3NextModel(
+            config, quant_config, prefix=add_prefix("model", prefix)
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=add_prefix("model.shared_head.head", prefix),
+            use_attn_tp_group=global_server_args_dict["enable_dp_lm_head"],
+        )
+        self.logits_processor = LogitsProcessor(config)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        input_embeds: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        if input_embeds is None:
+            input_embeds = self.model.embed_tokens(input_ids)
+
+        hidden_states = forward_batch.spec_info.hidden_states
+        # Some idle batch has 0 batch size. GemmaRMSNorm.forward would fail due to bs=0.
+        if not forward_batch.forward_mode.is_idle():
+            input_embeds = self.pre_fc_norm_embedding(input_embeds)
+            hidden_states = self.pre_fc_norm_hidden(hidden_states)
+        hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1))
+
+        hidden_states = self.model(
+            input_ids,
+            positions,
+            forward_batch,
+            hidden_states,
+        )
+
+        return self.logits_processor(
+            input_ids, hidden_states, self.lm_head, forward_batch
+        )
+
+    def load_weights(
+        self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False
+    ):
+        super().load_weights(weights, is_mtp=True)
+
+
+EntryClass = [Qwen3NextForCausalLMMTP]

sglang/srt/models/step3_vl.py

@@ -133,7 +133,7 @@ class Step3TextMoEMLP(nn.Module):
             use_grouped_topk=False,
         )
 
-        self.experts = get_moe_impl_class()(
+        self.experts = get_moe_impl_class(quant_config)(
             num_experts=config.moe_num_experts,
             top_k=config.moe_top_k,
             hidden_size=config.hidden_size,

sglang/srt/multimodal/processors/dots_vlm.py

@@ -0,0 +1,99 @@
+import asyncio
+import math
+import re
+from typing import Dict, List, Union
+
+from PIL import Image
+
+from sglang.srt.models.dots_vlm import DotsVLMForCausalLM
+from sglang.srt.multimodal.processors.base_processor import (
+    BaseMultimodalProcessor,
+    MultimodalSpecialTokens,
+)
+from sglang.srt.multimodal.processors.qwen_vl import resize_image_async
+
+
+class DotsVLMImageProcessor(BaseMultimodalProcessor):
+    models = [DotsVLMForCausalLM]
+
+    def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
+        super().__init__(hf_config, server_args, _processor, *args, **kwargs)
+        # The single, pre-expanded image token.
+        self.IMAGE_TOKEN = "<|img|><|imgpad|><|endofimg|>"
+        # The regex that matches expanded image tokens.
+        self.IMAGE_TOKEN_REGEX = re.compile(r"<\|img\|>(?:<\|imgpad\|>)+<\|endofimg\|>")
+
+        assert len(_processor.tokenizer.encode("<|img|>")) == 1
+        self.im_start_id = _processor.tokenizer.encode("<|img|>")[0]
+        self.im_end_id = _processor.tokenizer.encode("<|endofimg|>")[0]
+        self.image_token_id = _processor.tokenizer.encode("<|imgpad|>")[0]
+        self.IM_TOKEN_ID = self.image_token_id
+        self.IM_START_ID = self.im_start_id
+        self.IM_END_ID = self.im_end_id
+
+        vision_config = hf_config.vision_config
+        patch_size = vision_config.patch_size
+        merge_size = vision_config.spatial_merge_size
+
+        self.IMAGE_FACTOR = patch_size * merge_size
+        self.MIN_PIXELS = _processor.image_processor.min_pixels
+        self.MAX_PIXELS = _processor.image_processor.max_pixels
+        self.MAX_RATIO = 200
+        self.mm_tokens = MultimodalSpecialTokens(
+            image_token=self.IMAGE_TOKEN,
+            image_token_id=self.image_token_id,
+            image_token_regex=self.IMAGE_TOKEN_REGEX,
+        ).build(_processor)
+
+    async def process_mm_data_async(
+        self,
+        image_data: List[Union[str, bytes, Dict]],
+        input_text,
+        request_obj,
+        max_req_input_len,
+        *args,
+        **kwargs,
+    ):
+        if isinstance(image_data, str):
+            image_data = [image_data]
+
+        if (
+            isinstance(image_data, list)
+            and image_data
+            and isinstance(image_data[0], list)
+        ):
+            image_data = sum(image_data, [])
+
+        base_output = self.load_mm_data(
+            prompt=input_text,
+            image_data=image_data,
+            multimodal_tokens=self.mm_tokens,
+        )
+
+        # Qwen-specific: resize images if they are raw Image objects
+        if base_output.images and isinstance(base_output.images[0], Image.Image):
+            resize_tasks = [
+                resize_image_async(
+                    image,
+                    min_pixels=self.MIN_PIXELS,
+                    max_pixels=self.MAX_PIXELS,
+                    size_factor=self.IMAGE_FACTOR,
+                )
+                for image in base_output.images
+            ]
+            base_output.images = await asyncio.gather(*resize_tasks)
+
+        combined_mm_item, input_ids, _ = self.process_and_combine_mm_data(
+            base_output, self.mm_tokens
+        )
+
+        if combined_mm_item is None:
+            return None
+
+        return {
+            "input_ids": input_ids.tolist(),
+            "mm_items": combined_mm_item,
+            "im_start_id": self.im_start_id,
+            "im_end_id": self.im_end_id,
+            "im_token_id": self.image_token_id,
+        }

sglang/srt/multimodal/processors/glm4v.py

@@ -2,7 +2,6 @@ import re
 from typing import List, Union
 
 from decord import VideoReader
-from transformers.video_utils import VideoMetadata
 
 from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
 from sglang.srt.models.glm4v import Glm4vForConditionalGeneration
@@ -66,17 +65,18 @@ class Glm4vImageProcessor(SGLangBaseProcessor):
         total_num_frames = len(vr)
         duration = total_num_frames / video_fps if video_fps else 0
 
-        metadata = VideoMetadata(
-            total_num_frames=int(total_num_frames),
-            fps=float(video_fps),
-            duration=float(duration),
-            video_backend="decord",
-        )
-
         # Extract all frames
         indices = list(range(total_num_frames))
         frames = vr.get_batch(indices).asnumpy()
-        metadata.frames_indices = indices
+
+        # Return metadata as dict so transformers can properly create VideoMetadata objects
+        metadata = {
+            "total_num_frames": int(total_num_frames),
+            "fps": float(video_fps),
+            "duration": float(duration),
+            "video_backend": "decord",
+            "frames_indices": indices,
+        }
 
         return frames, metadata
 

sglang/srt/multimodal/processors/internvl.py

@@ -2,8 +2,10 @@
 
 import numpy as np
 import torch
-from decord import VideoReader, cpu
+import torchvision.transforms as T
+from decord import VideoReader, cpu, gpu
 from PIL import Image
+from torchvision.transforms import InterpolationMode
 
 from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
 from sglang.srt.models.interns1 import InternS1ForConditionalGeneration
@@ -48,99 +50,6 @@ class InternVLImageProcessor(BaseMultimodalProcessor):
             image_token_id=tokenizer.convert_tokens_to_ids(self.IMG_CONTEXT_TOKEN),
         ).build(_image_processor)
 
-    @staticmethod
-    def build_transform(input_size):
-        IMAGENET_MEAN = (0.485, 0.456, 0.406)
-        IMAGENET_STD = (0.229, 0.224, 0.225)
-
-        def resize_image(img, size):
-            return img.resize((size, size), Image.Resampling.BICUBIC)
-
-        def to_tensor(img):
-            # Convert PIL Image to numpy array
-            img_array = np.array(img).astype(np.float32) / 255.0
-            # Convert HWC to CHW format
-            img_array = img_array.transpose(2, 0, 1)
-            return torch.from_numpy(img_array)
-
-        def normalize(tensor, mean, std):
-            mean = torch.tensor(mean).view(-1, 1, 1)
-            std = torch.tensor(std).view(-1, 1, 1)
-            return (tensor - mean) / std
-
-        def transform(img):
-            img = img.convert("RGB") if img.mode != "RGB" else img
-            img = resize_image(img, input_size)
-            tensor = to_tensor(img)
-            tensor = normalize(tensor, IMAGENET_MEAN, IMAGENET_STD)
-            return tensor
-
-        return transform
-
-    @staticmethod
-    def dynamic_preprocess(
-        image, min_num=1, max_num=12, image_size=448, use_thumbnail=False
-    ):
-
-        def find_closest_aspect_ratio(
-            aspect_ratio, target_ratios, width, height, image_size
-        ):
-            best_ratio_diff = float("inf")
-            best_ratio = (1, 1)
-            area = width * height
-            for ratio in target_ratios:
-                target_aspect_ratio = ratio[0] / ratio[1]
-                ratio_diff = abs(aspect_ratio - target_aspect_ratio)
-                if ratio_diff < best_ratio_diff:
-                    best_ratio_diff = ratio_diff
-                    best_ratio = ratio
-                elif ratio_diff == best_ratio_diff:
-                    if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
-                        best_ratio = ratio
-            return best_ratio
-
-        orig_width, orig_height = image.size
-        aspect_ratio = orig_width / orig_height
-
-        # calculate the existing image aspect ratio
-        target_ratios = set(
-            (i, j)
-            for n in range(min_num, max_num + 1)
-            for i in range(1, n + 1)
-            for j in range(1, n + 1)
-            if i * j <= max_num and i * j >= min_num
-        )
-        target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
-
-        # find the closest aspect ratio to the target
-        target_aspect_ratio = find_closest_aspect_ratio(
-            aspect_ratio, target_ratios, orig_width, orig_height, image_size
-        )
-
-        # calculate the target width and height
-        target_width = image_size * target_aspect_ratio[0]
-        target_height = image_size * target_aspect_ratio[1]
-        blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
-
-        # resize the image
-        resized_img = image.resize((target_width, target_height))
-        processed_images = []
-        for i in range(blocks):
-            box = (
-                (i % (target_width // image_size)) * image_size,
-                (i // (target_width // image_size)) * image_size,
-                ((i % (target_width // image_size)) + 1) * image_size,
-                ((i // (target_width // image_size)) + 1) * image_size,
-            )
-            # split the image
-            split_img = resized_img.crop(box)
-            processed_images.append(split_img)
-        assert len(processed_images) == blocks
-        if use_thumbnail and len(processed_images) != 1:
-            thumbnail_img = image.resize((image_size, image_size))
-            processed_images.append(thumbnail_img)
-        return processed_images
-
     @staticmethod
     def get_index(bound, fps, max_frame, first_idx=0, num_segments=32):
         if bound:
@@ -160,27 +69,112 @@ class InternVLImageProcessor(BaseMultimodalProcessor):
 
     @staticmethod
     def load_video(video_path, bound=None, input_size=448, max_num=1, num_segments=32):
-        vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
+        try:
+            vr = VideoReader(video_path, ctx=gpu(0), num_threads=1)
+            use_gpu = True
+        except (RuntimeError, OSError) as e:
+            print(
+                f"[WARNING] Load video on gpu decoding failed: {e}. Falling back to CPU."
+            )
+            vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
+            use_gpu = False
+
         max_frame = len(vr) - 1
         fps = float(vr.get_avg_fps())
 
-        pixel_values_list, num_patches_list = [], []
-        transform = InternVLImageProcessor.build_transform(input_size=input_size)
+        pixel_values_list = []
+        num_patches_list = []
         frame_indices = InternVLImageProcessor.get_index(
             bound, fps, max_frame, first_idx=0, num_segments=num_segments
         )
+
         for frame_index in frame_indices:
-            img = Image.fromarray(vr[frame_index].asnumpy()).convert("RGB")
-            img = InternVLImageProcessor.dynamic_preprocess(
-                img, image_size=input_size, use_thumbnail=True, max_num=max_num
+            # Load frame
+            frame = vr[frame_index]
+            if use_gpu:
+                img = frame.cuda().permute(2, 0, 1).float() / 255.0
+            else:
+                img_np = frame.asnumpy()
+                img = torch.from_numpy(img_np).permute(2, 0, 1).cuda().float() / 255.0
+
+            # Using the mean and variance of the ImageNet dataset for all input images can lead to accuracy issues, while using the mean and variance of each input image is a more accurate choice.
+            mean = img.mean(dim=[1, 2], keepdim=True)
+            # Prevent division by zero; clamp to minimum value of 1e-6
+            std = img.std(dim=[1, 2], keepdim=True).clamp(min=1e-6)
+            img = (img - mean) / std
+
+            tiles = InternVLImageProcessor.dynamic_preprocess(
+                img, image_size=input_size, max_num=max_num, use_thumbnail=True
             )
-            pixel_values = [transform(tile) for tile in img]
-            pixel_values = torch.stack(pixel_values)
-            num_patches_list.append(pixel_values.shape[0])
-            pixel_values_list.append(pixel_values)
-        pixel_values = torch.cat(pixel_values_list)
+
+            pixel_values_list.append(tiles)
+            num_patches_list.append(tiles.shape[0])
+
+        pixel_values = torch.cat(pixel_values_list, dim=0)
         return pixel_values, num_patches_list
 
+    @staticmethod
+    def dynamic_preprocess(tensor, image_size=448, max_num=12, use_thumbnail=False):
+        C, H, W = tensor.shape
+        aspect_ratio = W / H
+
+        # Generate all possible aspect ratios
+        target_ratios = set(
+            (i, j)
+            for n in range(1, max_num + 1)
+            for i in range(1, n + 1)
+            for j in range(1, n + 1)
+            if i * j <= max_num
+        )
+        target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+        # Find closest ratio
+        best_ratio_diff = float("inf")
+        best_ratio = (1, 1)
+
+        for x, y in target_ratios:
+            target_ar = x / y
+            diff = abs(aspect_ratio - target_ar)
+            blocks = x * y
+            best_blocks = best_ratio[0] * best_ratio[1]
+
+            if diff < best_ratio_diff:
+                best_ratio_diff = diff
+                best_ratio = (x, y)
+            elif diff == best_ratio_diff and blocks > best_blocks:
+                best_ratio = (x, y)
+
+        target_w, target_h = image_size * best_ratio[0], image_size * best_ratio[1]
+        blocks = best_ratio[0] * best_ratio[1]
+
+        # Resize on GPU
+        resized = torch.nn.functional.interpolate(
+            tensor.unsqueeze(0),
+            size=(target_h, target_w),
+            mode="bicubic",
+            align_corners=False,
+        ).squeeze(0)
+
+        # Split into tiles
+        tiles = []
+        for i in range(blocks):
+            x = (i % best_ratio[0]) * image_size
+            y = (i // best_ratio[0]) * image_size
+            tile = resized[:, y : y + image_size, x : x + image_size]
+            tiles.append(tile)
+
+        # Add thumbnail if needed
+        if use_thumbnail and len(tiles) > 1:
+            thumb = torch.nn.functional.interpolate(
+                tensor.unsqueeze(0),
+                size=(image_size, image_size),
+                mode="bicubic",
+                align_corners=False,
+            ).squeeze(0)
+            tiles.append(thumb)
+
+        return torch.stack(tiles).to(torch.bfloat16)
+
     async def process_mm_data_async(
         self, image_data, input_text, request_obj, **kwargs
     ):
@@ -191,53 +185,71 @@ class InternVLImageProcessor(BaseMultimodalProcessor):
             discard_alpha_channel=True,
         )
 
-        def process_image_internvl(image, input_size=448, max_num=12):
-            transform = InternVLImageProcessor.build_transform(input_size=input_size)
-            images = InternVLImageProcessor.dynamic_preprocess(
-                image, image_size=input_size, use_thumbnail=True, max_num=max_num
-            )
-            pixel_values = [transform(image) for image in images]
-            pixel_values = torch.stack(pixel_values)
-            return pixel_values
-
         num_patches_list = []
         pixel_values = []
+
         # Process each input with allocated frames
-        for image_index, (image) in enumerate(base_output.images):
+        for image_index, image in enumerate(base_output.images):
             try:
                 # TODO: video input
-                raw_image = process_image_internvl(image)
-                pixel_value = [raw_image.to(torch.bfloat16)]
-                pixel_values += pixel_value
-                num_patches = raw_image.shape[0]
-                num_patches_list += [num_patches]
-
-            except FileNotFoundError as e:
-                print(e)
+                # Convert PIL to GPU tensor
+                if isinstance(image, Image.Image):
+                    img_np = np.array(image.convert("RGB"))
+                    tensor = (
+                        torch.from_numpy(img_np).permute(2, 0, 1).cuda().float() / 255.0
+                    )
+                else:
+                    tensor = image.cuda()  # assume already tensor
+
+                # Using the mean and variance of the ImageNet dataset for all input images can lead to accuracy issues, while using the mean and variance of each input image is a more accurate choice.
+                mean = tensor.mean(dim=[1, 2], keepdim=True)
+                # Prevent division by zero; clamp to minimum value of 1e-6
+                std = tensor.std(dim=[1, 2], keepdim=True).clamp(min=1e-6)
+                tensor = (tensor - mean) / std
+                tiles = self.dynamic_preprocess(
+                    tensor, image_size=448, max_num=12, use_thumbnail=True
+                )
+
+                pixel_values.append(tiles)
+                num_patches_list.append(tiles.shape[0])
+
+            except Exception as e:
+                print(f"[Error] Failed to process image {image_index}: {e}")
                 return None
 
+        # Concatenate all
         pixel_values = torch.cat(pixel_values, dim=0)
 
         original_placeholder = "<<<__IMG_CONTEXT_PLACEHOLDER__>>>"
         input_text = input_text.replace(self.IMG_CONTEXT_TOKEN, original_placeholder)
 
-        for idx, num_patches in enumerate(num_patches_list):
+        input_text_updated = input_text
+        for num_patches in num_patches_list:
             image_tokens = (
                 self.IMG_START_TOKEN
                 + self.IMG_CONTEXT_TOKEN * self.num_image_token * num_patches
                 + self.IMG_END_TOKEN
             )
-            input_text = input_text.replace(original_placeholder, image_tokens, 1)
+            input_text_updated = input_text_updated.replace(
+                original_placeholder, image_tokens, 1
+            )
 
-        input_text = input_text.replace(original_placeholder, self.IMG_CONTEXT_TOKEN)
+        input_text_updated = input_text_updated.replace(
+            original_placeholder, self.IMG_CONTEXT_TOKEN
+        )
 
-        input_ids = self.tokenizer(input_text, return_tensors="pt")[
+        # Tokenize
+        input_ids_tensor = self.tokenizer(input_text_updated, return_tensors="pt")[
             "input_ids"
         ].flatten()
+        input_ids = input_ids_tensor.tolist()
+
+        # Get image token offsets
         image_offsets = self.get_mm_items_offset(
-            input_ids=input_ids,
+            input_ids=input_ids_tensor.to("cuda"),
             mm_token_id=self.mm_tokens.image_token_id,
         )
+
         items = [
             MultimodalDataItem(
                 feature=pixel_values,
@@ -247,7 +259,7 @@ class InternVLImageProcessor(BaseMultimodalProcessor):
         ]
 
         return {
-            "input_ids": input_ids.tolist(),
+            "input_ids": input_ids,
             "mm_items": items,
             "im_start_id": self.img_start_token_id,
             "im_end_id": self.img_end_token_id,

sglang/srt/multimodal/processors/qwen_vl.py

@@ -67,10 +67,15 @@ def smart_resize(
     return h_bar, w_bar
 
 
-def resize_image(image, size_factor: int = IMAGE_FACTOR) -> Image.Image:
+def resize_image(
+    image,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+    size_factor: int = IMAGE_FACTOR,
+) -> Image.Image:
     width, height = image.size
-    min_pixels = MIN_PIXELS
-    max_pixels = MAX_PIXELS
+    min_pixels = min_pixels
+    max_pixels = max_pixels
     resized_height, resized_width = smart_resize(
         height,
         width,
@@ -97,8 +102,13 @@ def floor_by_factor(number: int, factor: int) -> int:
     return math.floor(number / factor) * factor
 
 
-async def resize_image_async(image):
-    return resize_image(image)
+async def resize_image_async(
+    image,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+    size_factor: int = IMAGE_FACTOR,
+):
+    return resize_image(image, min_pixels, max_pixels, size_factor)
 
 
 def smart_nframes(

sglang/srt/offloader.py

@@ -38,6 +38,10 @@ class BaseOffloader(ABC):
     def post_init(self):
         pass
 
+    @property
+    def forbid_copy_engine_usage(self):
+        return False
+
 
 class NoopOffloader(BaseOffloader):
     pass
@@ -233,6 +237,10 @@ class OffloaderV2(BaseOffloader):
         for i in range(self.prefetch_step):
             self.offloaders[i].start_onload()
 
+    @property
+    def forbid_copy_engine_usage(self):
+        return self.mode == "cpu"
+
 
 def _hook_module_forward_for_offloader(index, module, offloaders, prefetch_step):
     def _on_forward_end():
@@ -398,14 +406,30 @@ class _ShmCpuParamOffloader(_BaseParamOffloader):
         return self.shm_cpu_data.to("cuda", non_blocking=True)
 
 
+def update_param(param, new_tensor):
+    """Update parameter while keeping properties needed by Offloader (e.g. pinned host memory)."""
+
+    if param.device == new_tensor.device:
+        param.data = new_tensor
+    else:
+        assert param.device == torch.device(
+            "cpu"
+        ), f"{param.device=} {new_tensor.device=}"
+        param.data = _create_cpu_data(new_tensor, pin_memory=True)
+
+
 def _move_param_to_cpu(param, pin_memory: bool):
+    param.data = _create_cpu_data(param.data, pin_memory=pin_memory)
+
+
+def _create_cpu_data(data, pin_memory: bool):
     cpu_data = _empty_strided_like(
-        param.data,
+        data,
         device="cpu",
         pin_memory=pin_memory,
     )
-    cpu_data.copy_(param.data)
-    param.data = cpu_data
+    cpu_data.copy_(data)
+    return cpu_data
 
 
 def _move_param_to_meta(module, param_name):