[Bug]: After wake up from level 2 sleep, model cannot load weights properly

[fingertap]

Your current environment

The output of `python collect_env.py`

PyTorch version: 2.6.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04) 11.3.0
Clang version: Could not collect
CMake version: version 4.0.0
Libc version: glibc-2.35

Python version: 3.12.9 | packaged by Anaconda, Inc. | (main, Feb  6 2025, 18:56:27) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-153-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A800-SXM4-80GB
GPU 1: NVIDIA A800-SXM4-80GB
GPU 2: NVIDIA A800-SXM4-80GB
GPU 3: NVIDIA A800-SXM4-80GB
GPU 4: NVIDIA A800-SXM4-80GB
GPU 5: NVIDIA A800-SXM4-80GB
GPU 6: NVIDIA A800-SXM4-80GB
GPU 7: NVIDIA A800-SXM4-80GB

Nvidia driver version: 535.161.08
cuDNN version: Could not collect
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
Address sizes:                   43 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          192
On-line CPU(s) list:             0-79
Off-line CPU(s) list:            80-191
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC 7643 48-Core Processor
CPU family:                      25
Model:                           1
Thread(s) per core:              2
Core(s) per socket:              48
Socket(s):                       2
Stepping:                        1
Frequency boost:                 enabled
CPU max MHz:                     2300.0000
CPU min MHz:                     1500.0000
BogoMIPS:                        4591.58
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 aperfmperf pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate sme ssbd mba sev ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca
Virtualization:                  AMD-V
L1d cache:                       3 MiB (96 instances)
L1i cache:                       3 MiB (96 instances)
L2 cache:                        48 MiB (96 instances)
L3 cache:                        512 MiB (16 instances)
NUMA node(s):                    2
NUMA node0 CPU(s):               0-47,96-143
NUMA node1 CPU(s):               48-95,144-191
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] flashinfer-python==0.2.2.post1+cu124torch2.6
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-cusparselt-cu12==0.6.2
[pip3] nvidia-ml-py==12.570.86
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] pynvml==12.0.0
[pip3] pyzmq==26.4.0
[pip3] torch==2.6.0
[pip3] torchaudio==2.6.0
[pip3] torchvision==0.21.0
[pip3] transformers==4.51.1
[pip3] triton==3.2.0
[conda] flashinfer-python         0.2.2.post1+cu124torch2.6          pypi_0    pypi
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] nvidia-cublas-cu12        12.4.5.8                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.4.127                 pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.2.1.3                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.5.147               pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.6.1.9                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.3.1.170               pypi_0    pypi
[conda] nvidia-cusparselt-cu12    0.6.2                    pypi_0    pypi
[conda] nvidia-ml-py              12.570.86                pypi_0    pypi
[conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.4.127                 pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.4.127                 pypi_0    pypi
[conda] pynvml                    12.0.0                   pypi_0    pypi
[conda] pyzmq                     26.4.0                   pypi_0    pypi
[conda] torch                     2.6.0                    pypi_0    pypi
[conda] torchaudio                2.6.0                    pypi_0    pypi
[conda] torchvision               0.21.0                   pypi_0    pypi
[conda] transformers              4.51.1                   pypi_0    pypi
[conda] triton                    3.2.0                    pypi_0    pypi
ROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: 0.8.2
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
        GPU0    GPU1    GPU2    GPU3    GPU4    GPU5    GPU6    GPU7    NIC0    NIC1    NIC2    NIC3    NIC4    NIC5    NIC6    NIC7    NIC8    NIC9    CPU Affinity    NUMA Affinity   GPU NUMA ID
GPU0     X      NV8     NV8     NV8     NV8     NV8     NV8     NV8     NODE    PXB     PXB     PXB     SYS     SYS     PXB     NODE    SYS     SYS     0-47    0               N/A
GPU1    NV8      X      NV8     NV8     NV8     NV8     NV8     NV8     NODE    PXB     PXB     PXB     SYS     SYS     PXB     NODE    SYS     SYS     0-47    0               N/A
GPU2    NV8     NV8      X      NV8     NV8     NV8     NV8     NV8     PXB     NODE    NODE    NODE    SYS     SYS     NODE    PXB     SYS     SYS     0-47    0               N/A
GPU3    NV8     NV8     NV8      X      NV8     NV8     NV8     NV8     PXB     NODE    NODE    NODE    SYS     SYS     NODE    PXB     SYS     SYS     0-47    0               N/A
GPU4    NV8     NV8     NV8     NV8      X      NV8     NV8     NV8     SYS     SYS     SYS     SYS     NODE    PXB     SYS     SYS     PXB     NODE    48-79   1               N/A
GPU5    NV8     NV8     NV8     NV8     NV8      X      NV8     NV8     SYS     SYS     SYS     SYS     NODE    PXB     SYS     SYS     PXB     NODE    48-79   1               N/A
GPU6    NV8     NV8     NV8     NV8     NV8     NV8      X      NV8     SYS     SYS     SYS     SYS     PXB     NODE    SYS     SYS     NODE    PXB     48-79   1               N/A
GPU7    NV8     NV8     NV8     NV8     NV8     NV8     NV8      X      SYS     SYS     SYS     SYS     PXB     NODE    SYS     SYS     NODE    PXB     48-79   1               N/A
NIC0    NODE    NODE    PXB     PXB     SYS     SYS     SYS     SYS      X      NODE    NODE    NODE    SYS     SYS     NODE    PIX     SYS     SYS
NIC1    PXB     PXB     NODE    NODE    SYS     SYS     SYS     SYS     NODE     X      PIX     PIX     SYS     SYS     PIX     NODE    SYS     SYS
NIC2    PXB     PXB     NODE    NODE    SYS     SYS     SYS     SYS     NODE    PIX      X      PIX     SYS     SYS     PIX     NODE    SYS     SYS
NIC3    PXB     PXB     NODE    NODE    SYS     SYS     SYS     SYS     NODE    PIX     PIX      X      SYS     SYS     PIX     NODE    SYS     SYS
NIC4    SYS     SYS     SYS     SYS     NODE    NODE    PXB     PXB     SYS     SYS     SYS     SYS      X      NODE    SYS     SYS     NODE    PIX
NIC5    SYS     SYS     SYS     SYS     PXB     PXB     NODE    NODE    SYS     SYS     SYS     SYS     NODE     X      SYS     SYS     PIX     NODE
NIC6    PXB     PXB     NODE    NODE    SYS     SYS     SYS     SYS     NODE    PIX     PIX     PIX     SYS     SYS      X      NODE    SYS     SYS
NIC7    NODE    NODE    PXB     PXB     SYS     SYS     SYS     SYS     PIX     NODE    NODE    NODE    SYS     SYS     NODE     X      SYS     SYS
NIC8    SYS     SYS     SYS     SYS     PXB     PXB     NODE    NODE    SYS     SYS     SYS     SYS     NODE    PIX     SYS     SYS      X      NODE
NIC9    SYS     SYS     SYS     SYS     NODE    NODE    PXB     PXB     SYS     SYS     SYS     SYS     PIX     NODE    SYS     SYS     NODE     X 

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks

NIC Legend:

  NIC0: mlx5_0
  NIC1: mlx5_1
  NIC2: mlx5_2
  NIC3: mlx5_3
  NIC4: mlx5_4
  NIC5: mlx5_5
  NIC6: mlx5_6
  NIC7: mlx5_7
  NIC8: mlx5_8
  NIC9: mlx5_9

NCCL_IB_TC=186
NCCL_IB_PCI_RELAXED_ORDERING=1
NCCL_SOCKET_IFNAME=eth0
NCCL_NVLS_ENABLE=0
NCCL_IB_HCA==mlx5_6,mlx5_7,mlx5_8,mlx5_9
NCCL_IB_GID_INDEX=5
NCCL_PXN_DISABLE=1
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
NCCL_IB_QPS_PER_CONNECTION=8
NCCL_IB_TIMEOUT=21
LD_LIBRARY_PATH=/data/cuda/cuda-12.4/cuda/lib64:/usr/local/nvidia/lib64
NCCL_IB_DISABLE=0
NCCL_IB_RETRY_CNT=7
NCCL_CUMEM_ENABLE=0
TORCHINDUCTOR_COMPILE_THREADS=1
CUDA_MODULE_LOADING=LAZY

🐛 Describe the bug

I want to drop all the parameters of the LLM engine and then load it from another state dict. The following code shows that level 1 is OK, while level 2 is problematic (cannot load the weights properly).

import torch
import torch.distributed as dist
from vllm import SamplingParams, LLM
from functools import cached_property


class SleepLevelTwoWakeUpIssue:
    model_path = "Qwen/Qwen2.5-7B-Instruct"

    def run(self):
        self.log("\033[91mRaw weights:\033[0m")
        self.generate_and_print()
        self.sleep_and_wake_up(1)
        self.log("\033[91mAfter sleep level 1:\033[0m")
        self.generate_and_print()
        self.sleep_and_wake_up(1, load_weights=True)
        self.log("\033[91mAfter sleep level 1 and load weights:\033[0m")
        self.generate_and_print()
        self.sleep_and_wake_up(2)
        self.log("\033[91mAfter sleep level 2:\033[0m")
        self.generate_and_print()
        self.sleep_and_wake_up(2, load_weights=True)
        self.log("\033[91mAfter sleep level 2 and load weights:\033[0m")
        self.generate_and_print()

    def generate_and_print(self):
        prompts = [
            "Hello, how are you?",
            "France is famous for its",
            "The capital of USA is",
        ]
        sampling_params = SamplingParams(max_tokens=100, temperature=1.0, stop=["\n"], seed=0)
        outputs = self.llm.generate(prompts, sampling_params=sampling_params)
        for output in outputs:
            self.log(output.outputs[0].text)

    def sleep_and_wake_up(self, level: int, load_weights: bool = False):
        self.llm.sleep(level=level)
        self.llm.wake_up()
        if load_weights:
            model = self.llm.llm_engine.model_executor.driver_worker.worker.model_runner.model
            # model.load_weights(weights=self.state_dict.items())  # This leads to the same results
            named_parameters, named_buffers = self.named_parameters_and_named_buffers
            for name, param in named_parameters:
                model.load_weights(weights=[(name, param)])
            for name, buffer in named_buffers:
                model.load_weights(weights=[(name, buffer)])

    def log(self, message: str):
        if dist.get_rank() != 0:
            return
        print(message)

    @cached_property
    def state_dict(self):
        from transformers import AutoModelForCausalLM
        model = AutoModelForCausalLM.from_pretrained(self.model_path)
        return model.state_dict()

    @cached_property
    def named_parameters_and_named_buffers(self):
        from transformers import AutoModelForCausalLM
        model = AutoModelForCausalLM.from_pretrained(self.model_path)
        named_parameters = list(model.named_parameters())
        named_buffers = list(model.named_buffers())
        return named_parameters, named_buffers

    @cached_property
    def llm(self):
        return LLM(
            enable_sleep_mode=True,
            model=self.model_path,
            gpu_memory_utilization=0.8,
            distributed_executor_backend="external_launcher",
            tensor_parallel_size=4,
            max_model_len=16384,
            seed=0,
            disable_custom_all_reduce=True,
            dtype="bfloat16",
        )


if __name__ == "__main__":
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(dist.get_rank())
    SleepLevelTwoWakeUpIssue().run()
    dist.destroy_process_group()

Before submitting a new issue...

• 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.

[toslali-ibm]

+1 on this. I had to revert to sleep level 1 on my end as well.

@youkaichao

[fingertap]

Indeed, this bug influences many downstream users, including all users of verl and openrlhf.

[fingertap]

# vllm/device_allocator/cumem.py

def wake_up(self):
        """
        Wake up the allocator from sleep mode.
        All data that is previously offloaded will be loaded back to GPU 
        memory, and the rest of the data will have empty memory."""
        for ptr, data in self.pointer_to_data.items():
            handle = data.handle
            create_and_map(handle)
            if data.cpu_backup_tensor is not None:
                cpu_backup_tensor = data.cpu_backup_tensor
                cpu_backup_tensor = torch.empty_like(cpu_backup_tensor)  # Use another tensor as source
                if cpu_backup_tensor is not None:
                    size_in_bytes = cpu_backup_tensor.numel(
                    ) * cpu_backup_tensor.element_size()
                    cpu_ptr = cpu_backup_tensor.data_ptr()
                    libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes)
                    data.cpu_backup_tensor = None

