Fix ring ragged unsort gradient and clamp local group size.

This change fixes issues in the ragged sort implementation when using ring of experts:
- In `moe.py`, clamp `local_group_size` to `buffer_size` to prevent overflow.
- In `ragged_sort.py`, apply mask to `grad_sorted_tokens` in `_ring_ragged_unsort` backward pass to ensure gradients outside shard boundaries are zeroed out.
- In `ops.py` (megablox), add `group_offset` to `_gmm_fwd` and `_gmm_bwd`.

PiperOrigin-RevId: 949759227

Author: RissyRan

src/maxtext/kernels/megablox/ops.py

@@ -204,6 +204,7 @@ def _gmm_fwd(
         ),
         preferred_element_type=preferred_element_type,
         partial_sum=partial_sum,
+        group_offset=group_offset,
     )
   elif use_tokamax_backend:
     # manual_axis_type is for gmm with shard_map check_vma=True, needs tokamax > 0.0.12
@@ -331,6 +332,7 @@ def _gmm_bwd(
               tile_n=tiling[5],
           ),
           preferred_element_type=lhs_dtype,
+          group_offset=group_offset,
       )
 
       # tgmm_v2_op requires lhs and rhs to have the same dtype.

src/maxtext/kernels/ragged/ragged_sort.py

@@ -386,6 +386,10 @@ def _ring_ragged_unsort_bwd(res, g_out):
           flops_override=gather_flops_override,
           bytes_accessed_override=gather_bytes_accessed_override,
       )
+      # Mask out gradients that correspond to elements outside the valid shard
+      # output range.
+      mask = (jnp.arange(n) >= shard_output_start) & (jnp.arange(n) < shard_output_end)
+      grad_sorted_tokens = jnp.where(mask[:, None], grad_sorted_tokens, 0.0)
     else:
       # Slice the inverse permutation to match the packed local buffer.
       padded_idx_inv = jnp.pad(idx_inv, (0, buffer_size))
@@ -405,6 +409,10 @@ def _ring_ragged_unsort_bwd(res, g_out):
           flops_override=gather_flops_override,
           bytes_accessed_override=gather_bytes_accessed_override,
       )
+      # Mask out gradients for elements beyond the valid limit of the local buffer.
+      limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size)
+      mask = jnp.arange(buffer_size) < limit
+      grad_sorted_tokens = jnp.where(mask[:, None], grad_sorted_tokens, 0.0)
     return grad_sorted_tokens, None, None, None
 
   _ring_ragged_unsort.defvjp(_ring_ragged_unsort_fwd, _ring_ragged_unsort_bwd)

src/maxtext/layers/moe.py

@@ -919,6 +919,9 @@ def permute(
           local_num_experts,
           axis=0,
       )
+      # Clamp local_group_size to buffer_size to ensure we don't exceed buffer
+      # capacity by leveraging the helper _truncate_matrix.
+      local_group_size = _truncate_matrix(local_group_size[:, None], buffer_size)[:, 0]
       expert_indices = jnp.arange(local_num_experts)
       sorted_experts = jnp.repeat(
           expert_indices,