diff --git a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py index e8df934b5..8af6ad1a6 100644 --- a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py @@ -65,31 +65,9 @@ class LayerwiseOffloadManager: self._named_buffers: Dict[str, torch.Tensor] = {} # Store forward hooks for removal self._forward_hooks: List[Any] = [] - # GPU buffer pool: dtype -> numel -> [buffers] - self._gpu_buffer_pool: Dict[torch.dtype, Dict[int, List[torch.Tensor]]] = {} - # layer_idx -> {dtype: gpu_buffer} - self._layer_gpu_buffers: Dict[int, Dict[torch.dtype, torch.Tensor]] = {} - self._gpu_buffer_pool_max = max(2, 2 * self.prefetch_size) - # Keep layer 0 resident during the forward pass to avoid redundant reloads - self._resident_window = 1 self._initialize() - def _acquire_gpu_buffer(self, dtype: torch.dtype, numel: int) -> torch.Tensor: - pool_by_dtype = self._gpu_buffer_pool.setdefault(dtype, {}) - bucket = pool_by_dtype.get(numel) - if bucket: - return bucket.pop() - return torch.empty((numel,), dtype=dtype, device=self.device) - - def _release_gpu_buffer( - self, dtype: torch.dtype, numel: int, buffer: torch.Tensor - ) -> None: - pool_by_dtype = self._gpu_buffer_pool.setdefault(dtype, {}) - bucket = pool_by_dtype.setdefault(numel, []) - if len(bucket) < self._gpu_buffer_pool_max: - bucket.append(buffer) - def _match_layer_idx(self, name: str) -> int | None: m = self._layer_name_re.search(name) if not m: @@ -104,14 +82,12 @@ class LayerwiseOffloadManager: if not self.enabled: return - named_parameters = list(self.model.named_parameters()) - named_buffers = list(self.model.named_buffers()) - self._named_parameters = dict(named_parameters) - self._named_buffers = dict(named_buffers) + self._named_parameters = dict(self.model.named_parameters()) + self._named_buffers = dict(self.model.named_buffers()) # 1. collect and group tensors by layer and dtype layer_groups: Dict[int, Dict[torch.dtype, List[Tuple[str, torch.Tensor]]]] = {} - all_tensors = list(chain(named_parameters, named_buffers)) + all_tensors = chain(self._named_parameters.items(), self._named_buffers.items()) for name, tensor in all_tensors: layer_idx = self._match_layer_idx(name) if layer_idx is None or layer_idx >= self.num_layers: @@ -197,7 +173,9 @@ class LayerwiseOffloadManager: gpu_buffers: Dict[torch.dtype, torch.Tensor] = {} with torch.cuda.stream(self.copy_stream): for dtype, cpu_buffer in self._consolidated_cpu_weights[layer_idx].items(): - gpu_buffer = self._acquire_gpu_buffer(dtype, cpu_buffer.numel()) + gpu_buffer = torch.empty( + cpu_buffer.shape, dtype=dtype, device=self.device + ) gpu_buffer.copy_(cpu_buffer, non_blocking=non_blocking) gpu_buffers[dtype] = gpu_buffer @@ -218,10 +196,9 @@ class LayerwiseOffloadManager: ].view(meta["shape"]) self._gpu_layers.add(layer_idx) - self._layer_gpu_buffers[layer_idx] = gpu_buffers @torch.compiler.disable - def release_layer(self, layer_idx: int, force: bool = False) -> None: + def release_layer(self, layer_idx: int) -> None: """ lightweight release layer weights Basically set the reference count to the gpu weight tensor to zero. The weights on cpu is untouched @@ -232,7 +209,7 @@ class LayerwiseOffloadManager: # clear prefetch event, since it's useless and needs to be reset self._prefetch_events.pop(layer_idx, None) - if layer_idx < self._resident_window and not force: + if layer_idx <= 0: return if layer_idx not in self._gpu_layers: @@ -242,11 +219,6 @@ class LayerwiseOffloadManager: target = self.get_target_with_name(name) target.data = torch.empty((1,), device=self.device, dtype=meta["dtype"]) - layer_buffers = self._layer_gpu_buffers.pop(layer_idx, None) - if layer_buffers is not None: - for dtype, buffer in layer_buffers.items(): - self._release_gpu_buffer(dtype, buffer.numel(), buffer) - self._gpu_layers.discard(layer_idx) @torch.compiler.disable @@ -257,9 +229,7 @@ class LayerwiseOffloadManager: torch.cuda.current_stream().wait_stream(self.copy_stream) for layer_idx in list(self._gpu_layers): - self.release_layer(layer_idx, force=True) - - self._gpu_buffer_pool.clear() + self.release_layer(layer_idx) @torch.compiler.disable def load_all_layers(self) -> None: