Comfy-Org / Comfy-Org/ComfyUI

[Qwen35] Fix device mismatch RuntimeError for unmanaged parameters during offloading

Open Beginner friendly
#15,193 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.