From a23f6ea09032811f200a103a40ccd92d03fd5285 Mon Sep 17 00:00:00 2001 From: Kangrui Du <89372739+Rockdu@users.noreply.github.com> Date: Sat, 25 Jul 2026 13:35:53 -0700 Subject: [PATCH] [Diffusion] offload rollout weights to pinned host memory (#32032) Co-authored-by: Yihao Wang <42559837+AgainstEntropy@users.noreply.github.com> --- .../memory_occupation_controller.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py index 3bffb4024..4a3d55d76 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/memory_occupation_controller.py @@ -16,6 +16,15 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +def _module_to_pinned_cpu(module: torch.nn.Module) -> None: + # Async D2H into pinned host memory; caller synchronizes once after the batch. + for t in list(module.parameters()) + list(module.buffers()): + if t.device.type == "cuda": + pin = torch.empty(t.shape, dtype=t.dtype, device="cpu", pin_memory=True) + pin.copy_(t.data, non_blocking=True) + t.data = pin + + def _get_module_device(module: torch.nn.Module) -> str: """Return best-effort device string for a module.""" param = next(module.parameters(), None) @@ -113,9 +122,13 @@ class MemoryOccupationController: for name in names: module = modules[name] src_device_map[name] = _get_module_device(module) - module.to(device) + if device.startswith("cpu"): + _module_to_pinned_cpu(module) + else: + module.to(device, non_blocking=True) moved.append(name) _move_unregistered_tensors(module, device) + torch.cuda.synchronize() except Exception as e: logger.warning( f"[_move_modules] move failed, rollback started: target={device} moved={moved} error={e}",