Support explicit attention `head_dim` independent of `hidden_size / num_attention_heads`

Disclaimer: The below text was largely written by ChatGPT after a long discussion with me and was independently verified by me.

GPT-NeoX currently assumes that the dimension of each attention head is always:

head_dim = hidden_size // num_attention_heads

This forces the concatenated query-head width to equal the residual-stream width:

num_attention_heads * head_dim == hidden_size

That relationship is common, but it is not required by multi-head attention and is not true for several current architectures. In particular, the official Qwen3-0.6B, Qwen3-4B, and Qwen3-32B configurations specify head_dim: 128 even though hidden_size / num_attention_heads is respectively 64, 80, and 80.

GPT-NeoX should support an explicit per-head dimension, falling back to the current derivation when it is not configured.

This is required before the configurations proposed in #1399 can faithfully represent all Qwen3 dense models.

Background

These dimensions describe different parts of the transformer:

• hidden_size is the width of the residual stream entering and leaving each transformer block.
• num_attention_heads is the number of query heads.
• head_dim is the width of each query, key, and value head.
• num_kv_heads is the number of key/value heads used by grouped-query attention.

An attention implementation can use projections with these shapes:

q_proj: hidden_size -> num_attention_heads * head_dim
k_proj: hidden_size -> num_kv_heads * head_dim
v_proj: hidden_size -> num_kv_heads * head_dim
o_proj: num_attention_heads * head_dim -> hidden_size

The attention-internal width does not have to equal the residual width. The input projections can expand the representation before attention, and the output projection contracts it back before the residual addition.

Hugging Face’s Qwen3 implementation follows this design:

self.head_dim = getattr(
    config,
    "head_dim",
    config.hidden_size // config.num_attention_heads,
)

It sizes the query projection from num_attention_heads * head_dim, the key/value projections from num_key_value_heads * head_dim, and the output projection from num_attention_heads * head_dim back to hidden_size.

See the Qwen3 attention implementation in Transformers v4.51.

The quotient is a backward-compatible default, not a constraint that an explicit head_dim must equal.

Concrete Qwen3 examples

The official Qwen3 dense configurations use the following dimensions:

The existing behavior happens to match Qwen3-1.7B, Qwen3-8B, and Qwen3-14B because those models satisfy:

hidden_size == num_attention_heads * head_dim

It does not match Qwen3-0.6B, Qwen3-4B, or Qwen3-32B.

Qwen3-0.6B example

The required attention path is:

Residual input:
[..., 1024]
      │
      ├─ q_proj: 1024 → 16 × 128 = 2048
      ├─ k_proj: 1024 →  8 × 128 = 1024
      └─ v_proj: 1024 →  8 × 128 = 1024

Attention:
Q: [..., 16, 128]
K: [...,  8, 128]  ── each KV head serves 2 query heads
V: [...,  8, 128]

Concatenated attention output:
[..., 16 × 128] = [..., 2048]
      │
      └─ o_proj: 2048 → 1024

Residual addition:
[..., 1024] + [..., 1024]

There is no separate resize operation. The learned Q/K/V projections perform the expansion, and o_proj restores the residual width.

Current GPT-NeoX limitation

The root assumption is in ParallelSelfAttention.__init__:

self.hidden_size_per_attention_head = mpu.divide(
    neox_args.hidden_size,
    neox_args.num_attention_heads,
)

The argument schema exposes hidden_size, num_attention_heads, and num_kv_heads, but no independent head dimension. See neox_args.py.

The derived value propagates through the attention implementation:

The most visible projection assumption is conceptually:

self.query_key_value = ColumnParallelLinear(
    input_size=neox_args.hidden_size,
    output_size=neox_args.hidden_size + 2 * self.kv_hidden_size,
    ...
)

The first term assumes that the total query width equals hidden_size.

For Qwen3-0.6B, the resulting difference is:

Current GPT-NeoX:
QKV projection: 1024 → 1024 + 2×512
Output projection: 1024 → 1024

Required:
QKV projection: 1024 → 2048 + 2×1024
Output projection: 2048 → 1024

Proposed configuration behavior

Add an optional argument matching the Hugging Face name:

head_dim: Optional[int] = None

Semantics:

