Skip to content

Defrag breaks on hybrid attention models (Qwen3.5: linear_attention + full_attention) #154

Description

@joelteply

Problem

Qwen3.5-9B forge crashes during cycle 1 training:
```
RuntimeError: shape '[2, 256, -1, 256]' is invalid for input of size 196608
```

Root Cause

Qwen3.5 is a hybrid multimodal model:

  • Config is nested: `text_config.head_dim = 256`, `text_config.num_attention_heads = ...`
  • Layer types alternate: `linear_attention` and `full_attention` per layer
  • Defrag code reads top-level `config.head_dim` (None) and falls back to `hidden_size / num_heads = 128`
  • But the actual head_dim is 256, set in `text_config`
  • Defrag slices tensors based on the wrong head_dim, producing fractional head counts

Additional concern: `linear_attention` layers use a different attention mechanism (likely Mamba/SSM-style state space) and shouldn't be pruned the same way as standard `full_attention` layers.

Fix Path

  1. Read head_dim from `text_config` (or any nested config) when top-level is None
  2. Detect layer type from `config.text_config.layer_types[i]` and SKIP defrag on `linear_attention` layers
  3. Add Qwen3.5 to the validation harness Layer 3/4 tests
  4. Document supported architectures clearly

Workaround

Use single-cycle no-prune forge_pipeline.py for hybrid models until this is fixed.
Standard Qwen3 (9B-Instruct, 4B, etc.) should work — only Qwen3.5 has the linear_attention split.

Refs

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions