[Bug]: vLLM crashes with larger context sizes on TPUs

[francescov1]

Your current environment

The output of `python collect_env.py`

Collecting environment information...
INFO 09-10 05:05:28 importing.py:10] Triton not installed; certain GPU-related functions will not be available.
PyTorch version: 2.5.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 11 (bullseye) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.10.14 (main, Aug 13 2024, 02:16:06) [GCC 10.2.1 20210110] (64-bit runtime)
Python platform: Linux-6.1.85+-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Byte Order:                           Little Endian
Address sizes:                        48 bits physical, 48 bits virtual
CPU(s):                               224
On-line CPU(s) list:                  0-223
Thread(s) per core:                   2
Core(s) per socket:                   56
Socket(s):                            2
NUMA node(s):                         2
Vendor ID:                            AuthenticAMD
CPU family:                           25
Model:                                1
Model name:                           AMD EPYC 7B13
Stepping:                             0
CPU MHz:                              2449.998
BogoMIPS:                             4899.99
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            3.5 MiB
L1i cache:                            3.5 MiB
L2 cache:                             56 MiB
L3 cache:                             448 MiB
NUMA node0 CPU(s):                    0-55,112-167
NUMA node1 CPU(s):                    56-111,168-223
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Mitigation; safe RET
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP conditional; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr arat npt nrip_save umip vaes vpclmulqdq rdpid fsrm

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] pyzmq==26.2.0
[pip3] torch==2.5.0
[pip3] torch-xla==2.5.0+git17a4ef5
[pip3] torchvision==0.19.0a0+d23a6e1
[pip3] transformers==4.44.2
[conda] Could not collect
ROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: 0.6.0@1447c97e753919709b613590d7267c93d07d9382
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
Could not collect

🐛 Describe the bug

I am trying to deploy llama3.1 8B instruct on GKE 2x4 v5e TPUs. The vLLM server boots up normally and works properly if max-model-len is very low (ie 1024), but crashes with higher context sizes (tried 16k and above). I'm using the Dockerfile.tpu, and the command python3.10 -u -m vllm.entrypoints.openai.api_server --host 0.0.0.0 --model meta-llama/Meta-Llama-3.1-8B-Instruct --tensor-parallel-size 8 --swap-space 16 --disable-log-requests --max-model-len=32768

Here is my full log file: vllm-llama31-tpu.log

Here is my full k8s manifest:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: llama3-8b-instruct-vllm-deployment
  namespace: default
  labels:
    app: llama3-8b-instruct-vllm
spec:
  replicas: 1
  selector:
    matchLabels:
      app: llama3-8b-instruct-vllm
  template:
    metadata:
      labels:
        app: llama3-8b-instruct-vllm
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
      hostIPC: true
      hostNetwork: true
      containers:
      - name: llama3-8b-instruct-vllm
        image: us-central1-docker.pkg.dev/project-lighthouse-403916/lighthouse-production/vllm-server/vllm-tpu-image:latest
        command: ["/bin/sh", "-c", "python3.10 -u -m vllm.entrypoints.openai.api_server --host 0.0.0.0 --model meta-llama/Meta-Llama-3.1-8B-Instruct --tensor-parallel-size 8 --swap-space 16 --disable-log-requests --max-model-len=32768"]
        env:
        - name: HUGGING_FACE_HUB_TOKEN
          value: <hf-token>
        - name: VLLM_LOGGING_LEVEL
          value: "DEBUG"
        - name: VLLM_TRACE_FUNCTION
          value: "1"
        securityContext:
          privileged: true
        ports:
        - containerPort: 8000
          protocol: TCP
        resources:
          requests:
            google.com/tpu: 8
          limits:
            google.com/tpu: 8

Before submitting a new issue...

• [X] Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[youkaichao]

RuntimeError: Bad StatusOr access: RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space smem. Used 1.00M of 1.00M smem. Exceeded smem capacity by 1.7K.

looks like some TPU compilation error.

[JackCaoG]

Seem like a smem OOM compiler time error. Given that v5e per chip HBM is only 16G OOMing at 16k seq length is kind of expected.

[miladm]

What's the largest seq len that you can run?

[francescov1]

@miladm Was able to get it running with --max-model-len 8192, but 16k crashed it

[JackCaoG]

Hey @francescov1 can you use

XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1 XLA_SAVE_TENSORS_FMT="hlo" XLA_SAVE_TENSORS_FILE="/tmp/save1.hlo"

and share the /tmp/save1.hlo(i think it will be /tmp/save1.hlo.0 in this case) with us? The error here is that XLA compiler complains that it doesn't have SMEM space(we need to move parameter from HBM to SMEM before executing each op) for this HLO. We can open a bug to the XLA compiler team and see if there is anything we can do about it.

[francescov1]

@JackCaoG Here you go. File was too big to upload so put on gdrive https://drive.google.com/file/d/1oSY7QqCIsQTZS2tGxJbFV-R_rOepDwHl/view?usp=sharing

Theres a number of other save1.hlo.X files, let me know if you need any others

[JackCaoG]

Thanks @francescov1 I downloaded the HLO, from the HLO it seems like that VLLM is warming up the cache and throw a compilation error in the last compilation. And that compilation error is the RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space smem. Used 1.00M of 1.00M smem right? If so I can just take the last HLO and open a bug to the XLA team.

[francescov1]

Yep @JackCaoG that's what i saw too. Sounds good, let me know if you need anything else from me

and could you send the link to the issue or tag this one in it so i can track it

[WoosukKwon]

This must be fixed by #9350.