diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 123e5a03f..23efee357 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/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] = {} - - # Store original device states and move parameters to GPU if they're on CPU - for name, p in module.named_parameters(): - 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: + with stage_module_for_post_load( + module, + target_device, + pin_memory=target_device.type != "cpu" and is_pin_memory_available(), + ): 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__) diff --git a/python/sglang/srt/model_loader/post_load.py b/python/sglang/srt/model_loader/post_load.py new file mode 100644 index 000000000..6d0110f08 --- /dev/null +++ b/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] diff --git a/test/registered/unit/test_quantization_post_load.py b/test/registered/unit/test_quantization_post_load.py new file mode 100644 index 000000000..c290d3923 --- /dev/null +++ b/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()