[Qwen35] Fix device mismatch RuntimeError for unmanaged parameters during offloading
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 6h
- Merged PRs (30d)
- 155
Description
### Description
When running larger Qwen3.5 variants in llm_qwen3_5_text_gen (e.g., Qwen3.5-9B) that trigger partial CPU offloading in ComfyUI's memory management, executing a prompt results in a device mismatch `RuntimeError`.
Specifically in file qwen35.py , `self.A_log` and `self.dt_bias` inside `GatedDeltaNet` are plain `nn.Parameter` tensors that remain on the CPU (or wherever they were offloaded) when `a` (residing on the GPU) is computed. This causes a crash during the gating calculation at line 215:
`RuntimeError: Expected all tensors to be on the same device...`
### Cause
Unlike `conv1d.weight`, which is properly handled using `comfy.model_management.cast_to_device`, the linear attention parameters (`A_log` and `dt_bias`) lack explicit device casting during the forward pass.
### Solution
Use `comfy.model_management.cast_to_device` on `A_log` and `dt_bias` to ensure they match the execution device, similarly to how other layers handle memory offloading.
### Current code (crashable when offloaded to RAM) file qwen35.py
beta = b.sigmoid()
g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias.float())
### Proposed code correction for qwen35.py file
beta = b.sigmoid()
A_log_casted = comfy.model_management.cast_to_device(self.A_log, x.device, torch.float32)
dt_bias_casted = comfy.model_management.cast_to_device(self.dt_bias, x.device, torch.float32)
g = -A_log_casted.exp() * F.softplus(a.float() + dt_bias_casted)
Contributor guide
Research direction
Inspect qwen35.py around line 215 and compare the unmanaged GatedDeltaNet parameters with the existing device-casting path for conv1d.weight. Reproduce with a larger Qwen3.5 variant such as Qwen3.5-9B under partial CPU offloading, then verify that the prompt runs without a device mismatch RuntimeError.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 76/100