I am also confused by this behavior. As shown in the above code, if I use libcudart.cudaMemcpy to copy the weights from another tensor, the weight loader will also fail. I think I am missing some important thing here. Can you give any input? I would be delighted to contribute~ @erictang000

[erictang000]

@fingertap I also saw similar behavior where the outputs are not well defined after sleep level 2 - when i went through the self.pointer_to_data.items(), there was usually some pointer that didn't correspond to any parameters in model.named_parameters() (and I guess based on your test also not in model.named_buffers()). Seems like this is noted here: #15254 (comment) as well

[fingertap]

Thanks!

I also saw similar behavior where the outputs are not well defined after sleep level 2 - when i went through the self.pointer_to_data.items(), there was usually some pointer that didn't correspond to any parameters in model.named_parameters() (and I guess based on your test also not in model.named_buffers()).

I also noticed this. But I cannot think of an explanation to the model's mangled output if both the parameters and buffers are correct, if the model.load_weights supports loading the buffers.

[fingertap]

Hi @erictang000 !
I would like to look into this issue, but I don't have any clues now. Can you provide some hints on this so that I can dig into the issue?

[youkaichao]

Indeed, this bug influences many downstream users, including all users of verl and openrlhf.

the verl team should be aware of this, and they use sleep mode with level 1.

