Repository navigation
OLMoE3: nanochat init, muP and CompletedP parametrizations (OLMo-core base-model mode) - #6057
Merged
gagika merged 21 commits intoOct 30, 2026
Conversation
…CE, masked fraction, throughput For a metric-for-metric comparison with the OLMo-core hero run, which logs these and MaxText did not: - olmoe3_router_metrics: per MoE layer (learning/block_NN/...) and in total (learning/...), load_balancing_loss and router_z_loss, scaled and unscaled, load_imbalance and global_load_imbalance. load_balance_loss returns [lb, load imbalance, global load imbalance] when the flag is on (the decoder layer splits it; the loss is unchanged). The counts behind the ratios are per sequence, summed over any sequence-sharding axes and gathered over the batch axes inside the routed MoE's shard_map, so the local ratio is over consecutive groups of router_metrics_group_seqs sequences (16 = the hero's per-GPU batch) whatever the TPU mesh. Totals follow OLMo-core: max over layers for load_imbalance, sums otherwise (including its global_load_imbalance total). Per-layer values are reordered by layer index (pytree key order puts layers_10 before layers_2). - learning/ce_loss (lm_loss includes the z-loss), learning/ce_perplexity, learning/z_loss_unscaled. - learning/masked_instance_fraction: rows the repetition filter masked out entirely (OLMo-core's "masked instances (%)", which is a fraction). - perf/mfu (peak_device_tflops, or inferred from the device kind), perf/steps_per_sec, perf/total_tokens, perf/total_petaflops, perf/data_loading_seconds and _percent, checkpoint/save_duration_s. Gradient accumulation averages the router metrics and masked fraction over micro-batches (OLMo-core sums the counts first, so its imbalance under GA is from the step's pooled counts). Tests: olmoe3_global_lb_test (8 simulated devices) checks the metrics vector from inside the shard_map against whole-batch counts; a 4-step CPU training run of olmoe3-30m (FSDP 8, sparse path) logs every metric, with the per-layer losses summing to the totals and to moe_lb_loss, and ce_loss + z_loss = lm_loss.
… not the sum over micro-batches loss_fn returns each micro-batch's per-token z-loss; the accumulation summed it with the other aux values but, unlike moe_lb_loss and the rest, never normalized it, so learning/z_loss (and z_loss_unscaled) came out gradient_accumulation_steps times too large and learning/ce_loss = lm_loss - z_loss too small. The step's value is now the micro-batch means weighted by their token counts, which is what GA=1 logs. The loss and gradients were never affected (the z-loss inside xent_sum is accumulated with it). Test: with per-token z-losses 1.0 and 5.0 on micro-batches of 3 and 1 tokens, the logged z_loss is 2.0 (it was 6.0).
…PL) in the training loop OLMo-core's LMEvaluator on the v3-small-ppl-validation mix: every EOS-delimited document (EOS included) is one row, truncated to the sequence length and padded, never packed. Its per-set CE averages the per-token losses over label_mask, which includes each document's last token at zero loss, so it is sum of next-token CE over tokens 1..L-1 / sum of L and PPL = exp(CE). With olmo_ppl_eval_interval > 0, the loop builds those rows once (olmo_ppl_eval_dir/<set>/val/ part-0-00000.npy, one set per batch, the last batch padded with weightless rows) and, after every step that completes a multiple of the interval (OLMo-core evaluates after step N with N % interval == 0), runs them through its own jitted eval step. A batch's token-loss sum minus its z-loss is the set's numerator; the denominator is counted on the host. Metrics go out with that step's under OLMo-core's names, eval/lm/<set>-validation/CE loss and /PPL, plus eval/lm/time_seconds. The router uses the EMO eval pool there, as OLMo-core does (eval_batch puts the model in eval mode and EmoRouterV2._pool_sizes returns eval_document_expert_pool when not training). MaxText calls the model the same way in train and eval steps (and enable_dropout is off in training too), so _emo_mask_logits sampled pool sizes in eval. emo_eval_pools(), set while the eval program is traced, makes it use emo_eval_document_expert_pool. Tests (CPU): document splitting and truncation, row targets and masks, one set per batch with weightless fill rows and the sum-of-L denominator, the eval schedule, and that the EMO router samples pools in training but keeps the eval pool under emo_eval_pools(). A 4-step training run of olmoe3-30m with the eval every 2 steps on the real WikiText-103 validation file logs eval/lm/wikitext_103-validation/CE loss = 11.45 (untrained, about ln vocab).
…nscaled-LB division _collect_olmoe3_router_metrics took layer names only from path elements with .key (DictKey); GetAttrKey carries .name, so an attribute-style path would lose its layer index and sort wrongly. The unscaled load-balancing metrics divided by load_balance_loss_weight; the layers only sow the metrics with a positive weight, but the division is now guarded on its own. (Review comments on AI-Hypercomputer#5880; the vocab-tiling fallback comment is addressed in AI-Hypercomputer#5867.)
… OLMo-core does OLMo-core's OLMoDDP train module keeps the tokens of instances masked by the instance filter in its loss divisor (seq_len - 1 each), so its divisor is batch * (seq_len - 1) every step. MaxText divides by the trained tokens only. For a step-for-step replay of an OLMo-core run, the flag switches the GA 1 divisor to OLMo-core's; logged losses stay per trained token. Besides the gradient scale (the kept fraction, 0.1-0.4%), this matters numerically: the divisor is applied before the fused LM-head kernel rounds the logit gradients to bf16, and a divisor 0.1-0.4% below 2^24 makes that rounding systematically biased (1/N rounds down to 2^-24), while 2048 * 8191 is 0.01% below.
…t accumulation loss_fn returns the count its gradient is normalized by (grad_weights: OLMo-core's, filtered instances counted, or the trained tokens), uses it for the micro-batch aux-loss weights, and gradient_accumulation divides the accumulated gradient by its sum over micro-batches. The manual-GA path already differentiates the sum scaled by 1/(batch * seq_len), a power of two, so its bf16 rounding was unbiased; this only changes the final fp32 division to match OLMo-core's gradient scale. Logged losses stay per trained token.
metric_logger.record_train_metrics: keep perf/steps_per_sec and the OLMOE3_FAST_HOST_LOOP learning-rate path. olmoe3 load_balance_loss: the counts now come from bincount (4f4f717), so the router metrics build the one-hot expert mask themselves, only when olmoe3_router_metrics is on.
…from a set of sizes instead of a range Off by default (empty list keeps the uniform draw from [emo_min_document_expert_pool, emo_max_document_expert_pool]). With e.g. [8, 64, 256], each document's pool size is one of those, drawn uniformly; eval still uses emo_eval_document_expert_pool. Each size must satisfy top_k <= size <= num_experts.
Optional, one entry per emo_document_expert_pool_choices, summing to 1: each training document's pool size is then drawn from the choices with these probabilities (jax.random.choice) instead of uniformly. Empty keeps the uniform draw over the choices, so existing configs are unchanged. E.g. choices [32, 64, 256] with probs [0.25, 0.25, 0.5]. Validated at construction: length must match the choices, entries non-negative, sum 1 within 1e-6. Originally written by Ryan Wang.
Off by default ("normal" keeps the std-0.02 truncated normal, bit-identical
for the same seed). "nanochat" is OLMo-core's InitMethod.nanochat
(tanushy/parameterization-ladders): every matrix a truncated normal (+/-3 std)
with std 1/sqrt(fan_in) * min(1, sqrt(fan_out/fan_in)), the KDA short convs
1/sqrt(kernel_size), and the embedding an untruncated normal with std 1.0.
Norm gains, A_log, dt_bias and biases are unchanged. Init only.
Fan-in/out come from the kernel axes, except: the fused attention query kernel
holds the reference's w_q and w_g, so its fan-out is heads * head_dim; and the
[experts, in, out] expert kernels are per-expert matrices. fused_mlp and
prefuse_moe_weights (gate and up in one kernel) are rejected.
Tests: the reference's nanochat_std cases; a per-parameter check against the
rules written out by name; a golden file of the std OLMo-core's own init gives
each parameter of the same reduced hero, with names from OLMo-core's MaxText
checkpoint converter (same parameters and shapes, kind, std, truncation); and
the default init at std 0.02.
…2x shared expert olmo35-small (d_model 1536, 40 layers, 7:1 KDA:attention, head_dim 128, 512 experts top-16 of width 1536) under the change that makes olmo35-tiny-e256k8-sh2x from olmo35-tiny: 256 experts top-8 of width 3072 and a 3072-wide shared expert. 72,498,550,144 total / 4,041,217,408 active parameters; the count test checks the total.
… OLMo-core's base-model mode Off by default (empty mup_base_model_name). With a base model (e.g. olmo35-tiny-e256k8-sh2x as the base of olmo35-small-e256k8-sh2x), each weight matrix trains at lr * base_fan_in(role) / fan_in, its role being the layer kind (dense or MoE, KDA or attention) plus its path in the layer, so every matrix of the base trains at lr; the embedding at lr * mup_embedding_lr_multiplier, every other parameter (norms, biases, A_log, dt_bias, SSMax gains, depthwise convs) at lr. As OLMo-core's MuPConfig.from_base_model (tanushy/parameterization-ladders). mup_scale_weight_decay (default true) sets each parameter's decay to wd / m, so it decays by lr * wd at every width, the recommendation in scaling-ladders' MUP_WEIGHT_DECAY.md; false keeps wd (lr * m * wd, OLMo-core's default). adamw_mask still exempts its parameters (the embedding in our runs). mup_adamw keeps optax.adamw's three transforms and their state, so checkpoints load with muP on or off. olmoe3 with opt_type=adamw and scan_layers=False only. Tests (tests/unit/olmoe3_mup_test.py): the multipliers of olmo35-small-e256k8-sh2x and of the base itself against the ones OLMo-core computes for the same pair (golden file, names from OLMo-core's MaxText converter); rejection of unknown parameters, shape/config mismatches and layer kinds the base lacks; all multipliers 1 is optax.adamw bitwise, state included; a zero-gradient step decays by lr * wd (scaled) or lr * m * wd (coupled), masked parameters not at all; the first Adam step; coupled is m times AdamW through a warmup; and get_optimizer with skip-step.
…dPConfig in base-model mode Off by default (empty completedp_base_model_name). Complete(d)P (arXiv:2512.22382 Table 1) relative to a base model (e.g. olmo35-tiny-e256k8-sh2x as the base of olmo35-small-e256k8-sh2x) and run, as OLMo-core's CompletedPConfig.from_base_model and CompletedPOptimConfig (tanushy/parameterization-ladders): - both residual branches of every layer scaled by (layers / base layers)^-alpha (completedp_alpha, default 1.0); - the normal init of each matrix scaled by sqrt(base fan-in / fan-in), or by base fan-in / fan-in for the LM head and the readout matrices (router, KDA f_proj_1 / g_proj_1); - every parameter's AdamW lr, eps and weight decay scaled by its role (input, hidden matrix / readout / vector / head-shared vector, output), its width ratio and the depth ratio, plus OLMo-core's batch-size and duration rules (completedp_base_batch_tokens / _train_tokens; 0 = this run's) on lr, wd, eps and the betas. olmoe3 with opt_type=adamw, olmo_init_method=normal and scan_layers=False only; not with muP. adamw_mask still exempts its parameters. completedp_adamw keeps optax.adamw's three transforms and state structure. The role and fan-in of every olmoe3 parameter now live in models/olmoe3_roles.py, shared by muP and CompletedP (muP's multipliers unchanged). Tests (tests/unit/olmoe3_completedp_test.py): every parameter's role and lr / eps / wd multipliers, the betas and the residual multiplier of olmo35-small-e256k8-sh2x (alpha 1 and 0.5, without and with 2x batch / 4x tokens) and of the base against OLMo-core's (golden file, names from OLMo-core's MaxText converter); the init factors against OLMo-core's rule and the init stds of a reduced model; a layer's output is x + c * attn + c * ffn; all multipliers 1 is optax.adamw bitwise, state included; a zero-gradient step decays by lr m_lr wd m_wd; the first step uses each parameter's eps and the scaled betas; get_optimizer with skip-step; rejection of unsupported runs. The default init is bit-identical to before.
…-all drops From the [source shard, expert] counts the dispatch already uses: each destination shard's load against its receive buffer, so the rows dropped, the total, and the largest load over the balanced one. One device per EP group (expert index 0 of data index 0) prints an A2A_DROPS line per call. Off by default.
…TOR for the 8x8x48 preset Both default to the measured recipe (olmoe3-3p5b-e256k8-sh2x, 1.0625), so the preset is unchanged unless they're set. PRESET_MODEL_NAME=olmo35-small-e256k8-sh2x runs the 7:1 small at the same recipe, with GMM tiles for its 3072-wide experts (768 / 1536); PRESET_RAGGED_BUFFER_FACTOR sweeps the all-to-all receive buffer, with moe_log_a2a_drops=true to log the rows it drops.
… 256 / 512) instead of 128 Same attention and KDA widths; head_dim 128 runs poorly on v7x. The head layout of olmoe3-3p5b-e256k8-sh2x. 72,580,704,064 total / 4,123,371,328 active parameters.
… 3.5 embedding and attention features The 3.5B Option A shape (d_model 1792, 30 layers, 4:1 KDA:attention, E=256, K=8, F=3584, 2x shared expert) plus what olmo35-tiny / small have and it lacks: the embedding scaled by sqrt(d_model) and RMSNormed, per-head QK norm gains and scalable softmax. With these every parameter has a counterpart role in olmo35-tiny-e256k8-sh2x, so CompletedP can use tiny as its base. 63,130,193,648 total / 3,741,994,736 active parameters; the count test checks the total.
Regenerated from OLMo-core (f65712270) with small-sh2x as MaxText's small since f9b9dcf: 8/4 heads of head_dim 256, KDA 8 heads of key / value head dim 256 / 512 (72,580,704,064 params). The 16/8 x 128 goldens failed the target and init tests since then. The muP tally follows: 0.5 on the attention output and KDA f_proj_2 / g_proj_2.
…rom the base model's Off by default (OLMo-core's rules). With a muP or CompletedP base model: - full attention's logit scale is head_dim^-0.5 * sqrt(base head_dim / head_dim), after the QK norm (muP's 1/head_dim beyond the base; 1/sqrt(2) on top of 1/sqrt(256) for small on tiny); - CompletedP: a matrix is a readout iff its output width equals the base model's (small on tiny: KDA f_proj_1 / g_proj_1 become hidden, w_b a readout), and a head-shared gain's eps is multiplied by (heads / base heads) / m_N.
github-actions
Bot
added
the
skip-ci-tests
AkshitaB
requested review from
aireenmei,
darisoy,
gagika,
igorts-git,
richjames0,
shralex,
shuningjin and
zxhe-sean
as code owners
October 30, 2026 18:43
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request integrates advanced model scaling and initialization techniques from OLMo-core into the OLMoE3 architecture. By introducing muP and CompletedP parametrizations, it provides robust methods for scaling models while maintaining training stability. Additionally, the PR enhances observability through new performance metrics, router logging, and in-loop perplexity evaluations, and adds new model configurations to facilitate experimentation with these features. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
AkshitaB
requested review from
A9isha,
NuojCheng,
abhinavclemson,
bvandermoon,
dipannita08,
gobbleturk,
hengtaoguo,
huytransformer,
jiangjy1982,
khatwanimohit,
parambole,
vipannalla and
xibinliu
as code owners
October 30, 2026 18:43
… gated metrics, dropless ppl eval, static_loss_cotangent - Router metrics: each row is placed by its global layer index from its path (layers_<i>; scanned: layers_0/layer_<j> and a row per cycle of scanned_blocks/layer_<j>), and blocks are named from it (i + first_num_dense_layers). Scanned rows came out sublayer-major before. - ce_loss, ce_perplexity, z_loss_unscaled and masked_instance_fraction are behind olmoe3_router_metrics, so other models' train steps are unchanged. - OLMo-core's perplexity evals run on their own graphdef: every RoutedMoE (subclasses included) at eval_ragged_buffer_factor when > 0, else force_dropless, so padded documents don't overflow a training-sized buffer. Documents fill only the micro_batch_size_to_eval_on rows the eval step scores, so the denominator matches. - static_loss_cotangent (off by default): at GA=1, the gradient of the token-loss sum times a static 1 / (batch * seq len), divided by the loss divisor in f32 afterwards, as manual GA does, so 1 / divisor doesn't round in the bf16 backward. - Tests: olmoe3_router_metrics_test (synthetic and real scanned / unscanned paths), ppl graphdef and scored-rows tests, static_loss_cotangent tests; mock configs updated (EMo pool choices; getattr for enable_checkpointing). pyink 24.10.1 and pylint 3.3.8 clean on the touched files.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #5880 (
akshitab/olmo35-hero-metrics). Until that merges, this PR also shows its commits; the diff of this PR alone is AkshitaB/maxtext@akshitab/olmo35-hero-metrics...akshitab/olmo35-parametrizationSummary
Adds OLMo-core's width/depth parametrizations to OLMoE3 as options, with the default parametrization unchanged:
olmo_init_method=nanochat): OLMo-core's spectral-condition init.mup_base_model_name): OLMo-core'sMuPConfigin base-model mode. Each weight matrix trains atlr * base fan-in / fan-in. Weight decay is scaled by the multiplier (mup_scale_weight_decay, default true), and the embedding has its own multiplier.completedp_base_model_name): OLMo-core'sCompletedPConfigin base-model mode (arXiv:2512.22382). It adds:m_L^-alpha;completedp_base_batch_tokens/completedp_base_train_tokens). These are implemented but off by default.head_dim_transfer(off by default): rules for head dims that differ from the base model's.head_dim^-0.5 * sqrt(base head_dim / head_dim), applied after the QK norm (muP's 1/head_dim beyond the base).(H / H_base) / m_N.olmo35-small-e256k8-sh2x: d_model 1536, 40 layers, 7:1, head_dim 256, KDA 256 / 512; 72.58B total / 4.12B active.olmo35-gagik-fixed-e256k8-sh2x:olmoe3-3p5b-e256k8-sh2xplus the OLMo 3.5 embedding and attention features, so every parameter has a CompletedP role in tiny.Parameter roles and fan-ins live in
models/olmoe3_roles.py, which muP and CompletedP share. Both needscan_layers=Falseand AdamW, and raise otherwise.The branch also carries two small run-side commits:
moe_log_a2a_drops: an off-by-default debug probe for rows the expert-major all-to-all drops.PRESET_MODEL_NAME/PRESET_RAGGED_BUFFER_FACTORknobs for the v4 8x8x48 launcher preset.Tests
tests/unit/olmoe3_test.py: nanochat init, and parameter counts for the new configs.tests/unit/olmoe3_mup_test.py: multipliers checked against OLMo-core's goldens (tests/assets/olmoe3_mup_multipliers_olmo_core.json).tests/unit/olmoe3_completedp_test.pychecks, against OLMo-core's goldens (tests/assets/olmoe3_completedp_olmo_core.json):head_dim_transfer.tanushy/parameterization-ladders(f65712270), with names converted byolmo_core.nn.maxtext.All 48 parametrization and parameter-count tests pass on CPU. In the full
olmoe3_test.pysuite, only the 4 fused-KDA tests fail; they also fail on #5880's head.🤖 Generated with Claude Code