AI-Hypercomputer / AI-Hypercomputer/maxtext

partial_rotary_factor is ignored when use_mrope is true (affects qwen3.5-35b-a3b and qwen3.5-397b-a17b)

Đang mở
#4,616 0 bình luận 0 reaction 0 người được giao Xem trên GitHub
Ngôn ngữ chính
Python
Star
2.4k
Fork
607
Merge trung bình
2 ngày 19 giờ
Pull request đã merge (30 ngày)
158

Mô tả

### Summary

`partial_rotary_factor` is silently ignored whenever `use_mrope: true`. Two shipped configs set both, so RoPE is applied to the full head dimension instead of the leading quarter, with no error and no shape mismatch.

Affected:

- `src/maxtext/configs/models/qwen3.5-35b-a3b.yml` — `use_mrope: true`, `partial_rotary_factor: 0.25`
- `src/maxtext/configs/models/qwen3.5-397b-a17b.yml` — `use_mrope: true`, `partial_rotary_factor: 0.25`

### Cause

`Attention.init_rotary_embedding` routes to `Qwen3OmniMoeThinkerTextRotaryEmbedding` whenever `use_mrope` is set (`src/maxtext/layers/attentions.py`, the `elif self.use_mrope:` branch). It passes `embedding_dims=rope_embedding_dims`, which is the full `head_dim`, and never passes the partial factor — the class does not accept one.

`PartialRotaryEmbedding` implements the intended behaviour, but only the non-MRoPE path can reach it.

### Reproduction

On `main`, passing the factor fails outright:

```python
Qwen3OmniMoeThinkerTextRotaryEmbedding(
min_timescale=1, max_timescale=10000, embedding_dims=16,
cast_as_fprop_dtype=False, mrope_section=(2, 1, 1),
partial_rotary_factor=0.25, rngs=nnx.Rngs(params=0),
)
# TypeError: __init__() got an unexpected keyword argument 'partial_rotary_factor'
```

and the default construction rotates every channel:

```python
layer = Qwen3OmniMoeThinkerTextRotaryEmbedding(
min_timescale=1, max_timescale=10000, embedding_dims=16,
cast_as_fprop_dtype=False, mrope_section=(2, 1, 1), rngs=nnx.Rngs(params=0),
)
inputs = jax.random.normal(jax.random.PRNGKey(0), (2, 8, 4, 16), dtype=jnp.float32)
positions = jnp.broadcast_to(jnp.arange(8, dtype=jnp.int32), (2, 8))
out = layer(inputs, positions)

np.allclose(out[..., 4:], inputs[..., 4:])
# -> False
# With partial_rotary_factor 0.25 only the first 4 of 16 channels may be rotated;
# the remaining 12 must pass through untouched.
```

### Suggested fix

Give the MRoPE layer the same construction `PartialRotaryEmbedding` already uses — split off the leading `rotary_dim = int(head_dim * partial_rotary_factor)` channels, build `inv_freq` over `rotary_dim` rather than `head_dim`, and concatenate the untouched remainder back:

```python
self.head_dim = embedding_dims
self.partial_rotary_factor = partial_rotary_factor
self.rotary_dim = int(self.head_dim * self.partial_rotary_factor)
super().__init__(..., embedding_dims=self.rotary_dim, ...)
```

```python
if self.rotary_dim < self.head_dim:
inputs_rot, inputs_pass = jnp.split(inputs, [self.rotary_dim], axis=-1)
else:
inputs_rot, inputs_pass = inputs, None
...
x_out = self.apply_rotary(inputs_rot, cos_emb, sin_emb)
if inputs_pass is not None:
x_out = jnp.concatenate([x_out, inputs_pass], axis=-1)
```

plus one line in `init_rotary_embedding` to pass `config.partial_rotary_factor` through. The shape check should compare against `head_dim` rather than the (now reduced) `embedding_dims`.

A default of `1.0` keeps `qwen3-vl-*` and `qwen3-omni-*` bit-identical — those configs do not set the factor.

### Verification

With that change, four properties hold, checked against `PartialRotaryEmbedding` and against the unmodified layer:

1. channels beyond `rotary_dim` are passed through unchanged, leading ones are rotated;
2. a partial factor no longer produces the fully rotated result (the regression itself);
3. omitting the factor reproduces the previous full-rotation output exactly;
4. for text-only 1D positions — where MRoPE degenerates to ordinary RoPE — the MRoPE path agrees with `PartialRotaryEmbedding` channel for channel (`rtol=1e-5`).

Properties 2 and 4 fail on `main`.

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.