[AkshitaB]

This PR logs what the OLMo-core hero run olmo35-small-hero-20260907-emo logs, so the MaxText replay (olmo35-tiny on TPU v4 8x8x16) can be compared with it metric by metric. It also adds OLMo-core's in-loop perplexity evals. Each change is behind a flag whose default keeps current behavior.

It stacks on #5867: the branch is akshitab/olmo35-hero plus the three commits below, so until #5867 merges the diff shows its three commits too.

Changes

• OLMo-core training metrics (be14764a6, behind olmoe3_router_metrics and peak_device_tflops):
  • Per layer and totalled: learning/block_NN/{load_balancing_loss[_unscaled], router_z_loss[_unscaled], load_imbalance, global_load_imbalance}. Totals follow OLMo-core: max over blocks for load imbalance, sum for everything else.
  • Local load imbalance is measured over groups of router_metrics_group_seqs sequences, to match OLMo-core's per-rank value. The hero used 32 sequences per GPU early on.
  • Also: learning/ce_loss (lm loss minus z-loss), ce_perplexity, z_loss_unscaled, masked_instance_fraction, perf/mfu, steps_per_sec, total_tokens, total_petaflops, data-loading time and checkpoint save duration.
• z-loss under gradient accumulation (02e2120d9): learning/z_loss was the sum of the micro-batches' means, i.e. GA times too large. It's now the token-weighted mean over the step.
• Perplexity evals (b05d15d32, behind olmo_ppl_eval_interval): new input_pipeline/olmo_ppl_eval.py reproduces OLMo-core's LMEvaluator on the v3-small ppl validation sets.
  • One EOS-delimited document per row, padded and never packed. The CE denominator is the sum of document lengths, as in OLMo-core.
  • Logs eval/lm/<set>-validation/{CE loss, PPL}.
  • The step runs under a new olmoe3.emo_eval_pools() context, so the EMO router uses its fixed eval pool as OLMo-core does in eval mode, instead of sampling pool sizes.

Results

• Replay with every metric: 1000 steps on 1024 v4 chips.
  • The dropless replay's CE is within 0.22% of the hero (2.733 vs 2.739 at step 1000).
  • z-loss and LR match. Router metrics track the hero once the group size matches.
  • Grad norm is still 16–20% high and router z-loss 8% low late in the window.
• Perplexity evals on the step-1000 weights: 10 of the 11 sets are within ±0.5% of the hero's eval/lm/*. dolma_stack is 3.6% high, still under investigation.

Tests

On CPU, olmo_ppl_eval_test.py, olmoe3_global_lb_test.py and gradient_accumulation_nnx_test.py give 19 passed. The new tests are:

• test_router_metrics_match_whole_batch
• test_z_loss_is_the_token_weighted_mean_over_microbatches
• the five perplexity-eval tests

🤖 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 metrics to align MaxText with OLMo-core, including in-loop perplexity evaluations, detailed OLMoE3 router metrics, global load balancing, sequence-major vocabulary tiling, and a gradient accumulation loop barrier. The review feedback highlights a few important improvements: robustly extracting router metrics by checking both 'key' and 'name' attributes in PyTree paths, avoiding a potential division-by-zero error when the load-balancing loss weight is zero, and raising an explicit error instead of silently falling back to batch-major tiling when sequence length is not divisible by the number of vocabulary tiles.

[gemini-code-assist]

The current implementation of _collect_olmoe3_router_metrics only extracts keys from path if they have a key attribute (which is typical for DictKey). However, in JAX/Flax/NNX PyTree paths, some path elements can be GetAttrKey (which has a name attribute instead of key). If GetAttrKey is present, hasattr(k, "key") will be False, and the layer index won't be extracted correctly, leading to incorrect sorting of router metrics. Checking both key and name attributes ensures robust extraction across different PyTree path element types.

Suggested change

	keys = [str(k.key) for k in path if hasattr(k, "key")]
	keys = [str(k.key) if hasattr(k, "key") else str(k.name) if hasattr(k, "name") else "" for k in path]

[gemini-code-assist]

If config.load_balance_loss_weight is set to 0 (which is a valid configuration to disable load balancing loss), dividing by lb_weight will result in a division by zero (NaN or Inf) in the unscaled load balancing loss metrics. Wrapping the unscaled load balancing loss calculation in a check if lb_weight > 0: prevents this division by zero.

  if router_metrics is not None:
    lb_weight = config.load_balance_loss_weight
    lb, router_z, imbalance, global_imbalance = (router_metrics[:, i] for i in range(4))
    # Totals as OLMo-core forms them from the per-block values: max over blocks for load imbalance, sum for the
    # rest (its "global load imbalance" total is the sum of the per-block ratios).
    extra_metrics["learning/load_balancing_loss"] = jnp.sum(lb)
    if lb_weight > 0:
      extra_metrics["learning/load_balancing_loss_unscaled"] = jnp.sum(lb) / lb_weight
    extra_metrics["learning/router_z_loss"] = jnp.sum(router_z)
    extra_metrics["learning/router_z_loss_unscaled"] = jnp.sum(router_z) / _OLMOE3_ROUTER_Z_LOSS_WEIGHT
    extra_metrics["learning/load_imbalance"] = jnp.max(imbalance)
    extra_metrics["learning/global_load_imbalance"] = jnp.sum(global_imbalance)
    for i in range(router_metrics.shape[0]):
      block = f"learning/block_{i + 1:02d}"
      extra_metrics[f"{block}/load_balancing_loss"] = lb[i]
      if lb_weight > 0:
        extra_metrics[f"{block}/load_balancing_loss_unscaled"] = lb[i] / lb_weight
      extra_metrics[f"{block}/router_z_loss"] = router_z[i]
      extra_metrics[f"{block}/router_z_loss_unscaled"] = router_z[i] / _OLMOE3_ROUTER_Z_LOSS_WEIGHT
      extra_metrics[f"{block}/load_imbalance"] = imbalance[i]
      extra_metrics[f"{block}/global_load_imbalance"] = global_imbalance[i]

[gemini-code-assist]

If vocab_tiling_seq_major is set to True but the sequence length (hidden_states.shape[1]) is not divisible by num_vocab_tiling, the code silently falls back to batch-major tiling. This can lead to unexpected performance degradation due to all-gather resharding without any warning to the user. Raising a ValueError when this fallback occurs prevents silent performance traps.

  if config.vocab_tiling_seq_major and hidden_states.shape[1] % config.num_vocab_tiling != 0:
    raise ValueError(
        f"vocab_tiling_seq_major is True, but sequence length {hidden_states.shape[1]} "
        f"is not divisible by num_vocab_tiling {config.num_vocab_tiling}."
    )
  seq_major = config.vocab_tiling_seq_major