v0.3.4.post1 consumes more gpu momory?

[zeng-zc]

I've had a try on the latest release version v0.3.4.post1 recently. I suspect this version consumes more gpu memory which may be related to the cpu-time-hidden feature? I'm not very sure. I encounter OutOfMemory errors in v0.3.4.post1 which seldom appeared before in the older version.

lanching server commands in docker-compose.yaml:

    entrypoint: python3 -m sglang.launch_server
    command:
      --model-path Qwen/Qwen2-72B-Instruct
      --host 0.0.0.0 --port 30000
      --tp-size 8
      --mem-fraction-static 0.8

client side:

python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 1000 \
        --sharegpt-output-len 256 \
        --multi 

the OutOfMemory logs:

sglang  | [2024-10-23 18:12:14 TP0] Prefill batch. #new-seq: 1, #new-token: 8192, #cached-token: 1239, cache hit rate: 86.70%, token usage: 0.34, #running-req: 167, #queue-req: 0
sglang  | [2024-10-23 18:12:14 TP3] Traceback (most recent call last):
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 1133, in run_scheduler_process
sglang  |     scheduler.event_loop_normal()
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
sglang  |     return func(*args, **kwargs)
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 281, in event_loop_normal
sglang  |     result = self.run_batch(batch)
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py", line 727, in run_batch
sglang  |     logits_output, next_token_ids = self.tp_worker.forward_batch_generation(
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/managers/tp_worker.py", line 139, in forward_batch_generation
sglang  |     logits_output = self.model_runner.forward(forward_batch)
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner.py", line 588, in forward
sglang  |     return self.forward_extend(forward_batch)
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner.py", line 572, in forward_extend
sglang  |     return self.model.forward(
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
sglang  |     return func(*args, **kwargs)
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/models/qwen2.py", line 289, in forward
sglang  |     hidden_states = self.model(input_ids, positions, forward_batch, input_embeds)
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
sglang  |     return self._call_impl(*args, **kwargs)
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
sglang  |     return forward_call(*args, **kwargs)
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/models/qwen2.py", line 255, in forward
sglang  |     hidden_states, residual = layer(
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
sglang  |     return self._call_impl(*args, **kwargs)
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
sglang  |     return forward_call(*args, **kwargs)
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/models/qwen2.py", line 215, in forward
sglang  |     hidden_states = self.mlp(hidden_states)
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
sglang  |     return self._call_impl(*args, **kwargs)
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
sglang  |     return forward_call(*args, **kwargs)
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/models/qwen2.py", line 78, in forward
sglang  |     x, _ = self.down_proj(x)
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
sglang  |     return self._call_impl(*args, **kwargs)
sglang  |   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1562, in _call_impl
sglang  |     return forward_call(*args, **kwargs)
sglang  |   File "/sgl-workspace/sglang/python/sglang/srt/layers/linear.py", line 1143, in forward
sglang  |     output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
sglang  |   File "/usr/local/lib/python3.10/dist-packages/vllm/model_executor/layers/linear.py", line 135, in apply
sglang  |     return F.linear(x, layer.weight, bias)
sglang  | torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 128.00 MiB. GPU 3 has a total capacity of 39.56 GiB of which 48.56 MiB is free. Process 196119 has 39.51 GiB memory in use. Of the allocated memory 32.01 GiB is allocated by PyTorch, with 959.14 MiB allocated in private pools (e.g., CUDA Graphs), and 355.76 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

The machine is A100-SXM4-40GB with 8 gpus.