diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 449a84d899..7411e14975 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -884,6 +884,35 @@ class MoEGeneral(BaseModel): False, description="Whether to use Ring of Experts for sparse matmul expert parallelism.", ) + moe_bwd_inkernel_quant: bool = Field( + False, + description=( + "Quantize the MoE backward-gmm operands INSIDE the ragged kernels instead of with dense " + "XLA-level quantize ops. The dense quantize/amax ops process every row of the ragged " + "buffer at its STATIC size (worst-case at ragged_buffer_factor<=0); the in-kernel path " + "touches only the valid group_sizes rows/tiles, and no reduction ever reads " + "uninitialized buffer tail rows. drhs: tgmm quantizes BOTH operands in-kernel " + "(per-gm-tile-per-channel e4m3). Requires use_tokamax_gmm + use_gmm_v2 + an fp8 qwix " + "bwd_qtype; falls back to the XLA quantize otherwise. For performance improvement at " + "dropless (worst-case) ragged buffer sizes." + ), + ) + bwd_quantization_dtype: Literal["e5m2", "e4m3"] = Field( + "e5m2", + description=( + "fp8 dtype for the BACKWARD (gradient) quantization in the fp8_full qwix recipe: 'e5m2' " + "(default) or 'e4m3'." + ), + ) + moe_ring_cotangent_ag: bool = Field( + False, + description=( + "Run the BACKWARD cotangent all-gather of the ring-of-experts combine reduce-scatter on " + "a TensorCore Pallas ring kernel instead of the XLA collective (which can serialize on " + "the SparseCore collective-offload queue). Forward is unchanged. For performance " + "improvement; numerically equal to lax.all_gather." + ), + ) moe_quantize_token_all_gather: bool = Field( False, description="Whether to quantize token activations to FP8 before All-Gather across EP shards in Ring of Experts.", @@ -1321,6 +1350,15 @@ class RematAndOffload(BaseModel): RematLocation.REMAT, description="Remat policy for the first part of a gated MoE's output.", ) + moe_x_sorted: RematLocation = Field( + RematLocation.REMAT, + description=( + "Remat policy for the routed (post-dispatch, expert-sorted) MoE input plus its small " + "routing/metadata bundle. 'device' saves them across the remat boundary so the backward " + "does not re-run the dispatch token all-gather and sort; the expert GMMs re-run from the " + "saved tensor. Default 'remat' recomputes (existing behavior)." + ), + ) moe_mlpwi_1: RematLocation = Field( RematLocation.REMAT, description="Remat policy for the second part of a gated MoE's output.", @@ -3423,6 +3461,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "context", "mlpwi", "moe_mlpwi_0", + "moe_x_sorted", "moe_mlpwi_1", "moe_mlpwo", "mlpwi_0", @@ -4564,6 +4603,7 @@ def set_derived_values_and_validate(self) -> "RLConfig": "context", "mlpwi", "moe_mlpwi_0", + "moe_x_sorted", "moe_mlpwi_1", "moe_mlpwo", "mlpwi_0", diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index 886f7386be..a186e5207f 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -76,6 +76,7 @@ def gmm( use_manual_quantization: bool = False, # used in batchsplit use_gmm_v2: bool = False, partial_sum: jnp.ndarray | None = None, + bwd_inkernel_quant: bool = False, ): """Grouped matrix multiplication operation.""" if interpret is None: @@ -106,7 +107,7 @@ def gmm( gmm_fwd_bwd = lambda *args: _gmm_fwd(*args)[0] # pylint: disable=C3001 gmm_fwd_bwd = jax.custom_vjp( gmm_fwd_bwd, - nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15), + nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15, 17), ) gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs.dtype, rhs.dtype)) return gmm_fwd_bwd( @@ -127,6 +128,7 @@ def gmm( rhs_vma_axes, use_gmm_v2, partial_sum, + bwd_inkernel_quant, ) @@ -163,6 +165,7 @@ def _gmm_fwd( rhs_vma_axes: tuple = tuple(), use_gmm_v2: bool = False, partial_sum: jnp.ndarray | None = None, + bwd_inkernel_quant: bool = False, ) -> tuple[ jnp.ndarray, tuple[ @@ -462,6 +465,7 @@ def _gmm_bwd( lhs_vma_axes: tuple, rhs_vma_axes: tuple, use_gmm_v2: bool, + bwd_inkernel_quant: bool, residual: tuple[ jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray, @@ -493,14 +497,44 @@ def _gmm_bwd( # - dlhs_dout: the incoming gradient used to calculate dlhs. # - drhs_dout: the incoming gradient used to calculate drhs. + # moe_bwd_inkernel_quant: run the drhs tgmm with BOTH operands quantized in-kernel + # (per-gm-tile-per-channel), touching only rows covered by group_sizes. This subsumes three + # dense buffer-sized XLA ops -- the per-row lhs re-quantize, the drhs_dout *= lhs.scale + # multiply, and the per-N cotangent quantize -- and, because no reduction ever reads past the + # valid rows, removes the NaN hazard of amax over uninitialized ragged-buffer tail rows. + inkernel_drhs = ( + bwd_inkernel_quant + and use_tokamax_backend + and use_gmm_v2 + and quantization_rule is not None + and bool(quantization_rule.bwd_qtype) + ) + # 1. Scale Application & QArray Unwrapping dlhs_dout, drhs_dout, lhs, rhs = _bwd_prepare_inputs( - grad, residual_lhs, residual_rhs, group_sizes, use_gmm_v2, transpose_rhs, quantization_rule + grad, residual_lhs, residual_rhs, group_sizes, use_gmm_v2, transpose_rhs, quantization_rule, + skip_lhs_quant=inkernel_drhs, ) # 2. Backward Pass Quantization if quantization_rule: - dlhs_dout, drhs_dout = _bwd_quantize_gradient(dlhs_dout, drhs_dout, quantization_rule) + if inkernel_drhs: + # the in-kernel tgmm quantizes drhs_dout ITSELF; quantize ONLY the dlhs cotangent here + # (calling the two-sided helper would emit the dense drhs quantize just to discard it). + if quantization_rule.bwd_qtype: + dlhs_dout = qpl.quantize( + # pyrefly: ignore[bad-argument-type] + dlhs_dout, + quantization_rule.bwd_qtype, + channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0], + calibration_method=quantization_rule.bwd_calibration_method, + ) + if not isinstance(drhs_dout, qpl.QArray) and drhs_dout.dtype != lhs.dtype: + # tgmm requires equal operand widths; the in-kernel path reads the RAW cotangent, so + # carry it at the activation width (halves the kernel's cotangent read bytes vs f32). + drhs_dout = drhs_dout.astype(lhs.dtype) + else: + dlhs_dout, drhs_dout = _bwd_quantize_gradient(dlhs_dout, drhs_dout, quantization_rule) # 3. DLHS Gradient Execution dlhs = _compute_dlhs( @@ -534,6 +568,7 @@ def _gmm_bwd( interpret, rhs_vma_axes, quantization_rule, + inkernel_quant=inkernel_drhs, ) # 5. Output Formatting @@ -572,8 +607,14 @@ def _bwd_prepare_inputs( use_gmm_v2: bool, transpose_rhs: bool, quantization_rule: qwix.QtRule | None, + skip_lhs_quant: bool = False, ) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray, jnp.ndarray, jnp.ndarray]: - """Prepares backward operands.""" + """Prepares backward operands. + + `skip_lhs_quant=True` (bwd_inkernel_quant) keeps the lhs as the raw wide array: the drhs + tgmm quantizes BOTH operands in-kernel over valid gm tiles only, so the dense per-row XLA + quantize here (and the drhs_dout *= lhs.scale multiply below) would be buffer-sized overhead. + """ # dlhs_dout and drhs_dout can be different when quantization is enabled. dlhs_dout = grad @@ -597,7 +638,7 @@ def _bwd_prepare_inputs( # GMM2 FWD performs lhs quantization inside kernel, lhs is stored as unquantized dtype # in the residual tuple. In BWD, we explicitly quantize lhs. - if quantization_rule and quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray): + if quantization_rule and quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and not skip_lhs_quant: lhs = qpl.quantize( # pyrefly: ignore[bad-assignment] lhs, quantization_rule.act_qtype, @@ -816,12 +857,16 @@ def _compute_drhs( interpret: bool, rhs_vma_axes: tuple, quantization_rule: qwix.QtRule | None, + inkernel_quant: bool = False, ) -> jnp.ndarray: """Routes execution of DRHS based on backend choices.""" if use_tokamax_backend and not use_gmm_v2: drhs = _drhs_run_tokamax_v1(drhs_dout, lhs, group_sizes, rhs_dtype, use_manual_quantization) elif use_tokamax_backend and use_gmm_v2: - drhs = _drhs_run_tokamax_v2(drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling) + drhs = _drhs_run_tokamax_v2( + drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling, + quantize_operands=inkernel_quant, + ) else: drhs = _drhs_run_megablox( drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling, interpret, rhs_vma_axes @@ -888,6 +933,7 @@ def _drhs_run_tokamax_v2( num_actual_groups: int, rhs_dtype: jax.typing.DTypeLike, tiling: tuple, + quantize_operands: bool = False, ) -> jnp.ndarray: """Executes Tokamax TGMM V2 backend for DRHS = LHS^T @ DRHS_dout.""" drhs_rhs = drhs_dout.qvalue if isinstance(drhs_dout, qpl.QArray) else drhs_dout @@ -909,6 +955,7 @@ def _drhs_run_tokamax_v2( preferred_element_type=rhs_dtype, # pyrefly: ignore[bad-argument-type] group_offset=group_offset, tile_info=custom_drhs_tiling, + quantize_operands=quantize_operands and rhs_scale is None, ) diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py index a5f3f6338d..4b47b7b26a 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py @@ -191,6 +191,7 @@ def make_tgmm_configs( out_dtype: jnp.dtype, acc_dtype: jnp.dtype | None, target_zero_ref_bytes: int, + quantize_operands: bool = False, ): """Fills the GMM config for the TGMM kernel.""" assert out_dtype, "out_dtype cannot be None" @@ -233,15 +234,24 @@ def make_tgmm_configs( size_lhs_sublane=size_lhs_sublane, ) + # moe_bwd_inkernel_quant: quantize_operands puts a per-(gm-tile x channel) dynamic e4m3 + # quantize of BOTH operands inside the inner kernel (signalled to tgmm_inner_kernel via + # lhs_cfgs.quant_dtype). Mutually exclusive with a pre-computed per-N rhs_scale. + if quantize_operands: + assert rhs_scale is None, "quantize_operands is mutually exclusive with rhs_scale" + inkernel_q_dtype = jnp.float8_e4m3fn.dtype + else: + inkernel_q_dtype = None + rhs_quant_block_size_m = size_m rhs_cfgs = gmm_v2.InputConfigs( - quant_dtype=None, + quant_dtype=inkernel_q_dtype, quant_block_size=rhs_quant_block_size_m, dtype=rhs.dtype, has_scale=(rhs_scale is not None), ) lhs_cfgs = gmm_v2.InputConfigs( - quant_dtype=None, + quant_dtype=inkernel_q_dtype, quant_block_size=-1, dtype=lhs.dtype, ) @@ -334,12 +344,43 @@ def _matmul(is_new_group: bool, is_group_changing: bool): rhs_mask = jnp.logical_and(m_start_local <= rhs_iota, rhs_iota < m_end_local) rhs_masked = jnp.where(rhs_mask, tiled_rhs_ref[...], 0) - acc = jax.lax.dot_general( - lhs_masked, - rhs_masked, - (((0,), (0,)), ((), ())), - preferred_element_type=jnp.float32, - ) + if cfgs.lhs_cfgs.quant_dtype is not None: + # moe_bwd_inkernel_quant: quantize BOTH operands in-kernel with per-(gm-tile x channel) + # dynamic scales and rescale the partial product by the scale outer-product before + # accumulation (the tgmm_block per-segment pattern, fused into the tile loop -- no dense + # XLA-level quantize/amax over the ragged buffer, only valid gm tiles pay). The masked + # rows are zero, so they neither perturb the amax nor the product. + q_dtype = cfgs.lhs_cfgs.quant_dtype + dtype_max = float(jnp.finfo(q_dtype).max) + lhs_f = lhs_masked.astype(jnp.float32) + rhs_f = rhs_masked.astype(jnp.float32) + lhs_scale = jnp.max(jnp.abs(lhs_f), axis=0) / dtype_max # [tile_k] f32 + rhs_scale = jnp.max(jnp.abs(rhs_f), axis=0) / dtype_max # [tile_n] f32 + # A near-zero scale would give 0 * inf = NaN. An `== 0` guard is NOT enough: for any column + # whose amax is nonzero but below ~1.3e-36, `1/scale` OVERFLOWS f32 to inf, the guard does not + # fire, and every exactly-zero element in that column becomes 0*inf = NaN. Masked rows are set + # to exactly 0 above, so MORE masked rows = more NaN sites -- imbalanced (real) routing makes + # small groups and mostly-masked tiles, so it is strictly more exposed than a balanced + # synthetic router. Guard on the smallest scale with a finite reciprocal instead. + _recip_min = jnp.float32(1.0) / jnp.finfo(jnp.float32).max + lhs_inv = jnp.where(lhs_scale > _recip_min, 1.0 / lhs_scale, 0.0) + rhs_inv = jnp.where(rhs_scale > _recip_min, 1.0 / rhs_scale, 0.0) + lhs_q = (lhs_f * lhs_inv.reshape(1, -1)).astype(q_dtype) + rhs_q = (rhs_f * rhs_inv.reshape(1, -1)).astype(q_dtype) + acc = jax.lax.dot_general( + lhs_q, + rhs_q, + (((0,), (0,)), ((), ())), + preferred_element_type=jnp.float32, + ) + acc = acc * lhs_scale.reshape(-1, 1) * rhs_scale.reshape(1, -1) + else: + acc = jax.lax.dot_general( + lhs_masked, + rhs_masked, + (((0,), (0,)), ((), ())), + preferred_element_type=jnp.float32, + ) if not is_new_group: acc += acc_ref[...] @@ -642,6 +683,7 @@ def validate_tgmm_inputs( "precision", "preferred_element_type", "acc_dtype", + "quantize_operands", ], ) def tgmm_v2( @@ -658,6 +700,7 @@ def tgmm_v2( precision: jax.lax.Precision = jax.lax.Precision.DEFAULT, preferred_element_type: jnp.dtype | None = None, acc_dtype: jnp.dtype | None = None, + quantize_operands: bool = False, ): """Computes a transposed grouped matrix multiplication. @@ -710,6 +753,7 @@ def tgmm_v2( out_dtype=preferred_element_type, # pyrefly: ignore[bad-argument-type] acc_dtype=acc_dtype, target_zero_ref_bytes=target_zero_ref_bytes, + quantize_operands=quantize_operands, ) dims = cfgs.dims tiles = cfgs.tiles diff --git a/src/maxtext/kernels/ring_ag.py b/src/maxtext/kernels/ring_ag.py new file mode 100644 index 0000000000..40a6bcc4f0 --- /dev/null +++ b/src/maxtext/kernels/ring_ag.py @@ -0,0 +1,191 @@ +"""Bidirectional store-and-forward ring all-gather (TensorCore Pallas, ICI DMA). + +Validated on v7x: tiled semantics bit-exact vs ``jax.lax.all_gather(..., tiled=True)``. Written for +IN-shard_map use: `ring_all_gather` runs the pallas_call directly on the local shard +inside an enclosing shard_map, with a caller-supplied `collective_id` (a fixed id would collide +with other in-flight Pallas collectives' barrier semaphores) and an explicit `CostEstimate` so the latency-hiding scheduler can see the ~2x-bytes +DMA cost instead of treating the custom-call as free (kernels without a cost estimate are +invisible to the scheduler's overlap decisions). + +Why a RING (not the direct-to-owner broadcast `_direct_all_gather` in moe.py): the direct pattern +sends every shard's block over multi-hop paths to all peers (bisection congestion, measured +regressing at EP=8); the ring moves each block over neighbor links only, at the validated +159-179 GB/s. +""" + +import jax +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +MESH = pl.DeviceIdType.MESH + + +def _strides(sizes): + """Tiled chunk strides for axes ordered outermost..innermost. strides[k]=prod(sizes[k+1:]).""" + st = [1] * len(sizes) + for i in range(len(sizes) - 2, -1, -1): + st[i] = st[i + 1] * sizes[i + 1] + return st + + +def _full_index(ref, gather_dim, start, length, split_dim=None, s_start=0, s_len=None): + """Index tuple: whole array, a [start:start+length] window on gather_dim, and (for bidi) a + [s_start:s_start+s_len] window on split_dim. The split is on a NON-gather dim so the gather_dim + offsets stay tile-aligned (splitting the gather/concat dim breaks Mosaic tile alignment).""" + idx = [pl.ds(0, ref.shape[d]) for d in range(ref.ndim)] + idx[gather_dim] = pl.ds(start, length) + if split_dim is not None: + idx[split_dim] = pl.ds(s_start, s_len) + return tuple(idx) + + +def _pick_split_dim(shape, gather_dim): + """First non-gather dim with an even extent (to halve for bidirectional). None -> force uni.""" + for d in range(len(shape)): + if d != gather_dim and shape[d] % 2 == 0: + return d + return None + + +def _neighbor(all_axes, axis, delta): + """device_id MESH dict: step `delta` along `axis` (wrap), hold every other mesh axis fixed.""" + size = lax.axis_size(axis) + nxt = lax.rem(lax.axis_index(axis) + delta + size, size) + return {a: (nxt if a == axis else lax.axis_index(a)) for a in all_axes} + + +def _ring_stage(o_ref, all_axes, axes, sizes, strides, k, chunk, gather_dim, + send_sem, recv_sem, bidi, split_dim, pipe): + """One store-and-forward ring over axes[k]; fills sizes[k] blocks of cur_len rows. + + Bidirectional splits each block along split_dim (a NON-gather dim): +dir carries the upper half, + -dir the lower half, each over the FULL gather block -> gather_dim offsets stay tile-aligned. + + Pipelining: each direction's half is further sliced into `pipe` independent sub-chunks along + split_dim, each with its OWN sem. Store-and-forward forces depth-1 PER sub-chunk (can't forward + what hasn't arrived), but the 2*pipe sub-chunks progress independently and are issued + round-robin, so at steady state 2*pipe DMAs are in flight. pipe=1 is the baseline.""" + s = sizes[k] + if s == 1: + return + cur_rows = strides[k] * chunk # gather-dim block length in rows + # Base row of this axis-k group = (sum_{j 0 else 0) * chunk + my = lax.axis_index(axes[k]) + use_bidi = bidi and split_dim is not None + + right_n = _neighbor(all_axes, axes[k], +1) + left_n = _neighbor(all_axes, axes[k], -1) + bsem = pltpu.get_barrier_semaphore() + pl.semaphore_signal(bsem, inc=1, device_id=right_n, device_id_type=MESH) + if use_bidi: + pl.semaphore_signal(bsem, inc=1, device_id=left_n, device_id_type=MESH) + pl.semaphore_wait(bsem, 2 if use_bidi else 1) + + # Pipeline "lanes": (sem_idx, neighbor, hop_sign, split_start, split_len). hop_sign=-1 forwards + # source-coord (my-i) to the +1 neighbor; +1 forwards (my+i) to the -1 neighbor. + lanes = [] + if split_dim is not None: + E = o_ref.shape[split_dim] + if use_bidi: + half = E // 2 + assert half % pipe == 0, (E, pipe) + w = half // pipe + for j in range(pipe): + lanes.append((j, right_n, -1, half + j * w, w)) # +dir, upper half + lanes.append((pipe + j, left_n, +1, j * w, w)) # -dir, lower half + else: + assert E % pipe == 0, (E, pipe) + w = E // pipe + for j in range(pipe): + lanes.append((j, right_n, -1, j * w, w)) # uni, full split sliced into pipe + else: + lanes.append((0, right_n, -1, None, None)) # uni, no split_dim: single block + + def idx(start, ss, sl): + if ss is None: + return _full_index(o_ref, gather_dim, start, cur_rows) + return _full_index(o_ref, gather_dim, start, cur_rows, split_dim=split_dim, s_start=ss, s_len=sl) + + def block_start(u): + return base_rows + lax.rem(u + s, s) * cur_rows + + prev = [None] * len(lanes) + for i in range(s - 1): + for li, (si, nb, sign, ss, sl) in enumerate(lanes): + start = block_start(my - i if sign < 0 else my + i) + if prev[li] is not None: + prev[li].wait() + prev[li] = pltpu.async_remote_copy( + o_ref.at[idx(start, ss, sl)], o_ref.at[idx(start, ss, sl)], + send_sem.at[si], recv_sem.at[si], device_id=nb, device_id_type=MESH) + for d in prev: + if d is not None: + d.wait() + + +def _kernel(w_ref, o_ref, send_sem, recv_sem, local_sem, *, + all_axes, axes, sizes, strides, chunk, gather_dim, bidi, split_dim, pipe): + # 1. place our own shard at its tiled slot: chunk index = sum coord_j * strides[j]. + my_chunk = sum(lax.axis_index(axes[j]) * strides[j] for j in range(len(axes))) + my_start = my_chunk * chunk + local_copy = pltpu.make_async_copy( + w_ref.at[_full_index(w_ref, gather_dim, 0, chunk)], + o_ref.at[_full_index(o_ref, gather_dim, my_start, chunk)], + local_sem) + local_copy.start() + local_copy.wait() + # 2. nested rings, innermost axis first (largest k -> smallest stride). + for k in range(len(axes) - 1, -1, -1): + _ring_stage(o_ref, all_axes, axes, sizes, strides, k, chunk, gather_dim, + send_sem, recv_sem, bidi, split_dim, pipe) + + +def ring_all_gather(x, mesh, gather_axes, gather_dim, collective_id, *, bidi=True, pipe=1): + """In-shard_map ring all-gather of the LOCAL shard `x` over `gather_axes` (mesh axis names, + outermost..innermost), tiled on `gather_dim`. Semantics == + ``jax.lax.all_gather(x, gather_axes, axis=gather_dim, tiled=True)`` (bit-exact, pure data move). + + MUST be called INSIDE a shard_map spanning `mesh` (uses lax.axis_index on the mesh axes for MESH + device_id addressing). `collective_id` selects the barrier semaphore and must be DISTINCT from + every other concurrently in-flight Pallas collective's id.""" + if isinstance(gather_axes, str): + gather_axes = (gather_axes,) + sizes = tuple(mesh.shape[ax] for ax in gather_axes) + strides = _strides(sizes) + n_total = 1 + for s_ in sizes: + n_total *= s_ + chunk = x.shape[gather_dim] # local rows on the gather dim + out_shape = list(x.shape) + out_shape[gather_dim] = chunk * n_total + out_shape = tuple(out_shape) + all_axes = tuple(mesh.axis_names) + split_dim = _pick_split_dim(out_shape, gather_dim) if bidi else None + HBM = pltpu.MemorySpace.HBM + + def kern(w_ref, o_ref, send_sem, recv_sem, local_sem): + _kernel(w_ref, o_ref, send_sem, recv_sem, local_sem, + all_axes=all_axes, axes=tuple(gather_axes), sizes=sizes, strides=strides, + chunk=chunk, gather_dim=gather_dim, bidi=bidi, split_dim=split_dim, pipe=pipe) + + nsem = 2 * pipe # 2 directions x pipe sub-chunks (uni uses the first `pipe`) + # Cost estimate: each device receives + forwards ~2x the full gathered bytes over the ring. + # Without it the custom-call is invisible to the latency-hiding scheduler's overlap decisions. + full_bytes = 1 + for d_ in out_shape: + full_bytes *= d_ + full_bytes *= x.dtype.itemsize + return pl.pallas_call( + kern, + out_shape=jax.ShapeDtypeStruct(out_shape, x.dtype), + in_specs=[pl.BlockSpec(memory_space=HBM)], + out_specs=pl.BlockSpec(memory_space=HBM), + scratch_shapes=[pltpu.SemaphoreType.DMA((nsem,)), # send, per (dir,sub) lane + pltpu.SemaphoreType.DMA((nsem,)), # recv, per (dir,sub) lane + pltpu.SemaphoreType.DMA], # local copy + compiler_params=pltpu.CompilerParams(collective_id=collective_id), + cost_estimate=pl.CostEstimate(flops=0, bytes_accessed=2 * full_bytes, transcendentals=0), + )(x) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 1d2cf2e959..7fc6b0c517 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -31,6 +31,7 @@ from jax.sharding import Mesh, NamedSharding from jax.sharding import PartitionSpec as P from maxtext.common import common_types as ctypes +from maxtext.kernels.ring_ag import ring_all_gather from maxtext.common.common_types import ShardMode from maxtext.kernels import megablox as mblx from maxtext.layers import attentions, linears, nnx_wrappers, quantizations @@ -57,6 +58,41 @@ COMBINE = "combine" +# Distinct barrier-semaphore ids for the in-MoE Pallas ring collectives; must not collide with +# any other collective_id in flight in the same program. +_RING_CT_AG_COLLECTIVE_ID = 55 # moe_ring_cotangent_ag: backward combine-cotangent ring all-gather + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3)) +def _ring_ct_reduce_scatter(output, mesh, ep_name, collective_id): + """Combine reduce-scatter whose BACKWARD cotangent all-gather runs on the TC ring kernel. + + FORWARD: byte-identical to the stock path -- the plain + ``jax.lax.psum_scatter(output, ep_name, scatter_dimension=0, tiled=True)`` (also what a remat + recompute re-traces: the primal is the plain collective, so no Pallas DMA ever runs inside a + rematted region). BACKWARD: the autodiff transpose of the tiled psum_scatter is a tiled EP + all-gather of the combine cotangent; as an XLA collective it can ride the SparseCore + collective-offload queue and serialize behind other SC work. Here it runs on the bidirectional + store-and-forward TensorCore ring kernel instead (ICI DMAs issued from the TC, where the + backward has slack), numerically == ``lax.all_gather`` (pure tiled data move, bit-exact). + """ + return jax.lax.psum_scatter(output, ep_name, scatter_dimension=0, tiled=True) + + +def _ring_ct_rs_fwd(output, mesh, ep_name, collective_id): + return _ring_ct_reduce_scatter(output, mesh, ep_name, collective_id), None + + +def _ring_ct_rs_bwd(mesh, ep_name, collective_id, _res, ct): + return (ring_all_gather(ct, mesh, (ep_name,), 0, collective_id),) + + +_ring_ct_reduce_scatter.defvjp(_ring_ct_rs_fwd, _ring_ct_rs_bwd) + + + + + @struct.dataclass class RouteMetadata: """EP communication state needed to undo the forward all-to-all after expert computation.""" @@ -1598,6 +1634,7 @@ def extract_vma(tensor): use_gmm_v2=self.config.use_gmm_v2, partial_sum=partial_sum, interpret=megablox_interpret, + bwd_inkernel_quant=getattr(self.config, "moe_bwd_inkernel_quant", False), ) else: # jax.lax.ragged_dot @@ -2358,6 +2395,18 @@ def _moe_body( x, routing, route_metadata = route( x, logits, pre_bias_logits, rngs, input_ids=sharded_input_ids ) + # moe_x_sorted: tag the routed expert input and its small routing/metadata bundle for + # the remat policy. With moe_x_sorted=device the backward LOADS these instead of + # re-running route() -- removing the rematted dispatch token all-gather and sort from the + # backward (the up-projection weight gradient needs the sorted input anyway). The + # routing/metadata leaves (indices, group sizes, weights -- tiny) must be saved too, else + # the sort re-runs just to reproduce them. Tags are inert under the default + # moe_x_sorted=remat. + _cn = lambda t: adc.checkpoint_name(t, "moe_x_sorted") if isinstance(t, jax.Array) else t + # tree.map so a QArray x (moe_quantize_token_all_gather) gets its qvalue/scale leaves tagged + x = jax.tree.map(_cn, x) + routing = jax.tree.map(_cn, routing) + route_metadata = jax.tree.map(_cn, route_metadata) if self.config.mlp_bias: w0_bias, w1_bias, wo_bias = self.transform_bias( @@ -2411,12 +2460,20 @@ def _moe_body( self.moe_expert_input_dim // self.get_tensor_parallelism_size(), ), ) - output = jax.lax.psum_scatter( - output, - self._expert_parallelism_name, - scatter_dimension=0, - tiled=True, - ) + if ( + getattr(self.config, "moe_ring_cotangent_ag", False) + and isinstance(self._expert_parallelism_name, str) + ): + output = _ring_ct_reduce_scatter( + output, self.mesh, self._expert_parallelism_name, _RING_CT_AG_COLLECTIVE_ID + ) + else: + output = jax.lax.psum_scatter( + output, + self._expert_parallelism_name, + scatter_dimension=0, + tiled=True, + ) return output, routing.lb_loss, routing.bias_updates if self.get_expert_parallelism_size() > 1: diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index a275a0afa8..9b2b1b6a96 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -809,7 +809,9 @@ def get_fp8_full_qwix_rule_w_sparsity(config: Config): module_path="decoder/.*layers.*", weight_qtype=jnp.float8_e4m3fn, act_qtype=jnp.float8_e4m3fn, - bwd_qtype=jnp.float8_e5m2, + # bwd_quantization_dtype: e5m2 (default, wider exponent range) or e4m3 (finer + # mantissa; some fp8 recipes use e4m3 gradients). + bwd_qtype=jnp.float8_e4m3fn if config.bwd_quantization_dtype == "e4m3" else jnp.float8_e5m2, weight_calibration_method=config.weight_quantization_calibration_method, act_calibration_method=config.act_quantization_calibration_method, bwd_calibration_method=config.bwd_quantization_calibration_method,