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

[lettt325]

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:

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:

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:

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, ...)

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.