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)
- 主要语言
- Python
- 星标
- 2.4k
- 派生
- 607
- 平均合并
- 2 天 19 小时
- 30 天内合并 PR
- 158
描述
### 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`.
贡献指南
评估
这个 Issue 还没有评估数据。