[quantization] share bounded post-load device staging (#35180)

python/sglang/srt/model_loader/loader.py

@@ -42,6 +42,7 @@ import numpy as np
 import torch
 
 from sglang.srt.constants import GIB_BYTES
+from sglang.srt.model_loader.post_load import stage_module_for_post_load
 from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
     RemoteInstanceWeightLoaderBackend,
     get_remote_instance_transfer_engine_info_per_rank,
@@ -150,62 +151,16 @@ logger = logging.getLogger(__name__)
 @contextmanager
 def device_loading_context(module: torch.nn.Module, target_device: torch.device):
     if target_device.type == "cpu":
-        # If target is CPU, no need to move anything
         yield module
         return
 
-    original_infos: Dict[str, Dict] = {}
+    with stage_module_for_post_load(
-
+        module,
-    # Store original device states and move parameters to GPU if they're on CPU
+        target_device,
-    for name, p in module.named_parameters():
+        pin_memory=target_device.type != "cpu" and is_pin_memory_available(),
-        if p.device.type == "cpu":
+    ):
-            original_data = p.data
-            device_data = p.data.to(target_device)
-            original_infos[name] = dict(
-                device=p.device,
-                original_data=original_data,
-                device_data=device_data,
-            )
-            p.data = device_data
-        # Parameters already on target device are not touched
-
-    try:
         yield module
 
-    finally:
-        # Restore parameters to their original devices, ignoring new parameters
-        pin_memory = is_pin_memory_available()
-        for name, p in module.named_parameters():
-            if name in original_infos:
-                original_info = original_infos[name]
-                device_data = original_info["device_data"]
-                original_data = original_info["original_data"]
-                original_device: torch.device = original_info["device"]
-
-                if (
-                    (device_data.device == p.data.device)
-                    and (device_data.data_ptr() == p.data.data_ptr())
-                    and (device_data.shape == p.data.shape)
-                    and (device_data.dtype == p.data.dtype)
-                ):
-                    original_data.copy_(p.data.to(original_data.device))
-                    p.data = original_data
-                elif original_device.type == "cpu":
-                    # `torch.empty_like` does not support `pin_memory` argument
-                    cpu_data = torch.empty_strided(
-                        size=p.data.size(),
-                        stride=p.data.stride(),
-                        dtype=p.data.dtype,
-                        layout=p.data.layout,
-                        device="cpu",
-                        pin_memory=pin_memory,
-                    )
-                    cpu_data.copy_(p.data)
-                    p.data = cpu_data
-                else:
-                    p.data = p.data.to(original_device)
-        # New parameters or parameters already on target device are untouched
-
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/model_loader/post_load.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Device staging for post-load weight processing."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Iterator
+
+import torch
+from torch import nn
+
+__all__ = ["stage_module_for_post_load"]
+
+
+@dataclass(slots=True)
+class _TensorState:
+    tensor: torch.Tensor
+    original_data: torch.Tensor
+    origin: torch.device
+    staged_data: torch.Tensor | None = None
+
+
+_SlotKey = tuple[int, str, str]
+
+
+def _iter_registered_tensors(
+    module: nn.Module,
+) -> Iterator[tuple[nn.Module, str, str, torch.Tensor]]:
+    # inspect the registries directly so aliases and non-persistent buffers are
+    # retained. named_parameters()/named_buffers() remove duplicate objects.
+    for owner in module.modules():
+        for registry_name, registry in (
+            ("_parameters", owner._parameters),
+            ("_buffers", owner._buffers),
+        ):
+            for name, tensor in registry.items():
+                if tensor is not None:
+                    yield owner, registry_name, name, tensor
+
+
+def _slot_key(owner: nn.Module, registry_name: str, name: str) -> _SlotKey:
+    return id(owner), registry_name, name
+
+
+def _same_staged_data(current: torch.Tensor, staged: torch.Tensor) -> bool:
+    if current.layout != torch.strided or staged.layout != torch.strided:
+        return False
+    return (
+        current.device == staged.device
+        and current.data_ptr() == staged.data_ptr()
+        and current.shape == staged.shape
+        and current.dtype == staged.dtype
+        and current.stride() == staged.stride()
+    )
+
+
+def _copy_data_to_device(
+    data: torch.Tensor,
+    device: torch.device,
+    *,
+    pin_memory: bool,
+) -> torch.Tensor:
+    if data.device == device:
+        return data
+    if device.type == "cpu" and data.layout == torch.strided and not data.is_quantized:
+        result = torch.empty_strided(
+            size=data.size(),
+            stride=data.stride(),
+            dtype=data.dtype,
+            layout=data.layout,
+            device=device,
+            pin_memory=pin_memory,
+        )
+        result.copy_(data)
+        return result
+    return data.to(device)
+
+
+def _restore_tensor(
+    tensor: torch.Tensor,
+    destination: torch.device,
+    original_state: _TensorState | None,
+    *,
+    pin_memory: bool,
+) -> None:
+    if tensor.is_meta:
+        raise RuntimeError("Post-load processing produced a meta tensor")
+
+    if (
+        original_state is not None
+        and tensor is original_state.tensor
+        and original_state.staged_data is not None
+        and _same_staged_data(tensor.data, original_state.staged_data)
+    ):
+        original_state.original_data.copy_(tensor.data)
+        tensor.data = original_state.original_data
+        return
+
+    tensor.data = _copy_data_to_device(
+        tensor.data,
+        destination,
+        pin_memory=pin_memory,
+    )
+
+
+@contextmanager
+def stage_module_for_post_load(
+    module: nn.Module,
+    process_device: torch.device,
+    *,
+    pin_memory: bool = False,
+) -> Iterator[nn.Module]:
+    """Temporarily stage a module's registered state for a post-load hook.
+
+    Existing parameters and buffers are restored to their per-slot devices.
+    Replacements inherit the slot device unless they newly occupy a
+    non-persistent buffer slot; those remain on ``process_device`` because
+    state-dict based offload cannot move them on demand. Other new tensors use
+    their owner's unique original device, then the module's unique original
+    device, and otherwise remain on ``process_device``.
+
+    This context only owns tensor residency. Hook selection, invocation count,
+    and model/component lifecycle remain the caller's responsibility.
+    """
+
+    if process_device.type == "meta":
+        raise ValueError("process_device cannot be meta")
+
+    original_slots: dict[_SlotKey, _TensorState] = {}
+    original_nonpersistent_buffer_slots: set[_SlotKey] = set()
+    original_name_origins: dict[tuple[int, str], torch.device] = {}
+    owner_origins: dict[int, set[torch.device]] = {}
+    module_origins: set[torch.device] = set()
+    tensor_states: dict[int, _TensorState] = {}
+
+    # snapshot and validate all state before moving any tensor
+    for owner, registry_name, name, tensor in _iter_registered_tensors(module):
+        if tensor.is_meta:
+            raise RuntimeError(
+                f"Cannot post-process meta tensor {type(owner).__name__}.{name}"
+            )
+        state = tensor_states.get(id(tensor))
+        if state is None:
+            state = _TensorState(tensor, tensor.data, tensor.device)
+            tensor_states[id(tensor)] = state
+        key = _slot_key(owner, registry_name, name)
+        original_slots[key] = state
+        if registry_name == "_buffers" and name in owner._non_persistent_buffers_set:
+            original_nonpersistent_buffer_slots.add(key)
+        original_name_origins[(id(owner), name)] = tensor.device
+        owner_origins.setdefault(id(owner), set()).add(tensor.device)
+        module_origins.add(tensor.device)
+
+    try:
+        for state in tensor_states.values():
+            if state.origin != process_device:
+                state.staged_data = state.tensor.data.to(process_device)
+                state.tensor.data = state.staged_data
+        yield module
+    finally:
+        restore_plan: dict[
+            int, tuple[torch.Tensor, torch.device, _TensorState | None]
+        ] = {}
+        unique_module_origin = (
+            next(iter(module_origins)) if len(module_origins) == 1 else None
+        )
+        for owner, registry_name, name, tensor in _iter_registered_tensors(module):
+            key = _slot_key(owner, registry_name, name)
+            original_state = original_slots.get(key)
+            if original_state is None:
+                original_state = tensor_states.get(id(tensor))
+            original_name_origin = original_name_origins.get((id(owner), name))
+            is_new_nonpersistent_buffer = (
+                registry_name == "_buffers"
+                and name in owner._non_persistent_buffers_set
+                and key not in original_nonpersistent_buffer_slots
+            )
+            if is_new_nonpersistent_buffer:
+                destination = process_device
+            elif original_state is not None:
+                destination = original_state.origin
+            else:
+                destination = original_name_origin
+                if destination is None:
+                    origins = owner_origins.get(id(owner), set())
+                    destination = next(iter(origins)) if len(origins) == 1 else None
+                if destination is None:
+                    destination = unique_module_origin or process_device
+
+            if tensor.is_meta:
+                raise RuntimeError("Post-load processing produced a meta tensor")
+            previous = restore_plan.get(id(tensor))
+            if previous is not None:
+                if previous[1] != destination:
+                    raise RuntimeError(
+                        "Aliased post-load tensor has conflicting restore devices"
+                    )
+                if previous[2] is None and original_state is not None:
+                    restore_plan[id(tensor)] = (tensor, destination, original_state)
+                continue
+            restore_plan[id(tensor)] = (tensor, destination, original_state)
+
+        restore_errors: list[Exception] = []
+        for tensor, destination, original_state in restore_plan.values():
+            try:
+                _restore_tensor(
+                    tensor,
+                    destination,
+                    original_state,
+                    pin_memory=pin_memory,
+                )
+            except Exception as error:
+                restore_errors.append(error)
+        if len(restore_errors) == 1:
+            raise restore_errors[0]
+        if restore_errors:
+            raise RuntimeError(
+                f"Multiple post-load tensor restores failed: {restore_errors!r}"
+            ) from restore_errors[0]

test/registered/unit/test_quantization_post_load.py

@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+import torch
+from torch import nn
+
+from sglang.srt.model_loader.loader import device_loading_context
+from sglang.srt.model_loader.post_load import stage_module_for_post_load
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small")
+
+
+def _process_device() -> torch.device | None:
+    if torch.cuda.is_available():
+        return torch.device("cuda", torch.cuda.current_device())
+    return None
+
+
+PROCESS_DEVICE = _process_device()
+
+
+@unittest.skipUnless(PROCESS_DEVICE is not None, "requires CUDA")
+class TestModulePostLoadStaging(unittest.TestCase):
+    def test_preserves_unchanged_parameter_and_buffer_storage(self):
+        module = nn.Module()
+        module.weight = nn.Parameter(torch.arange(4.0).reshape(2, 2))
+        module.register_buffer("scale", torch.ones(2))
+        weight = module.weight
+        scale = module.scale
+        weight_ptr = weight.data_ptr()
+        scale_ptr = scale.data_ptr()
+
+        with stage_module_for_post_load(module, PROCESS_DEVICE):
+            self.assertEqual(module.weight.device, PROCESS_DEVICE)
+            self.assertEqual(module.scale.device, PROCESS_DEVICE)
+            module.weight.data.add_(1)
+            module.scale.add_(2)
+
+        self.assertIs(module.weight, weight)
+        self.assertIs(module.scale, scale)
+        self.assertEqual(module.weight.data_ptr(), weight_ptr)
+        self.assertEqual(module.scale.data_ptr(), scale_ptr)
+        torch.testing.assert_close(module.weight, torch.arange(4.0).reshape(2, 2) + 1)
+        torch.testing.assert_close(module.scale, torch.full((2,), 3.0))
+
+    def test_restores_replacements_and_new_state_with_mixed_residency(self):
+        module = nn.Module()
+        module.weight = nn.Parameter(torch.ones(2, 2))
+        module.register_buffer("resident", torch.ones(1, device=PROCESS_DEVICE))
+        module.child = nn.Module()
+        module.child.scale = nn.Parameter(torch.ones(2))
+
+        with stage_module_for_post_load(module, PROCESS_DEVICE):
+            for tensor in (*module.parameters(), *module.buffers()):
+                self.assertEqual(tensor.device, PROCESS_DEVICE)
+            module.weight = nn.Parameter(torch.ones(3, 2, device=PROCESS_DEVICE))
+            del module.child._parameters["scale"]
+            module.child.register_buffer(
+                "scale", torch.ones(3, device=PROCESS_DEVICE), persistent=False
+            )
+            module.new_parameter = nn.Parameter(torch.ones(1, device=PROCESS_DEVICE))
+            module.child.register_buffer(
+                "new_buffer", torch.ones(1, device=PROCESS_DEVICE)
+            )
+            module.child.register_buffer(
+                "runtime_buffer", torch.ones(1, device=PROCESS_DEVICE), persistent=False
+            )
+
+        # ordinary replacements inherit origin; nonpersistent buffers stay executable
+        self.assertEqual(module.weight.device.type, "cpu")
+        self.assertEqual(module.weight.shape, (3, 2))
+        self.assertEqual(module.child.scale.device, PROCESS_DEVICE)
+        self.assertIn("scale", module.child._buffers)
+        # the child's unique origin wins over the mixed module residency
+        self.assertEqual(module.child.new_buffer.device.type, "cpu")
+        # non-persistent runtime state is absent from offloader state_dicts
+        self.assertEqual(module.child.runtime_buffer.device, PROCESS_DEVICE)
+        # a new tensor owned by the mixed root remains on the process device
+        self.assertEqual(module.new_parameter.device, PROCESS_DEVICE)
+        self.assertEqual(module.resident.device, PROCESS_DEVICE)
+
+    def test_restores_existing_and_new_state_after_hook_error(self):
+        module = nn.Module()
+        module.weight = nn.Parameter(torch.ones(2, 2))
+        module.register_buffer("scale", torch.ones(1))
+
+        with self.assertRaisesRegex(RuntimeError, "hook failed"):
+            with stage_module_for_post_load(module, PROCESS_DEVICE):
+                module.weight.data.add_(1)
+                module.register_buffer(
+                    "workspace", torch.ones(1, device=PROCESS_DEVICE)
+                )
+                raise RuntimeError("hook failed")
+
+        self.assertEqual(module.weight.device.type, "cpu")
+        self.assertEqual(module.scale.device.type, "cpu")
+        self.assertEqual(module.workspace.device.type, "cpu")
+
+    def test_preserves_registered_alias(self):
+        module = nn.Module()
+        shared = nn.Parameter(torch.ones(2))
+        module.left = shared
+        module.right = shared
+
+        with stage_module_for_post_load(module, PROCESS_DEVICE):
+            self.assertIs(module.left, module.right)
+            module.left.data.mul_(2)
+
+        self.assertIs(module.left, shared)
+        self.assertIs(module.right, shared)
+        self.assertEqual(module.left.device.type, "cpu")
+        torch.testing.assert_close(module.left, torch.full((2,), 2.0))
+
+    def test_rejects_alias_with_conflicting_restore_devices_before_restoring(self):
+        module = nn.Module()
+        module.left = nn.Parameter(torch.ones(1))
+        module.child = nn.Module()
+        module.child.right = nn.Parameter(torch.ones(1, device=PROCESS_DEVICE))
+
+        with self.assertRaisesRegex(RuntimeError, "conflicting restore devices"):
+            with stage_module_for_post_load(module, PROCESS_DEVICE):
+                shared = nn.Parameter(torch.ones(1, device=PROCESS_DEVICE))
+                module.left = shared
+                module.child.right = shared
+
+        self.assertIs(module.left, module.child.right)
+        self.assertEqual(module.left.device, PROCESS_DEVICE)
+
+
+class TestModulePostLoadValidation(unittest.TestCase):
+    @unittest.skipUnless(PROCESS_DEVICE is not None, "requires CUDA")
+    def test_srt_cpu_loading_context_remains_a_noop(self):
+        module = nn.Module()
+        module.weight = nn.Parameter(torch.ones(1, device=PROCESS_DEVICE))
+
+        with device_loading_context(module, torch.device("cpu")):
+            self.assertEqual(module.weight.device, PROCESS_DEVICE)
+
+        self.assertEqual(module.weight.device, PROCESS_DEVICE)
+
+    def test_rejects_meta_before_moving_other_state(self):
+        module = nn.Module()
+        module.meta_weight = nn.Parameter(torch.empty(1, device="meta"))
+        module.register_buffer("scale", torch.ones(1))
+        scale_ptr = module.scale.data_ptr()
+
+        with self.assertRaisesRegex(RuntimeError, "meta tensor"):
+            with stage_module_for_post_load(module, torch.device("cpu")):
+                self.fail("context should not be entered")
+
+        self.assertEqual(module.scale.device.type, "cpu")
+        self.assertEqual(module.scale.data_ptr(), scale_ptr)
+
+
+if __name__ == "__main__":
+    unittest.main()