effective_head_dim = (
    neox_args.head_dim
    if neox_args.head_dim is not None
    else neox_args.hidden_size // neox_args.num_attention_heads
)

Existing configurations that omit head_dim must retain their current shapes and checkpoint compatibility.

When head_dim is explicitly configured, GPT-NeoX should not require:

hidden_size % num_attention_heads == 0

That divisibility requirement is only needed to derive head_dim. Tensor-parallel head-count and projection-partitioning constraints must still be validated.

A Qwen3 configuration would use:

"head_dim": 128,

Proposed implementation

Maintain separate values for the residual width and the attention-internal widths:

self.hidden_size_per_attention_head = (
    neox_args.head_dim
    if neox_args.head_dim is not None
    else mpu.divide(
        neox_args.hidden_size,
        neox_args.num_attention_heads,
    )
)

self.query_hidden_size = (
    neox_args.num_attention_heads
    * self.hidden_size_per_attention_head
)

num_kv_heads = (
    neox_args.num_kv_heads
    if neox_args.num_kv_heads is not None
    else neox_args.num_attention_heads
)

self.kv_hidden_size = (
    num_kv_heads
    * self.hidden_size_per_attention_head
)

The partitioned query width should also be tracked independently:

self.query_hidden_size_per_partition = mpu.divide(
    self.query_hidden_size,
    tensor_model_parallel_world_size,
)

Fused QKV projection

For grouped-query attention:

self.query_key_value = ColumnParallelLinear(
    neox_args=neox_args,
    input_size=neox_args.hidden_size,
    output_size=self.query_hidden_size + 2 * self.kv_hidden_size,
    ...
)

For ordinary multi-head attention:

output_size = 3 * self.query_hidden_size

The Q/K/V split sizes must remain based on the partitioned query-head count, partitioned KV-head count, and effective head_dim.

Attention output flattening

The concatenated attention heads must be flattened to the attention-internal query width, not the residual width:

new_context_layer_shape = context_layer.size()[:-2] + (
    self.query_hidden_size_per_partition,
)

Output projection

The row-parallel output projection should contract the attention-internal width back to the residual width:

self.dense = RowParallelLinear(
    input_size=self.query_hidden_size,
    output_size=neox_args.hidden_size,
    ...
)

Scaling and RoPE

Attention scaling must use:

effective_head_dim**-0.5

RoPE construction and application must also use the effective head dimension, subject to the existing rotary_pct behavior and even-dimension validation.

Transformer Engine backend

GPT-NeoX pins transformer-engine[pytorch]==1.12.

Transformer Engine 1.12 already supports an attention head dimension that differs from hidden_size / num_attention_heads. Its MultiheadAttention constructor accepts:

kv_channels: Optional[int] = None

Despite its name, kv_channels controls the per-head dimension of Q, K, and V. When it is omitted, TE derives it using:

kv_channels = hidden_size // num_attention_heads

When it is supplied, TE 1.12 calculates:

hidden_size_q = kv_channels * num_attention_heads
hidden_size_kv = kv_channels * num_gqa_groups

and constructs the projections as:

QKV:  hidden_size → hidden_size_q + 2 × hidden_size_kv
PROJ: hidden_size_q → hidden_size

