Repository navigation
397B v7x trainer: MoE wo FSDP sharding and valid-rows-only MoE kernels - #5760
Draft
khatwanimohit wants to merge 2 commits into
Draft
khatwanimohit wants to merge 2 commits into
khatwanimohit wants to merge 2 commits into
Conversation
…educe-scatter scatters a middle dim wo is [experts, mlp, embed]. Sharding its embed dimension over FSDP made XLA transpose each wo gradient to [embed, experts*mlp] and reduce-scatter its major dimension on the SparseCore, about 3x slower than wi's middle-dimension scatter for the same bytes (12.9 vs 4.2 ms per call on v7x), and that chain of expert-weight-gradient reduce-scatters ends every backward layer iteration. Mesh rules that define mlp_moe_out/embed_moe_out (cp-as-ep) now shard wo's mlp dimension over FSDP; rule sets without those names keep the old layout. The forward all-gather becomes [64,32,4096]->[64,1024,4096].
… buffer's valid rows Only rows [0, limit) of the ring's packed [buffer_size, ...] expert buffer hold tokens (about 56% on the bench's hot shard at ragged_buffer_factor 2.75). The grouped matmuls already skip the rest, but four XLA fusion families per layer chunk still streamed the whole buffer on every device: the unsort backward (f32[225280] + bf16[225280,4096]), the GLU forward and its backward recompute (bf16[225280,1024]), and the GLU backward, a select fusion writing both half gradients plus a separate pad/maximum pass that concatenated them (bf16[225280,2048]). - kernels/ragged/unsort_bwd_rows.py: a Pallas TensorCore kernel for the buffered branch of _ring_ragged_unsort_bwd. It computes bf16(f32(g) * w) and the f32 routing-weight row dot in 128x128 tiles over (1024, 2048) blocks, with `limit` as a prefetched scalar; blocks past the one holding row limit - 1 map to that block, so they are neither fetched nor written. The weights and the dot are lane-dense [rows / 128, 128] (free reshapes), turned into per-row columns with one 128x128 transpose each. The scatter of the dot is unchanged; off a TPU the XLA code runs. - kernels/ragged/glu_rows.py: forward and backward kernels for silu(gate) * up over the same rows, reading gate and up in place from the fused buffer. The backward writes [d_gate | d_up] in one pass, with the -0 -> +0 fix. They follow XLA's bf16 fusions exactly: the same op sequence (negate, exp, + 1, divide, multiply), in f32, rounding only the outputs (XLA keeps excess precision inside the fusion; op-by-op bf16 rounding differs from it in ~3% of elements). - moe.py: fused_glu(fused, glu, limit=None). _down_projection passes valid_token_count as the limit when the buffer is the ring's packed one and apply_ffn_activation is the plain bf16 SiLU GLU; the kernel path keeps the whole buffer as its residual, under the remat names of both halves. The dropless replay program (unpacked buffer) and every other path keep the XLA code. Rows left unwritten are never read: gmm_v2 reads lhs rows up to the sublane tile of row limit - 1, tgmm_v2 masks rows outside the groups with a select, and the dot scatter drops rows >= limit. Numerics: the weighted unsort gradient and both GLU results are bit-identical; the routing-weight dot changes summation order only. Checks: - v6e, bench shape [225280, 4096], limits 0, 1, 1023, 125952, 126156, 225280: unsort gradient bitwise equal to XLA (rows up to the end of limit's block), dot max |d| 5.3e-5. Interpret mode with the TPU7x chip description and a partial last block: bitwise as well. - v6e, GLU forward and backward against fused_glu's XLA path, every bf16 bit pattern as the gate (x64 random up/d, plus +-0, +-inf, NaN, denormals), 225280-row and partial-block buffers: 0 mismatching bits. - v6e timing at 126156 of 225280 valid rows (XLA -> kernel): unsort 4.71 -> 2.61 ms, GLU forward 1.08 -> 0.61 ms, GLU backward 1.92 -> 1.02 ms; at a full buffer the kernels are within 1-9% of XLA. - Tiny Qwen3.5 on v6e (ab_model.py, use_ragged_sort=true ragged_buffer_factor=2.75): loss, logits and all 66 gradient leaves bit-identical to the unmodified tree. - 397B v7x AOT (wi dlhs tiles 2048, moe_route=device, recompute limit 100): fwd_bwd temp 94.95 -> 92.25 GiB, dropless 56.69 GiB unchanged. module_0000 has 8 unsort_bwd_rows, 16 glu_fwd_rows and 8 glu_bwd_rows calls in place of the 16 activation, 8 unsort, 8 select and 8 concatenate fusions, with no new 225280-row copies.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces optimized Pallas kernels (glu_rows and unsort_bwd_rows) to process only the valid rows of the expert buffer in MoE layers, along with integration in moe.py and unit tests. The review feedback identifies a critical race condition in _unsort_bwd_rows_kernel when num_col_blocks > 1 due to unsynchronized concurrent writes to dot_ref. To fix this, it is recommended to add a column-block dimension to dot_ref and sum the partial dots in JAX. The reviewer also suggests adding defensive checks to ensure inputs are 2D arrays before unpacking, and expanding the unit tests to cover multi-column-block configurations.
Comment on lines
+72
to
+99
| def _unsort_bwd_rows_kernel(limit_ref, g_ref, y_ref, w_ref, grad_ref, dot_ref, *, num_col_blocks): # pylint: disable=too-many-positional-arguments | ||
| """One (1024, cols) block: grad = bf16(f32(g) * w) and the partial f32 row dot, in 128 x 128 tiles.""" | ||
| i, j = pl.program_id(0), pl.program_id(1) | ||
| limit = limit_ref[0] | ||
| row0 = i * _BLOCK_ROWS | ||
| num_col_tiles = g_ref.shape[1] // _LANES | ||
|
|
||
| def row_tile(masked, r, carry): | ||
| rows = pl.ds(pl.multiple_of(r * _LANES, _LANES), _LANES) | ||
| # Row a of the tile gets weight w[row0 + 128 r + a], broadcast along the lanes. | ||
| w_col = jnp.transpose(jnp.broadcast_to(w_ref[pl.ds(r, 1), :], (_LANES, _LANES))) | ||
| if masked: | ||
| row_ids = row0 + r * _LANES + lax.broadcasted_iota(jnp.int32, (_LANES, _LANES), 0) | ||
| valid = row_ids < limit | ||
| acc = jnp.zeros((_LANES, _LANES), jnp.float32) | ||
| for c in range(num_col_tiles): | ||
| cols = pl.ds(c * _LANES, _LANES) | ||
| g = g_ref[rows, cols].astype(jnp.float32) | ||
| if masked: | ||
| g = jnp.where(valid, g, 0.0) | ||
| grad_ref[rows, cols] = (g * w_col).astype(grad_ref.dtype) | ||
| acc = acc + g * y_ref[rows, cols].astype(jnp.float32) | ||
| # Lane-dense: entry a of the row sum over the transposed tile is row a's dot over this block's columns. | ||
| part = jnp.sum(jnp.transpose(acc), axis=0, keepdims=True) | ||
| if num_col_blocks > 1: | ||
| part = part + jnp.where(j == 0, 0.0, dot_ref[pl.ds(r, 1), :]) | ||
| dot_ref[pl.ds(r, 1), :] = part | ||
| return carry |
There was a problem hiding this comment.
There is a critical race condition in _unsort_bwd_rows_kernel when num_col_blocks > 1. Because j (the column block index) is a grid dimension, different grid blocks execute concurrently on the TPU. When num_col_blocks > 1, multiple grid blocks concurrently read and write to the same dot_ref slice without synchronization, leading to data corruption.\n\nTo eliminate this race condition, we can make dot_ref have an extra dimension for the column block (i.e., shape (num_col_blocks, rows // _LANES, _LANES)), so each grid block j writes to its own unique slice. We can then sum the partial dots across column blocks in JAX after the pallas_call.
def _unsort_bwd_rows_kernel(limit_ref, g_ref, y_ref, w_ref, grad_ref, dot_ref):\n \"\"\"One (1024, cols) block: grad = bf16(f32(g) * w) and the partial f32 row dot, in 128 x 128 tiles.\"\"\"\n i, j = pl.program_id(0), pl.program_id(1)\n limit = limit_ref[0]\n row0 = i * _BLOCK_ROWS\n num_col_tiles = g_ref.shape[1] // _LANES\n\n def row_tile(masked, r, carry):\n rows = pl.ds(pl.multiple_of(r * _LANES, _LANES), _LANES)\n # Row a of the tile gets weight w[row0 + 128 r + a], broadcast along the lanes.\n w_col = jnp.transpose(jnp.broadcast_to(w_ref[pl.ds(r, 1), :], (_LANES, _LANES)))\n if masked:\n row_ids = row0 + r * _LANES + lax.broadcasted_iota(jnp.int32, (_LANES, _LANES), 0)\n valid = row_ids < limit\n acc = jnp.zeros((_LANES, _LANES), jnp.float32)\n for c in range(num_col_tiles):\n cols = pl.ds(c * _LANES, _LANES)\n g = g_ref[rows, cols].astype(jnp.float32)\n if masked:\n g = jnp.where(valid, g, 0.0)\n grad_ref[rows, cols] = (g * w_col).astype(grad_ref.dtype)\n acc = acc + g * y_ref[rows, cols].astype(jnp.float32)\n # Lane-dense: entry a of the row sum over the transposed tile is row a's dot over this block's columns.\n part = jnp.sum(jnp.transpose(acc), axis=0, keepdims=True)\n dot_ref[0, pl.ds(r, 1), :] = part[None, :]\n return carry\n\n num_row_tiles = _BLOCK_ROWS // _LANES\n\n @pl.when(row0 + _BLOCK_ROWS <= limit)\n def _full():\n lax.fori_loop(0, num_row_tiles, functools.partial(row_tile, False), 0)\n\n @pl.when(jnp.logical_and(row0 + _BLOCK_ROWS > limit, i == _last_block(limit)))\n def _partial():\n lax.fori_loop(0, num_row_tiles, functools.partial(row_tile, True), 0)
Comment on lines
+112
to
+158
| def _unsort_bwd_rows_pallas(gathered, sorted_tokens, weights, limit, *, interpret=False): | ||
| """The Pallas TensorCore kernel; see `unsort_bwd_rows`.""" | ||
| rows, hidden = gathered.shape | ||
| block_cols = _block_cols(hidden) | ||
| num_row_blocks = pl.cdiv(rows, _BLOCK_ROWS) | ||
| num_col_blocks = hidden // block_cols | ||
|
|
||
| def data_map(i, j, limit_ref): | ||
| last = _last_block(limit_ref[0]) | ||
| return jnp.minimum(i, last), jnp.where(i <= last, j, num_col_blocks - 1) | ||
|
|
||
| def row_map(i, j, limit_ref): | ||
| del j | ||
| return jnp.minimum(i, _last_block(limit_ref[0])), 0 | ||
|
|
||
| data_spec = pl.BlockSpec((_BLOCK_ROWS, block_cols), data_map) | ||
| lane_dense_spec = pl.BlockSpec((_BLOCK_ROWS // _LANES, _LANES), row_map) | ||
| grad, dot = pl.pallas_call( | ||
| functools.partial(_unsort_bwd_rows_kernel, num_col_blocks=num_col_blocks), | ||
| grid_spec=pltpu.PrefetchScalarGridSpec( | ||
| num_scalar_prefetch=1, | ||
| grid=(num_row_blocks, num_col_blocks), | ||
| in_specs=[data_spec, data_spec, lane_dense_spec], | ||
| out_specs=[data_spec, lane_dense_spec], | ||
| ), | ||
| out_shape=[ | ||
| jax.ShapeDtypeStruct((rows, hidden), gathered.dtype), | ||
| jax.ShapeDtypeStruct((rows // _LANES, _LANES), jnp.float32), | ||
| ], | ||
| compiler_params=pltpu.CompilerParams( | ||
| dimension_semantics=("arbitrary", "arbitrary"), | ||
| vmem_limit_bytes=48 * 1024 * 1024, | ||
| ), | ||
| cost_estimate=pl.CostEstimate( | ||
| flops=4 * rows * hidden, | ||
| transcendentals=0, | ||
| bytes_accessed=3 * rows * hidden * gathered.dtype.itemsize + 8 * rows, | ||
| ), | ||
| interpret=interpret, | ||
| name="unsort_bwd_rows", | ||
| )( | ||
| jnp.reshape(limit.astype(jnp.int32), (1,)), | ||
| gathered, | ||
| sorted_tokens, | ||
| jnp.reshape(weights.astype(jnp.float32), (rows // _LANES, _LANES)), | ||
| ) | ||
| return grad, jnp.reshape(dot, (rows,)) |
There was a problem hiding this comment.
To support the race-free column-block reduction, we need to update _unsort_bwd_rows_pallas to define separate input and output block specs for the lane-dense tensors, declare the output shape of dot with a num_col_blocks dimension, and sum the partial dots across column blocks in JAX after the pallas_call.
def _unsort_bwd_rows_pallas(gathered, sorted_tokens, weights, limit, *, interpret=False):\n \"\"\"The Pallas TensorCore kernel; see `unsort_bwd_rows`.\"\"\"\n rows, hidden = gathered.shape\n block_cols = _block_cols(hidden)\n num_row_blocks = pl.cdiv(rows, _BLOCK_ROWS)\n num_col_blocks = hidden // block_cols\n\n def data_map(i, j, limit_ref):\n last = _last_block(limit_ref[0])\n return jnp.minimum(i, last), jnp.where(i <= last, j, num_col_blocks - 1)\n\n def row_map_in(i, j, limit_ref):\n del j\n return jnp.minimum(i, _last_block(limit_ref[0])), 0\n\n def row_map_out(i, j, limit_ref):\n return j, jnp.minimum(i, _last_block(limit_ref[0])), 0\n\n data_spec = pl.BlockSpec((_BLOCK_ROWS, block_cols), data_map)\n lane_dense_in_spec = pl.BlockSpec((_BLOCK_ROWS // _LANES, _LANES), row_map_in)\n lane_dense_out_spec = pl.BlockSpec((1, _BLOCK_ROWS // _LANES, _LANES), row_map_out)\n grad, dot = pl.pallas_call(\n _unsort_bwd_rows_kernel,\n grid_spec=pltpu.PrefetchScalarGridSpec(\n num_scalar_prefetch=1,\n grid=(num_row_blocks, num_col_blocks),\n in_specs=[data_spec, data_spec, lane_dense_in_spec],\n out_specs=[data_spec, lane_dense_out_spec],\n ),\n out_shape=[\n jax.ShapeDtypeStruct((rows, hidden), gathered.dtype),\n jax.ShapeDtypeStruct((num_col_blocks, rows // _LANES, _LANES), jnp.float32),\n ],\n compiler_params=pltpu.CompilerParams(\n dimension_semantics=(\"arbitrary\", \"arbitrary\"),\n vmem_limit_bytes=48 * 1024 * 1024,\n ),\n cost_estimate=pl.CostEstimate(\n flops=4 * rows * hidden,\n transcendentals=0,\n bytes_accessed=3 * rows * hidden * gathered.dtype.itemsize + 8 * rows,\n ),\n interpret=interpret,\n name=\"unsort_bwd_rows\",\n )(\n jnp.reshape(limit.astype(jnp.int32), (1,)),\n gathered,\n sorted_tokens,\n jnp.reshape(weights.astype(jnp.float32), (rows // _LANES, _LANES)),\n )\n return grad, jnp.reshape(jnp.sum(dot, axis=0), (rows,))
Comment on lines
+64
to
+66
| def kernel_supported(fused): | ||
| rows, width = fused.shape | ||
| return fused.dtype == jnp.bfloat16 and rows % 16 == 0 and width % 256 == 0 |
There was a problem hiding this comment.
Add a defensive check to ensure fused is a 2D array before unpacking its shape. Since fused_glu can theoretically be called with N-dimensional arrays, this prevents unexpected ValueError unpacking crashes.
def kernel_supported(fused):\n if fused.ndim != 2:\n return False\n rows, width = fused.shape\n return fused.dtype == jnp.bfloat16 and rows % 16 == 0 and width % 256 == 0
Comment on lines
+54
to
+64
| def _kernel_supported(gathered, sorted_tokens, weights): | ||
| rows, hidden = gathered.shape | ||
| return ( | ||
| gathered.dtype == jnp.bfloat16 | ||
| and sorted_tokens.dtype == jnp.bfloat16 | ||
| and sorted_tokens.shape == gathered.shape | ||
| and weights.shape == (rows,) | ||
| and rows % _LANES == 0 | ||
| and hidden % _LANES == 0 | ||
| and hidden % _block_cols(hidden) == 0 | ||
| ) |
There was a problem hiding this comment.
Add defensive checks to ensure gathered and sorted_tokens are 2D arrays before unpacking their shapes.
def _kernel_supported(gathered, sorted_tokens, weights):\n if gathered.ndim != 2 or sorted_tokens.ndim != 2:\n return False\n rows, hidden = gathered.shape\n return (\n gathered.dtype == jnp.bfloat16\n and sorted_tokens.dtype == jnp.bfloat16\n and sorted_tokens.shape == gathered.shape\n and weights.shape == (rows,)\n and rows % _LANES == 0\n and hidden % _LANES == 0\n and hidden % _block_cols(hidden) == 0\n )
Comment on lines
+52
to
+64
| @pytest.mark.parametrize("limit", [0, 1000, 1700, _ROWS]) | ||
| def test_unsort_bwd_rows(limit): | ||
| """Weighted gradient bit-identical and row dot to f32 summation order on rows < limit, zeros up to the block end.""" | ||
| keys = jax.random.split(jax.random.PRNGKey(0), 3) | ||
| gathered, tokens = _bf16(keys[0], (_ROWS, _HIDDEN)), _bf16(keys[1], (_ROWS, _HIDDEN)) | ||
| weights = jax.random.uniform(keys[2], (_ROWS,), jnp.float32, -1.0, 1.0) | ||
| interpret = True if jax.default_backend() == "cpu" else None | ||
| grad, dot = unsort.unsort_bwd_rows(gathered, tokens, weights, jnp.int32(limit), interpret=interpret) | ||
| ref_grad, ref_dot = unsort.unsort_bwd_rows_reference(gathered, tokens, weights, jnp.int32(limit)) | ||
| block_end = min(_ROWS, max(-(-limit // 1024), 1) * 1024) | ||
| np.testing.assert_array_equal(_bits(grad, block_end), _bits(ref_grad, block_end)) | ||
| np.testing.assert_allclose(np.asarray(dot)[:limit], np.asarray(ref_dot)[:limit], rtol=0, atol=1e-4) | ||
| np.testing.assert_array_equal(np.asarray(dot)[limit:block_end], 0.0) |
There was a problem hiding this comment.
The current unit test hardcodes _HIDDEN = 256, which always results in num_col_blocks = 1. Parametrizing the test with both 256 and 4096 hidden dimensions will ensure that the multi-column-block reduction path (num_col_blocks > 1) is thoroughly tested and verified.
@pytest.mark.parametrize(\"limit\", [0, 1000, 1700, _ROWS])\n@pytest.mark.parametrize(\"hidden\", [256, 4096])\ndef test_unsort_bwd_rows(limit, hidden):\n \"\"\"Weighted gradient bit-identical and row dot to f32 summation order on rows < limit, zeros up to the block end.\"\"\"\n keys = jax.random.split(jax.random.PRNGKey(0), 3)\n gathered, tokens = _bf16(keys[0], (_ROWS, hidden)), _bf16(keys[1], (_ROWS, hidden))\n weights = jax.random.uniform(keys[2], (_ROWS,), jnp.float32, -1.0, 1.0)\n interpret = True if jax.default_backend() == \"cpu\" else None\n grad, dot = unsort.unsort_bwd_rows(gathered, tokens, weights, jnp.int32(limit), interpret=interpret)\n ref_grad, ref_dot = unsort.unsort_bwd_rows_reference(gathered, tokens, weights, jnp.int32(limit))\n block_end = min(_ROWS, max(-(-limit // 1024), 1) * 1024)\n np.testing.assert_array_equal(_bits(grad, block_end), _bits(ref_grad, block_end))\n np.testing.assert_allclose(np.asarray(dot)[:limit], np.asarray(ref_dot)[:limit], rtol=0, atol=1e-4)\n np.testing.assert_array_equal(np.asarray(dot)[limit:block_end], 0.0)This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #5759. Draft, pending a bench run.
mlp_moe_out/embed_moe_outkeep the old layout.