Repository navigation
Native TP breaks on Qwen3.5/3.6 hybrids: reshape error in linear-attention layers #48306
Description
Found while testing native TP training on recent Qwen models.
Reproduction
# torchrun --nproc-per-node 8 repro.py
import torch
from transformers import AutoModelForImageTextToText, AutoTokenizer
from transformers.distributed.configuration_utils import DistributedConfig
model_id = "Qwen/Qwen3.6-35B-A3B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForImageTextToText.from_pretrained(
model_id, dtype=torch.bfloat16, distributed_config=DistributedConfig(tp_size=8)
)
batch = tokenizer(["Some text " * 100], return_tensors="pt").to(model.device)
model(**batch)Outcome
Model loads and shards (10.5 GB/rank for 35B-A3B, consistent with TP8), but the forward fails:
RuntimeError: shape '[4, 178, -1, 256]' is invalid for input of size 45568
The reshape uses the full-attention head_dim (256) on the output of a projection that the TP plan sharded, in a linear-attention layer whose projections have their own head dims (linear_key_head_dim, linear_value_head_dim). The auto TP plan appears not to cover the linear-attention (gated delta) layers of the qwen3_5/qwen3_6 architecture correctly.
Not investigated further; non-hybrid MoE shards and runs fine with the same code. Secondary observation on the same path: torch.optim.AdamW defaults fail on a TP model with aten._foreach_mul_.Scalar got mixed torch.Tensor and DTensor (some params are not DTensors); foreach=False works around it (#48204 already handles this class of problem on the Trainer side with per-parameter AdamW).
System info
transformers main (a353632), torch 2.13.0+cu130, Python 3.12.12, 8x NVIDIA H100 80GB, torchrun single node.
Activity
remi-or commentedon Aug 27, 2026
cc. @3outeille for TP and @ArthurZucker since we talked about this recently
3outeille commentedon Aug 27, 2026
If I am not mistaken im seeing every
linear_attentionlayers getting acolwise_gather_outputplanI ran your script with
tp_size=2and it works which makes me realize thattp_size <= num_key_value_heads. So the issue is more that we dont have KV replication. A PR from a contributor to bring KV replication is on the way #47811github-actions commentedon Sep 25, 2026
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.
remi-or commentedon Sep 26, 2026
Issue should be fixed by #48848. Commenting to avoid stale issue.