See the TE 1.12 constructor and dimension calculation, QKV projection construction, and attention/output projection construction. The corresponding GQA split and reshape logic also uses the configured per-head dimension (source.

For Qwen3-0.6B, passing:

hidden_size=1024
num_attention_heads=16
num_gqa_groups=8
kv_channels=128

produces:

hidden_size_q  = 16 × 128 = 2048
hidden_size_kv =  8 × 128 = 1024

QKV projection: 1024 → 2048 + 2×1024 = 4096
Output projection: 2048 → 1024

This matches the required Qwen3 attention geometry.

GPT-NeoX’s current TEMultiheadAttention wrapper does not pass kv_channels, so TE falls back to hidden_size // num_attention_heads. The wrapper also independently derives its RoPE dimension using the same quotient.

The wrapper should instead calculate the effective head dimension once:

effective_head_dim = (
    neox_args.head_dim
    if neox_args.head_dim is not None
    else neox_args.hidden_size // neox_args.num_attention_heads
)

and pass it to Transformer Engine:

super(TEMultiheadAttention, self).__init__(
    hidden_size=neox_args.hidden_size,
    num_attention_heads=neox_args.num_attention_heads,
    kv_channels=effective_head_dim,
    num_gqa_groups=self.num_kv_heads,
    ...
)

RoPE should then use the same effective value. Since the TE superclass assigns it to self.hidden_size_per_attention_head, the wrapper can reuse that attribute instead of overwriting it with a newly derived quotient.

No Transformer Engine upgrade is required for this head-dimension behavior. TE 1.12 already sizes the QKV projection, GQA tensors, attention scaling, output projection, and inference KV cache from kv_channels.

Checkpoint conversion

Hugging Face checkpoint import/export code must read hf_config.head_dim when present instead of unconditionally deriving it.

Converters must support weights with these shapes:

q_proj.weight: [num_attention_heads * head_dim, hidden_size]
k_proj.weight: [num_kv_heads * head_dim, hidden_size]
v_proj.weight: [num_kv_heads * head_dim, hidden_size]
o_proj.weight: [hidden_size, num_attention_heads * head_dim]

This is especially important because Qwen3 checkpoints with expanded query widths cannot be reshaped into the current NeoX projection dimensions without losing parameters.

Tensor-parallel considerations

The implementation should validate the dimensions that are actually sharded:

• Query heads must remain compatible with the tensor-parallel size.
• KV-head sharding or replication must preserve the existing GQA semantics.
• The partitioned fused-QKV output must split on whole head boundaries.
• query_hidden_size_per_partition must be used when flattening local attention output.
• The row-parallel output projection must accept the local partition of query_hidden_size.

The residual width itself does not need to equal the total query width.

Backward compatibility

When head_dim is omitted:

head_dim = hidden_size // num_attention_heads
query_hidden_size = hidden_size

Therefore:

• Existing projection shapes remain unchanged.
• Existing checkpoints remain loadable.
• Existing YAML files require no modification.
• Existing attention scaling and RoPE dimensions remain unchanged.
• Parameter counts and tensor-parallel layouts remain unchanged.

This can be introduced as a backward-compatible extension.

Tests and acceptance criteria

• Add an optional documented head_dim configuration argument.

• Preserve current behavior exactly when head_dim is absent.

• Only require hidden_size % num_attention_heads == 0 when head_dim must be derived.

• Size Q, K, V, and output projections from the effective head dimension.

• Flatten attention output using the query-attention width rather than the residual width.

• Use the effective head dimension for attention scaling, RoPE, cache shapes, and reshaping.

• Support the existing GQA path.

• Provide a clear error or full support in the Transformer Engine path.

• Update relevant Hugging Face checkpoint converters.

• Test a Qwen3-0.6B-style configuration:

  复制

  hidden_size=1024
  num_attention_heads=16
  num_kv_heads=8
  head_dim=128
  query_hidden_size=2048
  kv_hidden_size=1024

• Test forward and backward passes with tensor-parallel sizes 1 and 2.

• Test both standard and FlashAttention paths.

• Test inference KV-cache allocation and decoding.

• Add a regression test proving that omitted head_dim retains existing shapes.

• Ideally add shape or checkpoint-conversion tests for the Qwen3-4B and Qwen3-32B geometries.

Existing ecosystem precedent

This is an established configuration pattern rather than a Qwen3-only special case:

• Hugging Face Transformers issue #37187 requested changing Qwen2-MoE from an unconditional hidden_size // num_attention_heads derivation to an explicit head_dim with that quotient as the fallback.
• The corresponding Transformers PR #37188 was merged. During review, maintainers confirmed that the behavior matches Qwen2/Qwen3 and noted that explicit head dimensions had become standard.
• A Gemma 3 model discussion covers the same apparent discrepancy: hidden_size=640, four query heads, and an explicit head_dim=256 rather than the derived value of 160. A Google-org response confirms that the explicit model configuration is authoritative.
• Modern Megatron Bridge Gemma 3 configurations use independent values such as hidden_size=640, num_attention_heads=4, and kv_channels=256.
• vLLM’s attention implementation supports an optional explicit head_dim, computes q_size = num_heads * head_dim, and constructs o_proj from num_heads * head_dim back to hidden_size.
• The distinction between an architectural constraint and a library implementation constraint has also been discussed in the PyTorch multi-head attention discussion.