MILES row-parallel fc2/down_proj reduction order does not match SGLang
Background
When testing strict train/inference consistency for Qwen3-4B, we found that the MLP fc2/down_proj outputs are not bitwise identical between training and rollout.
Root cause
The mismatch comes from a different accumulation order in the row-parallel matmul.
For Qwen3-4B fc2/down_proj with TP=2, the local K dimension is:
K = 4864
BLOCK_K = 128
T = K / BLOCK_K = 38
SGLang matmul_tp_inv computes:
FIRST_LEVEL_BLOCK = 19
LEVEL_K = 2
So its effective reduction order is:
out =
(((p0 + p1) + p2 + ... + p18)
+ (((p19 + p20) + p21 + ... + p37))
MILES currently simulates this path with a global pairwise tree over all partials, effectively closer to:
out =
(((p0 + p1) + (p2 + p3)) + ...)
These two orders are not bitwise equivalent for bf16 accumulation. As a result, even when each 128-wide partial matmul is identical, the final fc2/down_proj output can differ because the outer reduction order is different.
MILES row-parallel fc2/down_proj reduction order does not match SGLang
Background
When testing strict train/inference consistency for Qwen3-4B, we found that the MLP
fc2/down_projoutputs are not bitwise identical between training and rollout.Root cause
The mismatch comes from a different accumulation order in the row-parallel matmul.
For Qwen3-4B
fc2/down_projwith TP=2, the local K dimension is:SGLang
matmul_tp_invcomputes:So its effective reduction order is:
MILES currently simulates this path with a global pairwise tree over all partials, effectively closer to:
These two orders are not bitwise equivalent for bf16 accumulation. As a result, even when each 128-wide partial matmul is identical, the final
fc2/down_projoutput can differ because the outer reduction order is different.