[AkshitaB]

These are the changes that let MaxText replay the OLMo 3.5 hero run (OLMo-core olmo35-small-hero-20260907-emo, which is MaxText's olmo35-tiny) step for step on TPU v4. The replay starts from the hero's step-0 weights and trains on its batches in the same order, with its recipe. Every new behavior is behind a flag whose default keeps MaxText's current behavior.

Changes

Matching the OLMo-core recipe

• moe_global_load_balancing: OLMo-core's global load balancing. Inside the routed MoE's shard_map, each device used to balance only its own tokens, and the logged moe_lb_loss was device 0's value. With the flag, expert counts are averaged across devices with no gradient through them, and so is the loss. A new test on 8 simulated devices checks that the value and gradient equal the whole-batch loss.
• olmo_segment_ids_from_eos and olmoe3_mixers_ignore_segment_ids: EMO expert pools per EOS-delimited document, as in OLMo-core. Attention and KDA still see the whole instance, since the hero has no document masking.
• moe_aux_loss_reduction=sum: OLMo-core sums the routers' auxiliary losses over layers; MaxText averages them. At the same weight that was 15× less balancing gradient for this model.
• warmup_step_offset=1: warmup counts steps from one, as OLMo-core does.
• olmo_global_indices_path (olmo_grain loader): reads OLMo-core's global-indices file, so global batch k is exactly OLMo-core's batch k.
• Gradient accumulation reports z_loss and the MoE aux losses per micro-batch. They were logged gradient_accumulation_steps times too large.

KDA numerics

• The chunk system is now inverted by forward substitution, replacing Newton-Schulz and the power series. On trained keys the old inverse overflowed, which caused the NaN at hero-replay steps 210–220. The q/k L2-norm epsilon now matches FLA.
• Includes Fix sub-block KDA inverse miscompile on TPU v4 #5665, the sub-block inverse miscompile fix on TPU v4. That PR is still open against this base; this commit drops out of the diff if it merges first.

Memory with gradient accumulation (v4 8x8x16, EP 4 × FSDP 256)

• vocab_tiling_seq_major: splits vocab tiles along the sequence instead of the flattened batch. That avoids a full all-gather of the micro-batch's hidden states. With GA=2, temporaries drop from 49.9 GB to 26.3 GB, so the run fits.
• ga_params_loop_barrier (opt-in, off by default): stops XLA from hoisting loop-invariant work on the params out of the GA loop. On its own it takes temporaries from 49.9 GB to 44.3 GB.
• The gradient-sharding constraint inside the GA scan didn't reduce memory on its own. It's kept because it's harmless and makes the intended layout explicit.

Results

The first 1000 steps on a v4 8x8x16 (1024 chips, GA 2, dropless) track the hero:

• CE: within 0.4% in every window. The largest single-step gap is 0.062; at step 1000, 2.7425 vs 2.7389.
• z-loss and LR: match.
• Still open: the aux loss matches at step 0 but runs 11–72% above the hero after that, and the grad norm is 7–32% high. Under investigation.

Tests

On CPU:

• olmoe3_test.py: 37 passed, 6 skipped (TPU-only).
• olmo_grain loader tests: 33 passed.
• olmoe3_global_lb_test.py: passes.
• gradient_accumulation_nnx_test.py: 8 passed.
• tiling_test.py CPU tests, including the new seq-major test: pass on 1 and 8 devices.
• configs_value tests: 51 passed.

🤖 Generated with Claude Code

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request introduces several features and stability fixes to align MaxText with OLMo-core, including support for global load balancing, custom document segmentation from EOS tokens, sequence-major vocabulary tiling, and a more stable unit lower triangular matrix inversion using blocked forward substitution to prevent NaN losses. However, the review identified a critical shape mismatch bug in the ShiftToInputsTargets transform where inputs_seg is computed using the full tokens array (length L + 1) instead of inputs (length L), which will cause a runtime crash when segment_ids_from_eos is enabled. A corresponding assertion error in the unit tests was also flagged.

[gemini-code-assist]

There is a shape mismatch bug here. tokens has length L + 1 (since inputs = tokens[:-1] has length L). Therefore, inputs_seg computed as (np.cumsum(eos) - eos + 1) will have length L + 1. However, inputs_segmentation is expected to have length L to match the length of inputs, inputs_position, and targets_segmentation. This will cause a shape mismatch crash at runtime when segment_ids_from_eos is enabled. You should slice inputs_seg to length L (i.e., [:-1]) or compute it on inputs directly.

Suggested change

	if self.segment_ids_from_eos is None:
	inputs_seg = np.ones(L, dtype=np.int32)
	else:
	eos = (tokens == self.segment_ids_from_eos).astype(np.int32)
	inputs_seg = (np.cumsum(eos) - eos + 1).astype(np.int32)
	if self.segment_ids_from_eos is None:
	inputs_seg = np.ones(L, dtype=np.int32)
	else:
	eos = (inputs == self.segment_ids_from_eos).astype(np.int32)
	inputs_seg = (np.cumsum(eos) - eos + 1).astype(np.int32)

[gemini-code-assist]

The test assertions here incorrectly assume that inputs_segmentation and targets_segmentation have length 9 (the length of arr). However, since inputs is tokens[:-1], the length L is 8. Thus, inputs_segmentation and targets_segmentation must have length 8. Comparing them to arrays of length 9 will fail with a shape mismatch once the production bug in ShiftToInputsTargets is fixed.

Suggested change

	def test_shift_to_inputs_targets_segment_ids_from_eos(self):
	"""Documents numbered as OLMo-core's segment_ids_from_eos (EOS ends its own document), plus one."""
	eos = 100257
	arr = np.array([eos, 5, 6, eos, 7, eos, eos, 8, 9], dtype=np.uint32)
	out = ShiftToInputsTargets(segment_ids_from_eos=eos).map(
	{"tokens": arr, "instance_id": 0, "file_id": 0, "instance_mask": True}
	)
	# OLMo-core: eos.cumsum() - eos = [0, 1, 1, 1, 2, 2, 3, 4, 4]; MaxText reserves 0 for padding.
	np.testing.assert_array_equal(out["inputs_segmentation"], np.array([1, 2, 2, 2, 3, 3, 4, 5, 5], dtype=np.int32))
	# The loss mask is the default one.
	np.testing.assert_array_equal(out["targets_segmentation"], np.array([1] * 8 + [0], dtype=np.int32))
	default = ShiftToInputsTargets().map({"tokens": arr, "instance_id": 0, "file_id": 0, "instance_mask": True})
	for key in ("inputs", "targets", "inputs_position", "targets_segmentation"):
	np.testing.assert_array_equal(out[key], default[key])
	def test_shift_to_inputs_targets_segment_ids_from_eos(self):
	"""Documents numbered as OLMo-core's segment_ids_from_eos (EOS ends its own document), plus one."""
	eos = 100257
	arr = np.array([eos, 5, 6, eos, 7, eos, eos, 8, 9], dtype=np.uint32)
	out = ShiftToInputsTargets(segment_ids_from_eos=eos).map(
	{"tokens": arr, "instance_id": 0, "file_id": 0, "instance_mask": True}
	)
	# OLMo-core: eos.cumsum() - eos = [0, 1, 1, 1, 2, 2, 3, 4, 4]; MaxText reserves 0 for padding.
	np.testing.assert_array_equal(out["inputs_segmentation"], np.array([1, 2, 2, 2, 3, 3, 4, 5], dtype=np.int32))
	# The loss mask is the default one.
	np.testing.assert_array_equal(out["targets_segmentation"], np.array([1] * 7 + [0], dtype=np.int32))
	default = ShiftToInputsTargets().map({"tokens": arr, "instance_id": 0, "file_id": 0, "instance_mask": True})
	for key in ("inputs", "targets", "inputs_position", "targets_segmentation"):
	np.testing.assert_array_equal(out[key], default[key])

[AkshitaB]

Superseded by #5867, which targets agagik-olmo35 and drops the commits that branch already has.