[Qwen3.5] Missing `linear_attn` entries in `base_model_tp_plan` causes OOM and shape error at TP>1
Author: tengomucho
Pull request: no
Labels: name: bug

### System Info

(Manually generated, working on an experimental framework)
transformers version: 5.12.1
Python version: 3.12.13
Huggingface_hub version: 1.20.1
Safetensors version: 0.8.0
Accelerate version: 1.14.0
Accelerate config: not found
DeepSpeed version: not installed
PyTorch version (accelerator?): 2.12.1 (NA)
Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker @vasqu @Cyrilvallez 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

`Qwen3_5TextConfig.base_model_tp_plan` only covers `self_attn.*` and `mlp.*` layers:

```python
base_model_tp_plan = {
    "layers.*.self_attn.q_proj": "colwise",
    "layers.*.self_attn.k_proj": "colwise",
    "layers.*.self_attn.v_proj": "colwise",
    "layers.*.self_attn.o_proj": "rowwise",
    "layers.*.self_attn.q_norm": "replicated_with_grad_allreduce",
    "layers.*.self_attn.k_norm": "replicated_with_grad_allreduce",
    "layers.*.mlp.gate_proj": "colwise",
    "layers.*.mlp.up_proj": "colwise",
    "layers.*.mlp.down_proj": "rowwise",
}
```

Qwen3.5 is a hybrid model: roughly 75% of its decoder layers are `linear_attention` layers backed by `Qwen3_5GatedDeltaNet` (Gated Delta Net / GDA). These layers have their own projection matrices (`in_proj_qkv`, `in_proj_z`, `in_proj_b`, `in_proj_a`, `out_proj`) that are absent from `base_model_tp_plan`. This produces two failures at TP>1:
**Failure 1: OOM when loading the model at TP>1**

The `linear_attn` weights are not sharded, so each rank holds a full copy. For `Qwen3.5-9B` at TP=4 this causes an out-of-memory error before any inference starts.

**Failure 2: RuntimeError during `model.generate()`**

Adding the `linear_attn` projections with the same `"colwise"` / `"rowwise"` pattern as `self_attn` does not fix this. `Qwen3_5GatedDeltaNet.forward` passes `in_proj_qkv`'s output directly through a depthwise `Conv1d` (with `groups=conv_dim`, where `conv_dim = key_dim*2 + value_dim = 8192`). With plain `"colwise"` TP, each rank receives only `conv_dim/TP = 2048` channels, but the `Conv1d` weight expects `8192`. A `torch.split([key_dim, key_dim, value_dim])` immediately after the conv has the same requirement.

```
RuntimeError: convolution(): expected the second dimension of the weight tensor
of shape [8192, 1, 4] to be 0 (2048 in channels divided by 8192 groups), got 1
```

### Root cause

`base_model_tp_plan` in `Qwen3_5TextConfig` is incomplete: it has no entries for any `linear_attn.*` submodule. Because roughly 75% of the model's layers use `linear_attention`, a large fraction of the parameters are not sharded at all when `tp_plan="auto"` is requested.

### Proposed fix

**File to edit: `src/transformers/models/qwen3_5/modular_qwen3_5.py`** (`configuration_qwen3_5.py` is auto-generated from it. The header says "Do NOT edit this file manually"; CI regenerates it from the modular source.)

Add `"colwise_gather_output"` entries for all `linear_attn` projections. `"colwise_gather_output"` is `ColwiseParallel(gather_output=True)`: the weight matrix is sharded across ranks (fixing the OOM), and the output activation is all-gathered before the next operation, so the forward pass sees the full tensor.

```python
base_model_tp_plan = {
    # ... existing self_attn and mlp entries unchanged ...
    "layers.*.linear_attn.in_proj_qkv": "colwise_gather_output",
    "layers.*.linear_attn.in_proj_z":   "colwise_gather_output",
    "layers.*.linear_attn.in_proj_b":   "colwise_gather_output",
    "layers.*.linear_attn.in_proj_a":   "colwise_gather_output",
    "layers.*.linear_attn.out_proj":    "colwise_gather_output",
}
```

