--cpu-offload crashes at warmup: KV cache allocated on meta device for offloaded layers ("Cannot copy out of meta tensor") #105

Summary

With --cpu-offload on a GPU that actually needs offloading (8GB RTX 3070), inference crashes during warmup with:

NotImplementedError: Cannot copy out of meta tensor; no data!

The crash happens for both moshi.server and moshi.offline, with NO_TORCH_COMPILE=1 and with/without NO_CUDA_GRAPH=1. This likely affects everyone asking about 8GB VRAM support (#10, #29).

Environment

• Windows 11, RTX 3070 8GB, driver 595.97
• Python 3.12, torch 2.13.0+cu130 (also reproduced on torch 2.4.1+cu124)
• accelerate 1.14.0 (also reproduced on 1.3.0)
• Model: nvidia/personaplex-7b-v1

Repro

python -m moshi.offline --cpu-offload \
  --voice-prompt voice.wav \
  --input-wav assets/test/input_assistant.wav \
  --output-wav out.wav --output-text out.json

On a GPU large enough that infer_auto_device_map places every module on GPU, the bug does not trigger — which is presumably why it isn't seen on development hardware.

Root cause

StreamingMultiheadAttention._init_streaming_state allocates the streaming KV cache on the device of the layer's weights:

device = self.in_proj_weight.device

Under accelerate's cpu-offload, offloaded layers' parameters are meta placeholders (real weights live in the offload weights map and are restored per-forward by AlignDevicesHook). So the KV cache is allocated on the meta device, and the first forward through the hook dies in send_to_device with the error above.

Fix that works for me

In moshi/modules/transformer.py, resolve the meta device to the hook's execution device:

device = self.in_proj_weight.device
if device.type == "meta":
    hook = getattr(self, "_hf_hook", None)
    exec_dev = getattr(hook, "execution_device", None)
    device = torch.device(exec_dev if exec_dev is not None else "cuda")

With this patch, --cpu-offload runs to completion on the 3070 (19 modules on GPU / 23 on CPU per infer_auto_device_map) and produces correct audio output.

Happy to open a PR if useful.