the problem is, when updating weights, usually people only update weights, but they forgot to update buffers as well, like the rope cos-sin cache. that will also be discarded when sleeping with level 2.

[fingertap]

Thanks!

the verl team should be aware of this, and they use sleep mode with level 1.

Yes. I mean they are using level 1 due to this issue, too, even when the weights are actually discarded (as the weights are updated later by the actor).

but they forgot to update buffers as well

I am aware of this and I tried loading buffers with load_weights, as shown in the example above, but it still produces mangled text.

I think I must miss something important about the memory allocation. There must be something I fail to recover from the sleep mode.

The fact that level 1 sleep is OK indicates using load_weights to load parameters works as expected. One possible cause of this level 2 issue is that the buffers cannot be loaded properly with load_weights.

[youkaichao]

I tried loading buffers with load_weights

load_weights does not work for buffers. i think you can keep a copy of the buffer after creating the model, and then copy these buffers back after updating weights.

something like:

buffers = module.named_buffers()
buffers = {k: v.clone() for k, v in buffers.items()} # clone the buffer, they are not managed by the cumem allocator.

# run something
# sleep
# wake up

# deal with the buffers manually.
empty_buffers = module.named_buffers()
for k, v in empty_buffers.items():
    v.copy_(buffers[k])

[fingertap]

This can be a workaround!
To solve this issue for good, can we introduce a new data tag "buffers"?

[youkaichao]

tag does not work here. buffers and parameters are created interleavingly, and there're no easy ways to know if an allocation is for buffers or for parameters.

if the above example code works, we can clone the buffer when we sleep at level 2, and copy the buffer when we wake up.

please give it a try and see if it works.

[fingertap]

i think you can keep a copy of the buffer after creating the model, and then copy these buffers back after updating weights.

Confirmed this works!