[Bug] PyTorch engine: Qwen3.5 hybrid-GDN + AWQ checkpoint crashes on TP>1 - 'AwqLinear' object has no attribute 'block_size'
- Dominant language
- Python
- Stars
- 8.1k
- Forks
- 748
- Avg merge
- 6d 2h
- Merged PRs (30d)
- 54
Description
### Checklist
- [x] 1. I have searched related issues but cannot get the expected help.
- [x] 2. The bug has not been fixed in the latest version.
- [x] 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.
### Describe the bug
**Environment**: lmdeploy 0.16.0 (latest), Python 3.12.3, Ubuntu 24.04, CUDA 12.8 runtime.
**Hardware**: 2 x RTX 2080 Ti 22GB (Turing SM75), NVLink, TP=2.
**Model**: community true-AWQ export of Qwen3.8-27B (`quantization_config: {"quant_method": "awq", "bits": 4, "group_size": 128}`, separate MTP tensors). Qwen3.8-27B is a hybrid model: 48x `linear_attention` + 16x `full_attention` layers.
**What happens**: with `--backend pytorch --tp 2`, the engine subprocess dies during weight loading (`RPCServerDeadError`). Root cause frames below.
**Root cause analysis**
`Qwen3_5TextModel._patch_qkv_weight_loader()` overrides `mod.weight_loader` for the fused GDN QKV projection (non-uniform sections `[key_dim, key_dim, value_dim]`). When a loaded 2D tensor has `shape[0] < sum(sections)` it assumes a fp8-style **scale** shrunk along dim0 and computes `bs = mod.block_size`. For AWQ checkpoints the module is `AwqLinear`, which carries **`group_size`** instead of `block_size` -> AttributeError.
Two extra hazards worth addressing together:
1. `AwqLinear.setup_loaders()` routes ALL four params (`qweight`, `scales`, `qzeros`, `bias`) through `mod.weight_loader`. AWQ `qweight` is transposed-packed (`dim0 ~= in_features/pack`), so the shape heuristic `shape[0] < sum(sections)` can misclassify a `qweight` as a scale and mis-split it -> silent corruption instead of a crash.
2. Suggested fix: branch explicitly on `param._weight_type` (`'qweight' | 'scales' | 'qzeros' | 'bias'`, already set by `setup_loaders()`), shrink scales/qzeros by `group_size`, and handle the transposed `qweight` separately.
TP=1 does not trigger the bug (`mod.is_tp == False` -> `default_weight_loader`). TurboMind backend loads the same artifacts fine; this affects the PyTorch engine path only. Happy to test any fix - failing artifact and hardware are on hand.
### Reproduction
lmdeploy serve api_server /path/to/Qwen3.8-27B-AWQ-MTP --backend pytorch --tp 2 --server-port 8002 Crashes while loading safetensors shards; the engine child process exits before the server comes up.
### Environment
```Shell
- lmdeploy 0.16.0 (pip, latest), torch 2.x cu12x
- Python 3.12.3, Ubuntu 24.04
- 2x RTX 2080 Ti 22GB modded (SM75), NVLink
- Model: shawnw3i/Qwen3.8-27B-AWQ-MTP (autoawq format, group_size 128)
```
### Error traceback
```Shell
File ".../lmdeploy/pytorch/models/qwen3_5.py", line 1249, in load_weights
load_weight(param, loaded_weight)
File ".../lmdeploy/pytorch/weight_loader/model_weight_loader.py", line 22, in load_weight
param.weight_loader(param, loaded_weight, **kwargs)
File ".../lmdeploy/pytorch/models/qwen3_5.py", line 503, in qkv_weight_loader
bs = mod.block_size
AttributeError: 'AwqLinear' object has no attribute 'block_size'
(via ray::RayWorkerWrapper.build_model() -> RPCServerDeadError: PyTorch ZMQ engine process is not alive.)
```
Contributor guide
Research direction
Start in lmdeploy/pytorch/models/qwen3_5.py at Qwen3_5TextModel._patch_qkv_weight_loader() and the qkv_weight_loader traceback, then inspect AwqLinear.setup_loaders() and load_weight in lmdeploy/pytorch/weight_loader/model_weight_loader.py. Reproduce with the Qwen3.8-27B-AWQ-MTP command using --backend pytorch --tp 2. Done means AWQ qweight, scales, qzeros, and bias load without misclassification and the engine completes weight loading instead of crashing.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, distributed-systems
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100