Add weight for ragged gather kernel and enable fan out in bwd ragged sort

[NuojCheng]

Description

This PR fixes the usage of ragged_gather kernel under _ring_ragged_unsort_bwd. As a result, it

• Removes the bitcast_add in backward pass
• Reduce HBM usage

Previously, this kernel was applied in non-fan-out way, since the backward pass the unsort requires weight balance, while the ragged gather kernel did not support. This PR adds this support and enable a fan-out style of this kernel, successfully avoid materializing the full padded-shape tensor in fp32.

FIXES: b/522526813

Tests

• Performance:
  • Not use ragged buffer: 17.79 s before 17.80s after. HBM usage gets 5GB lower.
  • Use ragged buffer: 17.53 s before 17.53s after. HBM usage gets 14GB lower.
• Solve ragged buffer size memory issue: EP=8 and EP=16 successful compilation https://paste.googleplex.com/6743176042905600
• EP=8 and FSDP=64 use ragged_buffer_factor=3.0: xprof

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 77.77778% with 12 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/kernels/ragged/ragged_gather.py	71.42%	11 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[gobbleturk]

this is a bit strange, why does emb dim matter?

[NuojCheng]

ragged gather reduce kernel hardcoded num of column partitions being 8. Since the partition size should be multiple of 128, it means emb size should be multiple of 1024 for using this kernel. I updated the emb dim in test case for better demonstrating this knowledge, see context in

maxtext/src/maxtext/kernels/ragged/ragged_gather_reduce.py

Lines 434 to 458 in 61c225f

	# This kernel partitions the output's columns into `num_column_partitions` and
	# partition the output's rows into `num_row_partitions` and run each
	# {row_partition} x {column_partition} combination on a separate SC subcore
	# for parallelism. With such work partitioning, we guarantee that there won't
	# be write collision (from different subcores) to the any output row X column.
	#
	# Each column partition should be multiple of 128 (number of lanes) due to
	# DMA requirements. Unless requiring padding on the column dimension, larger
	# column partitions (thus smaller row partitions given fixed num_cores) is
	# more preferable because large row partition may lead to imbalanced load
	# (valid_rows_mask may have more rows in some partitions than others).
	# Most LLM's hidden size is multiple of 1024, `num_column_partitions=8` should
	# work well in practice without requiring padding on the column size.
	num_column_partitions = 8
	assert num_cores % num_column_partitions == 0
	num_rows_partitions = num_cores // num_column_partitions
	aligned_hidden_size = _align_to(hidden_size, 128 * num_column_partitions)
	col_size = aligned_hidden_size // num_column_partitions
	row_tile_size = num_simd_lanes
	padded_input_size = _align_to(
	input_size,
	math.lcm(num_rows_partitions * row_tile_size, reduce_group_size),
	)
	pad_input_size = padded_input_size - input_size

.

[gobbleturk]

hmm this is a hard constraint for a lot of models, would this fail fast with a clear error message on these sizes? Ideally we can somehow support

[gobbleturk]

I would create a bug to track

[NuojCheng]

Agree, added in b/524661949