`conv1d` is a `nn.Conv1d`, not a `nn.Linear`, so it does not appear in the TP plan. With `colwise_gather_output` on `in_proj_qkv`, the activations are gathered before `conv1d` runs, so the conv layer needs no TP plan entry.

### Why not `"colwise"` / `"rowwise"` (the self_attn pattern)?

The `self_attn` pattern works because `q/k/v_proj` outputs feed into independent per-head computations that decompose naturally across ranks, and `o_proj` with `"rowwise"` expects a sharded input from the preceding colwise-sharded value projection. The GDA forward pass does not have this property. `in_proj_qkv`'s output is a single concatenated `[Q, K, V]` tensor fed directly into a depthwise conv and then split. Neither operation is rank-decomposable. `"colwise_gather_output"` handles this by sharding the weights and inserting an all-gather before those ops.


### Expected behavior

I would expect the model instantiation and generate to work.


Comments:
Not looking too deeply into this yet but can we not do
```py
base_model_tp_plan = {
    # ... existing self_attn and mlp entries unchanged ...
    "layers.*.linear_attn.in_proj_qkv": "colwise",
    "layers.*.linear_attn.in_proj_z":   "colwise",
    "layers.*.linear_attn.in_proj_b":   "colwise",
    "layers.*.linear_attn.in_proj_a":   "colwise",
    "layers.*.linear_attn.conv1d":   "colwise_conv",  # TODO: need to check on which dim the conv needs to be sharded 
    "layers.*.linear_attn.out_proj":    "rowwise",
}
```
Currently the code is not suited for this but in theory it should work. We still have `num_heads * head_dim` which should be properly shardable. The biggest issue is probably https://github.com/huggingface/transformers/blob/5bbed98a97f1d3b3ef5e15936a126e719c98573a/src/transformers/models/qwen3_5/modeling_qwen3_5.py#L504-L512
but other than that, the code already looks "TP friendly". 

Thanks for taking a look @vasqu. The colwise/colwise_conv/rowwise approach is the right direction for actually sharding computation end-to-end through the GDA block, not just sharding weights and gathering at the boundary.

I went with colwise_gather_output because the current forward pass has three spots that would break with plain colwise:

1. torch.split at uses absolute channel counts:
```python
query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
```
2. With colwise on in_proj_qkv, mixed_qkv after the conv has conv_dim/TP channels, but the split sums to conv_dim. It'd need [key_dim//TP, key_dim//TP, value_dim//TP], same range you flagged.
3. dt_bias and A_log are replicated params of size num_v_heads, but in_proj_a with colwise produces a of size num_v_heads/TP:
```python
g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
```
4. Shape mismatch on the addition. These would need to be sliced per-rank.
5. colwise_conv doesn't exist yet AFAIK.

colwise_gather_output sidesteps all three by gathering before the conv, so the rest of the forward sees full-size tensors and needs no changes. It works today at the cost of each rank running the full conv+GDA rather than a sharded version.

Very valid points. Trying to summarize some points together because they are closely related
- `1. + 2. + 5.` Agree this is what I also flagged, some portions of the code are not TP friendly written, mainly the split (but that's not too bad to fix imo). The conv dim ofc also needs to be split to account for the new ranges; this does not exist and was only an idea - it should probably readjust the dimension on which we slice.
- `3.` This one I missed actually and would need bigger changes because afaik we only operate on nn.Modules so far. 
- `4.` Relates to 3. I think(?) because this won't match anymore

I guess we can still get the benefit of avoiding hosting the full weights but we are far from efficient 😓 can you check #46847 in that case

Yes @vasqu I agree with you that a `colwise_conv` implementation would be a better solution, but I just approved #46847 because it solves the problem right now. 
I think we should merge that and think about a `colwise_conv` implementation in a separate issue.