Directory Diff Report

Left: experiments/grug/moe

Right: experiments/grug/moe_pipeline

Filter: .py, .md

Generated: 2026-09-02 17:47:02Z

Inline Diffs: 13 files

Changed

Added

Removed

Unchanged

File	Status	Lines Added	Lines Deleted
README.md	changed	+50	-189
model.py	changed	+29	-22
train.py	changed	+375	-566
benchmark.py	added	+99	-0
pipeline.py	added	+625	-0
adamh.py	removed	+0	-78
agent.md	removed	+0	-207
heuristic.py	removed	+0	-264
launch.py	removed	+0	-255
launch_cw_scale.py	removed	+0	-238
launch_datakit_moe_mix.py	removed	+0	-466
optimizer.py	removed	+0	-314
test_optimizer.py	removed	+0	-37
__init__.py	unchanged	+0	-0

README.md

changed

	experiments/grug/moe/README.md			experiments/grug/moe_pipeline/README.md
n	1	# grug-moe	n	1	# grug-moe-pipeline
	2			2
n	3	Current best recipe for training Mixture-of-Experts models in the grug	n	3	Canonical pipeline-parallel Grug MoE implementation. This directory is a copy-paste
	4	template. Model, optimizer, train loop, and launch wiring all live in		4	variant of [`experiments/grug/moe`](../moe/): its model and execution path are local
	5	`experiments/grug/moe/` so the MoE variant can iterate independently from		5	so pipeline-specific changes do not turn the ordinary MoE implementation into a
	6	the dense base template.		6	shared trainer framework.
	7			7
n	8	## Architecture	n	8	[`pipeline.py`](./pipeline.py) splits the local transformer into stage pytrees and
				9	builds automatic JaxPP ZeroBubble or DualPipeV optimizer steps. [`train.py`](./train.py)
				10	owns model initialization, the training loop, and the Fray dispatch entry point.
				11	[`benchmark.py`](./benchmark.py) maps environment variables onto that trainer for
				12	repeatable performance runs. One logical stage maps to each physical pipeline rank with
				13	`PIPELINE_SCHEDULE=automatic_zero_bubble`. Set `PIPELINE_SCHEDULE=automatic_dualpipe_v`,
				14	`PIPELINE_STAGES=2P`, and `PIPELINE_PHYSICAL_STAGES=P` to fold two logical stages onto
				15	each physical rank in JaxPP's V-shaped placement. DualPipeV requires at least `2P`
				16	microbatches. `PIPELINE_LAYERS_PER_STAGE` contains one positive layer count per logical
				17	stage.
	9			18
n	10	Experts.	n	19	The best validated Snowball 67B-A2B throughput point uses eight H100x8 replicas,
	11	- `num_experts = 256` routed pool; `num_experts_per_token = 4` active per token.		20	sixteen logical stages, batch 256, 32 microbatches, sequence length 8192, and layer
	12	- One always-on shared dense MLP per block, in parallel with the routed		21	counts `1,2,2,2,2,2,2,2,1,1,1,1,2,2,2,1`:
	13	experts (contributes to every token).
	14			22
n	15	Router.	n	23	```bash
	16	- Linear projection of `hidden_dim → num_experts` in fp32 (cast back at		24	export XLA_FLAGS='--xla_gpu_executable_terminate_timeout=300 --xla_gpu_enable_command_buffer='
	17	the end). Top-k, softmax, and QB statistics all run in fp32.		25	export PIPELINE_SCHEDULE=automatic_dualpipe_v
	18	- A `stop_gradient` bias term gets added to the router logits before top-k;		26	export PIPELINE_PHYSICAL_STAGES=8
	19	the bias is updated each step from the previous step's QB-β statistics		27	export PIPELINE_STAGES=16
	20	(see `train.py::_apply_qb_betas`).		28	export PIPELINE_BATCH=256
	21	- QB load balancing: per-expert β is the top-k logit threshold averaged		29	export PIPELINE_MICROBATCHES=32
	22	across the batch. On the next step, `router_bias := -β`, pushing		30	export PIPELINE_SEQ_LEN=8192
	23	rarely-selected experts up and over-selected experts down. Replaces an aux		31	export PIPELINE_HIDDEN_DIM=2560
	24	load-balancing loss — the bias mechanism is invisible to gradients.		32	export PIPELINE_INTERMEDIATE_DIM=1280
	25	- Combine weights: sigmoid on the unbiased router logits of the K		33	export PIPELINE_SHARED_EXPERT_INTERMEDIATE_DIM=2560
	26	selected experts, then renormalised to sum to 2.5		34	export PIPELINE_LAYERS=26
	27	(`_ROUTING_RENORM_SUM`).		35	export PIPELINE_LAYERS_PER_STAGE=1,2,2,2,2,2,2,2,1,1,1,1,2,2,2,1
	28			36	export PIPELINE_EXPERTS=256
	29	Attention (`CausalSelfAttention.__call__`).		37	export PIPELINE_TOP_K=4
	30	- GQA: default ratio 4:1 (`num_kv_heads = num_heads / 4`).		38	export PIPELINE_EXPERT_AXIS=8
	31	- Half-RoPE: rotary embeddings applied only to the first half of Q/K per		39	export PIPELINE_HEADS=20
	32	head (`q[..., :head_dim/2]`, `k[..., :head_dim/2]`); the second half is		40	export PIPELINE_KV_HEADS=5
	33	rope-free on every layer.		41	export PIPELINE_VOCAB_SIZE=128256
	34	- PKO (Partial Key Offset) is wired up but disabled by default		42	export PIPELINE_SLIDING_WINDOW=2048
	35	(`disable_pko=True`). When enabled it would shift the rope-free second half		43	export PIPELINE_QK_MULT=1.5703
	36	of K back by one position on every-4th + last "long" layers (zero at doc		44	export PIPELINE_MP=params=bfloat16,compute=bfloat16,output=bfloat16
	37	starts, then rms-norm); short layers always skip PKO.		45	export PIPELINE_ATTENTION=gpu_fa4_cute
	38	- NoPE on long layers: every-4th + last layers run with rotary embedding		46	export PIPELINE_MOE=ring
	39	skipped entirely (`disable_long_rope=True`); short layers keep half-RoPE.		47	export PIPELINE_REMAT=recompute_all
	40	- Sliding window: long layers run full causal attention		48	uv run --extra pipeline python -m experiments.grug.moe_pipeline.benchmark
	41	(`sliding_window = None`). Short layers run `cfg.sliding_window` (default
	42	2048).
	43
	44	To reduce risk of long-context extension, we have excluded PKO and use NoPE on
	45	long layers.
	46	- XSA (Exclusive Self-Attention): after attention, subtract the component
	47	of each head's output parallel to its `aligned_v`: `z = y − (yᵀv / ‖v‖²)·v`
	48	per head. Followed by a headwise sigmoid gate.
	49
	50	Norms.
	51	- GatedNorm: rank-128 low-rank gate on RMS-normalised input pre-attention
	52	and pre-MLP. Acts as a learned per-token gate over the hidden dimension.
	53
	54	Optimizer + schedule ([`heuristic.py`](./heuristic.py) +
	55	[`optimizer.py`](./optimizer.py)).
	56	- MuonH (`GrugMoeMuonHConfig`, registered as `grug_moe_muonh_v1`):
	57	Newton-Schulz orthogonalisation + Frobenius-hyperball scale-invariant
	58	updates on the matrix + GatedNorm group.
	59	- `adamh` (lm_head only).
	60	- `adam` (token_embed, router, attn_gate, biases, 1-D norm weights).
	61	- No gradient clipping (`max_grad_norm = None`).
	62	- 1% warmup, linear decay to 0, `min_lr_ratio = 0`.
	63	- LR scaling (fit on the May Recipe sweep, issue #5951, R²=0.996):
	64	`muonh_lr = 18.31 · tokens^-0.395 · dim^-0.150 · sqrt(B)`
	65	(equivalently `adam_lr = 0.06602 · tokens^-0.395 · dim^-0.150 · sqrt(tpb)`).
	66
	67	Loss.
	68	- Cross-entropy on the next-token logits.
	69	- Final-logit z-loss (`GrugTrainerConfig.z_loss_weight = 1e-4` by
	70	default): adds `z_loss_weight · mean(logsumexp(logits)²)` to stabilise the
	71	lm-head softmax.
	72	- Router z-loss off by default (`router_z_loss_coef = 0.0`).
	73	- No auxiliary load-balancing loss; QB router bias does the balancing.
	74
	75	Other.
	76	- Expert parallelism: `ragged_all_to_all` or ring-based via
	77	`levanter.grug.grug_moe.moe_mlp` (default: ring). Default capacity factor 1.0.
	78
	79	## Scaling heuristic
	80
	81	[`MoeHeuristic`](./heuristic.py) turns `(budget, hidden_dim)` into model +
	82	optimizer hyperparameters via `build_model_config(hidden_dim, seq_len)` and
	83	`build_optimizer_config(batch_size, tokens, hidden_dim, seq_len)`. May Recipe
	84	refit on the MuonH LR sweep (issue #5951; 17 cells, R²=0.996). All formulas
	85	anchor at seq_len = 4096 and write batch effects in terms of
	86	`tokens_per_batch = batch_size · seq_len`.
	87
	88	- LR: `adam_lr = 0.06602 · tokens^-0.395 · dim^-0.150 · sqrt(B)`,
	89	`muonh_lr = (13/3) · adam_lr`.
	90	- Compute budget: `C = 3 · flops_per_token(no_lm_head) · tokens`.
	91	- Epsilon: `epsilon_coeff · sqrt(tokens / tokens_per_batch)`.
	92	- Beta1: fixed at 0.9062.
	93	- Beta2: `clip(0.999^(tpb/131072), 0.95, 0.9999)` (constant-token half-life).
	94	- Layer count: `num_layers ≈ dim / (64 + 4·log2(dim) − 9)` (rounded).
	95	- GQA: largest divisor of `num_heads ≤ num_heads / 4`.
	96
	97	For the earlier AdamH-tuned heuristic (64 experts, used on the May 2026 1e23
	98	hero run), see the pre-rename
	99	[`heuristic.py` on main](https://github.com/marin-community/marin/blob/8586719b524bf7743ec5034403c7e834505fe73e/experime
	>	nts/grug/moe/heuristic.py).
	100
	101	## Compute-optimal baseline
	102
	103	For each hidden dim we picked a compute budget at the parabola optimum of
	104	its isoflop curve (V2 / May Recipe: drop-1e18 fit, issue #6074; V1 / v16:
	105	issue #4447) and ran the cell at that budget. These are the baseline runs
	106	that ablation experiments compare against.
	107
	108	### May Recipe (drop-1e18 fit, issue #6074) — current baseline
	109
	110	Reference runs on v4-32 us-central2 with EP=1, MuonH optimizer
	111	(`muonh_lr` from `heuristic.MoeHeuristic.build_optimizer_config`), 1pct-noclip
	112	schedule, no permanent step-interval checkpoints.
	113
	114	The Paloma macro numbers below were measured under the previous defaults:
	115	`seq_len = 4096`, PKO on for long layers, partial half-RoPE on every layer
	116	(no NoPE on long layers), and final-logit z-loss off (`z_loss_weight = 0`).
	117
	118	\| Budget \| Dim \| Layers \| bs \| Steps \| Tokens \| Paloma macro \| v4-32 tok/s \| Runtime \| Run \|
	119	\|----------\|--------\|--------\|-----\|----------\|---------\|--------------\|-------------\|---------\|-----\|
	120	\| 3.82e17 \| d512 \| 6 \| 32 \| 10,980 \| 1.44e9 \| 3.5422 \| 433,986 \| 1.25h \| [moe_may_compute_opt_d5
	>	12_ep1](https://wandb.ai/marin-community/marin_moe/runs/moe_may_compute_opt_d512_ep1) \|
	121	\| 2.81e18 \| d768 \| 8 \| 64 \| 16,875 \| 4.42e9 \| 3.2273 \| 294,726 \| 4.83h \| [moe_may_compute_opt_d7
	>	68_ep1](https://wandb.ai/marin-community/marin_moe/runs/moe_may_compute_opt_d768_ep1) \|
	122	\| 1.16e19 \| d1024 \| 11 \| 128 \| 16,080 \| 8.43e9 \| 3.0195 \| 219,720 \| 11.72h \| [moe_may_compute_opt_d1
	>	024_ep1](https://wandb.ai/marin-community/marin_moe/runs/moe_may_compute_opt_d1024_ep1) \|
	123	\| 3.46e19 \| d1280 \| 13 \| 256 \| 14,325 \| 1.50e10 \| 2.8857 \| 171,912 \| 25.68h \| [moe_may_compute_opt_d1
	>	280_ep1](https://wandb.ai/marin-community/marin_moe/runs/moe_may_compute_opt_d1280_ep1) \|
	124
	125	Fitted scaling law on these 4 cells (`L_inf=1.6` pinned, α pinned to the v16
	126	exponent of 0.0941):
	127
	128	```
	129	loss(C) = 1.6 + 88.32 · C^-0.0941
	130	```		49	```
	131			50
t	132	~2.22× equal-TPS compute-equivalent speedup vs v16 baseline at every budget.	t	51	A 14-step run with four warmups measured 3.476 s median step time, 603,327 tokens/s,
	133			52	19.46% analytic MFU, and 59.153 GB peak memory per device. A matched batch-128,
	134	### v16 baseline (historical reference)		53	16-microbatch control measured 543,490 tokens/s and 17.53% analytic MFU. Both used
	135			54	eight examples per microbatch. The larger batch improved median throughput by 11.0%
	136	Older AdamH MoE sweep (`group=isoflop-moe-v16` in `marin-community/dial_moe`,		55	without increasing peak memory; final loss differed by 4.0e-5 after 14 synthetic-data
	137	issue #4447). Kept here because the v16 scaling law		56	steps.
	138	(`loss = 1.6 + 95.18 · C^-0.0941`) is the reference curve in
	139	`experiments/grug/moe/agent.md` for gate-1 / gate-2 effective speedup
	140	calculations. Empirically validated at 1e21: predicted macro 2.606, measured
	141	d2560-v2 run came in at 2.599.
	142
	143	\| Budget \| Dim \| Layers \| Paloma macro \| Tokens \| v5p-8 avg tok/s \| v5p-8 runtime \| Run \|
	144	\|----------\|----------\|--------\|-------------\|---------\|-----------------\|---------------\|-----\|
	145	\| 2.19e17 \| d512 \| 6 \| 3.8104 \| 8.37e8 \| 405,630 \| 0.6h \| [moe-v16-compute-opt-d512-2.1
	>	9e+17](https://wandb.ai/marin-community/dial_moe/runs/moe-v16-compute-opt-d512-2.19e+17) \|
	146	\| 1.70e18 \| d768 \| 8 \| 3.4339 \| 2.71e9 \| 273,532 \| 2.8h \| [moe-v16-compute-opt-d768-1.7
	>	0e+18](https://wandb.ai/marin-community/dial_moe/runs/moe-v16-compute-opt-d768-1.70e+18) \|
	147	\| 9.00e18 \| d1024 \| 11 \| 3.1605 \| 6.63e9 \| 175,165 \| 10.5h \| [moe-v16-compute-opt-d1024-9.
	>	00e+18](https://wandb.ai/marin-community/dial_moe/runs/moe-v16-compute-opt-d1024-9.00e+18) \|
	148	\| 2.83e19 \| d1280 \| 13 \| 3.0065 \| 1.24e10 \| 128,277 \| 26.8h \| [moe-v16-compute-opt-d1280-2.
	>	83e+19](https://wandb.ai/marin-community/dial_moe/runs/moe-v16-compute-opt-d1280-2.83e+19) \|
	149
	150	## Promotion criteria
	151
	152	Changes can be promoted to this recipe when they demonstrate some combination
	153	of the following. Typically point 1 is sufficient.
	154
	155	1. Passes gate 1 and gate 2 as defined in [`agent.md`](./agent.md) —
	156	effective speedup > 1 at all compute-optimal baseline points, and lower
	157	projected macro_loss at 1e21 and 1e23.
	158	2. Low curvature around the minimum of each isoflop curve — stable
	159	behavior across under- and over-trained regimes, in particular the
	160	overtrained regime.
	161	3. Stability and scaling improvements — better routing balance, controlled
	162	norm growth, fewer activation outliers. Anything that makes the recipe more
	163	robust to scaling, even if loss is neutral at small scale.
	164
	165	Most promotable changes will land in one of three files:
	166
	167	- [`model.py`](./model.py) — architecture tweaks (routing, norms, attention,
	168	activation functions, expert layout, etc.).
	169	- [`heuristic.py`](./heuristic.py) — scaling heuristics (LR formula
	170	coefficients, depth/width formula, GQA ratio, per-batch-size epsilon/beta2
	171	scaling).
	172	- [`optimizer.py`](./optimizer.py) — optimizer internals (MuonH / AdamH
	173	components, parameter-group partitioning, per-group learning rates,
	174	weight decay).
	175
	176	Some discretionary factors may influence the promotion decision even when the
	177	loss criteria are met — for example, impact on training memory footprint,
	178	inference latency / KV-cache size, serving compatibility, or interaction effects with other promotable changes.
	179
	180	## Files
	181
	182	- [`model.py`](./model.py) — `GrugModelConfig` + transformer implementation.
	183	- [`optimizer.py`](./optimizer.py) — `GrugMoeAdamHConfig` and
	184	`GrugMoeMuonHConfig` wrappers with expert-param-group awareness.
	185	- [`train.py`](./train.py) — `GrugTrainState`, `train_step`, `_apply_qb_betas`,
	186	`run_grug` (dispatches a Fray job).
	187	- [`heuristic.py`](./heuristic.py) — `MoeHeuristic` (MuonH, May Recipe
	188	refit) and `build_from_heuristic` entry point. Current default.
	189	- [`launch.py`](./launch.py) — `GrugMoeLaunchConfig`, the `grug_moe_baseline()`
	190	lazy `Checkpoint`, and `StepRunner` wiring.
	191	- [`launch_cw_scale.py`](./launch_cw_scale.py) — CoreWeave scale-test launcher.
	192	- [`adamh.py`](./adamh.py) — shared AdamH utilities.
	193	- [`test_optimizer.py`](./test_optimizer.py) — unit tests for the AdamH
	194	parameter-group mask.
	195	- [`agent.md`](./agent.md) — agent guide for running ablation experiments on Iris.

model.py

changed

	experiments/grug/moe/model.py			experiments/grug/moe_pipeline/model.py
	9			9
	10	import dataclasses		10	import dataclasses
	11	from dataclasses import dataclass		11	from dataclasses import dataclass
n			n	12	from enum import StrEnum
	12	from typing import Any, Literal		13	from typing import Any, Literal
	13			14
	14	import equinox as eqx		15	import equinox as eqx
	19	from haliax import Axis		20	from haliax import Axis
	20	from haliax.jax_utils import named_call		21	from haliax.jax_utils import named_call
	21	from jax import core, random		22	from jax import core, random
n			n	23
				24	try:
				25	from jax.shard_map import shard_map
				26	except ModuleNotFoundError as error:
				27	if error.name != "jax.shard_map":
				28	raise
				29	from jax.experimental.shard_map import shard_map
	22	from jax.sharding import NamedSharding, get_abstract_mesh, reshard		30	from jax.sharding import NamedSharding, get_abstract_mesh, reshard
	23	from jax.sharding import PartitionSpec as P		31	from jax.sharding import PartitionSpec as P
n	24		n
	25	try:
	26	from jax.shard_map import shard_map
	27	except ModuleNotFoundError:
	28	from jax.experimental.shard_map import shard_map
	29	from jaxtyping import Array, Float, Int, PRNGKeyArray		32	from jaxtyping import Array, Float, Int, PRNGKeyArray
	30	from levanter.compat.hf_checkpoints import HFCheckpointConverter		33	from levanter.compat.hf_checkpoints import HFCheckpointConverter
	31	from levanter.grug.attention import (		34	from levanter.grug.attention import (
	105			108
	106	@dataclass(frozen=True)		109	@dataclass(frozen=True)
	107	class GrugModelConfig:		110	class GrugModelConfig:
n	108	"""Hyperparameters for the grug MoE transformer.	n	111	"""Architecture and execution settings for the Grug MoE transformer."""
	109
	110	Architecture choices (GatedNorm, XSA, QB routing) are hardcoded.
	111	Only shape/size knobs live here. All layers are MoE.
	112	"""
	113			112
	114	vocab_size: int		113	vocab_size: int
	115	hidden_dim: int = 512		114	hidden_dim: int = 512
	403	@named_call		402	@named_call
	404	def __call__(self, x: Float[Array, "... D"]) -> Float[Array, "... D"]:		403	def __call__(self, x: Float[Array, "... D"]) -> Float[Array, "... D"]:
	405	gate_hidden = jnp.einsum("...d,dr->...r", x, self.w_down)		404	gate_hidden = jnp.einsum("...d,dr->...r", x, self.w_down)
n	406	# TODO: silu activation here isn't explored, just cargo-culted from Qwen. Likely low-hanging ablation fruit	n
	407	# (e.g. compare no activation, relu, etc.).
	408	gate_hidden = jax.nn.silu(gate_hidden)		405	gate_hidden = jax.nn.silu(gate_hidden)
	409	gate = jax.nn.sigmoid(jnp.einsum("...r,rd->...d", gate_hidden, self.w_up))		406	gate = jax.nn.sigmoid(jnp.einsum("...r,rd->...d", gate_hidden, self.w_up))
	410	return x * gate.astype(x.dtype)		407	return x * gate.astype(x.dtype)
	609	local_tokens = s_minus_alpha.shape[0] // num_devices		606	local_tokens = s_minus_alpha.shape[0] // num_devices
	610	qb_count = max(1, local_tokens * self.cfg.num_experts_per_token // self.cfg.num_experts)		607	qb_count = max(1, local_tokens * self.cfg.num_experts_per_token // self.cfg.num_experts)
	611			608
n	612	def _local_qb_beta(s_ma):	n	609	def _local_qb_beta(local_s_minus_alpha):
	613	topk_vals, _ = jax.lax.top_k(s_ma.T, qb_count)		610	topk_vals, _ = jax.lax.top_k(local_s_minus_alpha.T, qb_count)
	614	beta = topk_vals[:, -1]		611	beta = topk_vals[:, -1]
	615	return jax.lax.pmean(beta, axis_name=BATCH_AXES)		612	return jax.lax.pmean(beta, axis_name=BATCH_AXES)
	616			613
	635	routed = rearrange(routed_flat, "(b s) d -> b s d", b=b, s=s)		632	routed = rearrange(routed_flat, "(b s) d -> b s d", b=b, s=s)
	636	routed = reshard(routed, _batch_spec())		633	routed = reshard(routed, _batch_spec())
	637	return routed, router_stats		634	return routed, router_stats
n			n	635
				636
				637	class LayerAttentionMode(StrEnum):
				638	SHORT = "short"
				639	LONG = "long"
	638			640
	639			641
	640	class Block(eqx.Module):		642	class Block(eqx.Module):
	669	self,		671	self,
	670	x: Float[Array, "B S D"],		672	x: Float[Array, "B S D"],
	671	mask: AttentionMask \| jax.Array,		673	mask: AttentionMask \| jax.Array,
n	672	use_pko: bool = False,	n	674	attention_mode: LayerAttentionMode,
	673	disable_rope: bool = False,
	674	) -> tuple[Float[Array, "B S D"], dict[str, jax.Array]]:		675	) -> tuple[Float[Array, "B S D"], dict[str, jax.Array]]:
n			n	676	is_long = attention_mode == LayerAttentionMode.LONG
	675	with jax.named_scope("attention"):		677	with jax.named_scope("attention"):
	676	attn_in = self.attn_gated_norm(self.rms_attn(x))		678	attn_in = self.attn_gated_norm(self.rms_attn(x))
n	677	x = x + self.attn(attn_in, mask, use_pko=use_pko, disable_rope=disable_rope)	n	679	x = x + self.attn(
				680	attn_in,
				681	mask,
				682	use_pko=is_long and not self.attn.cfg.disable_pko,
				683	disable_rope=is_long and self.attn.cfg.disable_long_rope,
				684	)
	678	with jax.named_scope("routed_moe"):		685	with jax.named_scope("routed_moe"):
	679	mlp_in = self.mlp_gated_norm(self.rms_mlp(x))		686	mlp_in = self.mlp_gated_norm(self.rms_mlp(x))
	680	mlp_out, router_stats = self.mlp(mlp_in)		687	mlp_out, router_stats = self.mlp(mlp_in)
	766	is_last = i == num_blocks - 1		773	is_last = i == num_blocks - 1
	767	is_long = i % 4 == 3 or is_last		774	is_long = i % 4 == 3 or is_last
	768	layer_mask = long_mask if is_long else short_mask		775	layer_mask = long_mask if is_long else short_mask
n	769	use_pko = is_long and not cfg.disable_pko	n
	770	disable_rope = is_long and cfg.disable_long_rope
	771	hidden, router_stats = eqx.filter_checkpoint(block, policy=remat_policy)(		776	hidden, router_stats = eqx.filter_checkpoint(block, policy=remat_policy)(
n	772	hidden, layer_mask, use_pko, disable_rope	n	777	hidden,
				778	layer_mask,
				779	LayerAttentionMode.LONG if is_long else LayerAttentionMode.SHORT,
	773	)		780	)
	774	moe_router_stats.append(router_stats)		781	moe_router_stats.append(router_stats)
	775			782
	823	)		830	)
	824	# No load-balancing loss; router z-loss only.		831	# No load-balancing loss; router z-loss only.
	825	num_moe_layers = router_metrics["router_z_loss_per_layer"].shape[0]		832	num_moe_layers = router_metrics["router_z_loss_per_layer"].shape[0]
t	826	rzl = jnp.sum(router_metrics["router_z_loss_per_layer"]) / num_moe_layers	t	833	router_z_loss = jnp.sum(router_metrics["router_z_loss_per_layer"]) / num_moe_layers
	827	aux_loss = self.config.router_z_loss_coef * rzl		834	aux_loss = self.config.router_z_loss_coef * router_z_loss
	828	loss = cross_entropy_loss + aux_loss if reduction != "none" else cross_entropy_loss		835	loss = cross_entropy_loss + aux_loss if reduction != "none" else cross_entropy_loss
	829	if return_router_metrics:		836	if return_router_metrics:
	830	summarized_metrics = _summarize_router_metrics(router_metrics)		837	summarized_metrics = _summarize_router_metrics(router_metrics)

train.py

changed

	experiments/grug/moe/train.py			experiments/grug/moe_pipeline/train.py
f	1	# Copyright The Marin Authors	f	1	# Copyright The Marin Authors
	2	# SPDX-License-Identifier: Apache-2.0		2	# SPDX-License-Identifier: Apache-2.0
	3			3
n	4	import dataclasses	n	4	"""Pipeline-parallel training for the Grug MoE pipeline variant."""
	5	import functools		5
	6	import logging		6	# GRUG NOVERIFY: JaxPP's MPMD lowering requires a multi-process accelerator mesh.
				7
				8	import json
				9	import os
				10	import statistics
	7	import time		11	import time
n	8	from dataclasses import dataclass, field	n	12	from dataclasses import dataclass
				13	from enum import StrEnum
				14	from types import SimpleNamespace
				15	from typing import cast
	9			16
n	10	import equinox as eqx	n
	11	import jax
	12	import jax.numpy as jnp
	13	import jmp
	14	import levanter.callbacks as callbacks
	15	import levanter.tracker
	16	import optax
	17	from fray.cluster import ResourceConfig		17	from fray.cluster import ResourceConfig
n	18	from haliax import Axis	n
	19	from haliax.partitioning import set_mesh
	20	from jax.sharding import Mesh, NamedSharding
	21	from jax.sharding import PartitionSpec as P
	22	from jax.tree_util import register_dataclass
	23	from jaxtyping import PRNGKeyArray
	24	from levanter.callbacks.state_adapter import StateCallbackRunner
	25	from levanter.callbacks.watch import WatchConfig, compute_watch_stats
	26	from levanter.data.dataset import AsyncDataset
	27	from levanter.data.loader import DataLoader
	28	from levanter.data.mixture import MixtureDataset, rescale_mixture_schedule_for_batch_schedule
	29	from levanter.data.text.datasets import LmDataConfig
	30	from levanter.data.text.examples import GrugLmExample, grug_lm_example_from_named
	31	from levanter.eval import TaggedEvaluator, cb_tagged_evaluate
	32	from levanter.grug.sharding import compact_grug_mesh
	33	from levanter.models.lm_model import LmExample
	34	from levanter.optim.config import AdamConfig, OptimizerConfig
	35	from levanter.schedule import BatchSchedule
	36	from levanter.trainer import TrainerConfig
	37	from levanter.training_control import TrainingDashboard
	38	from levanter.utils.flop_utils import lm_flops_per_token
	39	from levanter.utils.jax_utils import parameter_count
	40	from levanter.utils.logging import LoadingTimeTrackerIterator
	41			18
n	42	from experiments.grug.checkpointing import init_weights_only_from_checkpoint, restore_grug_state_from_checkpoint	n
	43	from experiments.grug.dispatch import dispatch_grug_training_run		19	from experiments.grug.dispatch import dispatch_grug_training_run
n	44	from experiments.grug.moe.model import BATCH_AXES, GrugModelConfig, Transformer	n
	45	from experiments.grug.sharding_dump import dump_grug_state_sharding_run_artifact
	46			20
n	47	# This file intentionally mirrors `experiments/grug/base/train.py` with	n
	48	# variant-specific model/loss/FLOP wiring, per the grug copy-first workflow in
	49	# `.agents/skills/change-grug/`.
	50			21
n	51	logger = logging.getLogger(__name__)	n	22	class PipelineSchedule(StrEnum):
				23	ZERO_BUBBLE = "automatic_zero_bubble"
				24	DUALPIPE_V = "automatic_dualpipe_v"
				25
				26	@property
				27	def automatic_schedule(self):
				28	from experiments.grug.moe_pipeline.pipeline import AutomaticPipelineSchedule # noqa: PLC0415
				29
				30	if self == self.DUALPIPE_V:
				31	return AutomaticPipelineSchedule.DUALPIPE_V
				32	if self == self.ZERO_BUBBLE:
				33	return AutomaticPipelineSchedule.ZERO_BUBBLE
				34	raise ValueError(f"unknown pipeline schedule: {self.value}")
	52			35
	53			36
	54	@dataclass(frozen=True)		37	@dataclass(frozen=True)
n	55	class GrugTrainerConfig:	n	38	class GrugPipelineTrainConfig:
	56	"""Runtime knobs for grug training."""		39	"""Model, topology, and loop settings for pipeline-parallel training."""
	57			40
n	58	trainer: TrainerConfig = field(default_factory=lambda: TrainerConfig(use_explicit_mesh_axes=True))	n	41	stages: int
	59	data_seed: int \| None = None		42	physical_stages: int
	60	log_every: int = 1		43	microbatches: int
	61	ema_beta: float \| None = None # EMA coefficient for eval/checkpoint model; None disables EMA.		44	batch_size: int
	62	z_loss_weight: float = 1e-4 # Weight on final-logit logsumexp z-loss stabilization term.		45	seq_len: int
	63			46	hidden_dim: int
	64	# Grug builds its own compact (replica_dcn, data, expert, model) mesh instead of using		47	intermediate_dim: int
	65	# the Trainer's logical axis mapping; `data` absorbs whatever these two leave free.		48	shared_expert_intermediate_dim: int
	66	# Defaults reproduce the historical layout: no expert parallelism and full replication		49	num_layers: int
	67	# across slices (replica_axis_size=None -> jax.process_count()), i.e. parameters		50	num_experts: int
	68	# replicated per slice and sharded only over the intra-slice `data` axis. For a model		51	top_k: int
	69	# too large to replicate within one slice, set replica_axis_size=1 (FSDP across every
	70	# slice) and expert_axis_size>1 (expert parallelism over the intra-slice devices).
	71	expert_axis_size: int = 1		52	expert_axis_size: int
	72	replica_axis_size: int \| None = None		53	vocab_size: int
	73	sharding_dump_path: str \| None = None		54	num_heads: int
				55	num_kv_heads: int
				56	sliding_window: int
				57	qk_mult: float
				58	steps: int
				59	warmup_steps: int
				60	profile_start_step: int
				61	profile_steps: int
				62	profile_run_id: str \| None
				63	layer_counts: tuple[int, ...] \| None
				64	memory_threshold: int \| None
				65	mp_policy_string: str
				66	remat_mode: str
				67	attention_implementation: str
				68	moe_implementation: str
				69	schedule: PipelineSchedule
	74			70
n	75		n	71	def __post_init__(self) -> None:
	76	@dataclass(frozen=True)		72	if self.schedule == PipelineSchedule.ZERO_BUBBLE and self.stages != self.physical_stages:
	77	class GrugEvalConfig:		73	raise ValueError(
	78	"""Perplexity eval settings for grug training."""		74	"automatic_zero_bubble requires one logical stage per physical stage; "
	79			75	f"got {self.stages} logical and {self.physical_stages} physical stages"
	80	eval_batch_size: int = 512		76	)
	81	steps_per_eval: int \| None = 1000		77	if self.schedule == PipelineSchedule.DUALPIPE_V:
	82	max_eval_batches: int \| None = None		78	if self.stages != 2 * self.physical_stages:
	83	prefix: str = "eval"		79	raise ValueError(
	84	eval_current: bool = True		80	"automatic_dualpipe_v requires two logical stages per physical stage; "
	85	eval_ema: bool = True		81	f"got {self.stages} logical and {self.physical_stages} physical stages"
	86	compute_bpb: bool = True		82	)
	87	shuffle: bool = True # Mix eval domains within batches to reduce expert-capacity drops.		83	if self.microbatches < self.stages:
				84	raise ValueError(
				85	f"automatic_dualpipe_v requires at least {self.stages} microbatches, got {self.microbatches}"
				86	)
	88			87
	89			88
	90	@dataclass(frozen=True)		89	@dataclass(frozen=True)
	91	class GrugRunConfig:		90	class GrugRunConfig:
n	92	"""Top-level config for grug training."""	n	91	"""Resources and identity for one pipeline-parallel training run."""
	93			92
n	94	model: GrugModelConfig	n	93	run_id: str
	95	data: LmDataConfig		94	train: GrugPipelineTrainConfig
	96	resources: ResourceConfig		95	resources: ResourceConfig
n	97	optimizer: OptimizerConfig = field(default_factory=AdamConfig)	n
	98	trainer: GrugTrainerConfig = field(default_factory=GrugTrainerConfig)
	99	eval: GrugEvalConfig \| None = field(default_factory=GrugEvalConfig)
	100	# GPU processes per task: > 1 runs one JAX process per GPU (multi-controller)
	101	# via the iris.hooks.multigpu_main supervisor instead of one process per node.
	102	processes_per_task: int = 1
	103			96
	104			97
n	105	def build_train_dataset(	n	98	def _log(event: str, **values) -> None:
	106	data_config: LmDataConfig,		99	print(f"{event} {json.dumps(values, sort_keys=True, default=str)}", flush=True)
				100
				101
				102	def _validate_local_mesh(
	107	*,		103	*,
n	108	max_seq_len: int,	n	104	local_device_count: int,
	109	batch_schedule: BatchSchedule,		105	expert_axis_size: int,
	110	key: PRNGKeyArray,		106	batch_size: int,
	111	) -> MixtureDataset[GrugLmExample]:		107	microbatches: int,
	112	pos = Axis("position", max_seq_len)		108	) -> None:
	113	mix_key, shuffle_key = jax.random.split(key)		109	if local_device_count % expert_axis_size != 0:
	114	weights = data_config.train_weights		110	raise ValueError(f"local device count {local_device_count} must be divisible by expert axis {expert_axis_size}")
	115	if isinstance(weights, list):		111	if batch_size % microbatches != 0:
	116	weights = rescale_mixture_schedule_for_batch_schedule(weights, batch_schedule)		112	raise ValueError(f"batch size {batch_size} must be divisible by {microbatches} microbatches")
	117			113	microbatch_size = batch_size // microbatches
	118	initial_batch_size = batch_schedule.batch_size_at_step(0)		114	if microbatch_size % local_device_count != 0:
	119	datasets = data_config.train_sets(pos, key=shuffle_key, initial_batch_size=initial_batch_size)		115	raise ValueError(
	120	return MixtureDataset(		116	f"microbatch size {microbatch_size} must be divisible by the {local_device_count} devices in each stage"
	121	datasets=datasets,
	122	weights=weights,
	123	stop_strategy=data_config.stop_strategy,
	124	key=mix_key,
	125	block_size=data_config.mixture_block_size,
	126	)
	127
	128
	129	def build_train_loader(
	130	dataset: AsyncDataset[GrugLmExample],
	131	*,
	132	batch_schedule: BatchSchedule,
	133	mesh: Mesh,
	134	) -> DataLoader[GrugLmExample]:
	135	# DataLoader uses this batch axis mapping to shard batches across the distributed mesh.
	136	# `compact_grug_mesh` always carries (replica_dcn, data, expert, model); length-1 axes
	137	# are kept so we can name "expert" unconditionally.
	138	return DataLoader(
	139	dataset,
	140	batch_schedule.schedule,
	141	mesh=mesh,
	142	axis_resources={"__BATCH__": BATCH_AXES},
	143	batch_axis_name="__BATCH__",
	144	allow_nondivisible_batch_size=False,
	145	)
	146
	147
	148	def build_tagged_evaluator(
	149	*,
	150	data_config: LmDataConfig,
	151	max_seq_len: int,
	152	mesh: Mesh,
	153	eval_cfg: GrugEvalConfig,
	154	) -> TaggedEvaluator[LmExample \| GrugLmExample, Transformer] \| None:
	155	pos = Axis("position", max_seq_len)
	156	tagged_eval_sets = data_config.tagged_eval_sets(pos)
	157	if len(tagged_eval_sets) == 0:
	158	logger.warning("No evaluation datasets provided.")
	159	return None
	160
	161	max_examples_per_dataset = None
	162	if eval_cfg.max_eval_batches is not None:
	163	max_examples_per_dataset = eval_cfg.max_eval_batches * eval_cfg.eval_batch_size
	164
	165	tokenizer = data_config.the_tokenizer if eval_cfg.compute_bpb else None
	166	# `compact_grug_mesh` always carries (replica_dcn, data, expert, model); length-1 axes
	167	# are kept so we can name "expert" unconditionally.
	168	eval_axis_mapping = {"batch": BATCH_AXES}
	169	eval_batch = Axis("batch", eval_cfg.eval_batch_size)
	170	eval_array_sharding = NamedSharding(mesh, P(BATCH_AXES, None))
	171
	172	def eval_loss_fn(model: Transformer, batch: LmExample \| GrugLmExample) -> tuple[jax.Array, jax.Array, jax.Array]:
	173	if isinstance(batch, LmExample):
	174	batch = grug_lm_example_from_named(batch)
	175	per_pos_loss = model.next_token_loss(
	176	batch.tokens,
	177	batch.loss_weight,
	178	mask=batch.attn_mask,
	179	reduction="none",
	180	logsumexp_weight=None,
	181	)		117	)
n	182	per_pos_loss = jax.sharding.reshard(per_pos_loss, eval_array_sharding)	n
	183	per_pos_weight = jax.sharding.reshard(batch.loss_weight, eval_array_sharding)
	184	per_pos_token_id = jnp.pad(batch.tokens[:, 1:], ((0, 0), (0, 1)))
	185	return per_pos_loss, per_pos_weight, per_pos_token_id
	186			118
n	187	return TaggedEvaluator(	n
	188	EvalBatch=eval_batch,
	189	tagged_eval_sets=tagged_eval_sets,
	190	loss_fn=eval_loss_fn,
	191	tokenizer=tokenizer,
	192	device_mesh=mesh,
	193	axis_mapping=eval_axis_mapping,
	194	max_examples_per_dataset=max_examples_per_dataset,
	195	shuffle=eval_cfg.shuffle,
	196	)
	197			119
n			n	120	def _run_grug_local(config: GrugPipelineTrainConfig) -> None:
				121	os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
	198			122
n	199	def _compute_flops(	n	123	from iris.runtime.jax_init import initialize_jax # noqa: PLC0415
	200	*,		124
				125	initialize_jax()
				126
				127	import jax # noqa: PLC0415
				128	import jax.numpy as jnp # noqa: PLC0415
				129	import jmp # noqa: PLC0415
				130	import numpy as np # noqa: PLC0415
				131	import optax # noqa: PLC0415
				132	from fray.device_flops import device_flops # noqa: PLC0415
				133	from jax.sharding import NamedSharding # noqa: PLC0415
				134	from jax.sharding import PartitionSpec as P # noqa: PLC0415
				135	from jaxpp.array import MpmdArray # noqa: PLC0415 # pyrefly: ignore[missing-import]
				136	from levanter.data.text.examples import GrugLmExample # noqa: PLC0415
				137	from levanter.pipeline import reshape_batch_into_microbatches # noqa: PLC0415
				138	from levanter.utils.flop_utils import lm_flops_per_token # noqa: PLC0415
				139
				140	from experiments.grug.moe_pipeline.model import BATCH_AXES, GrugModelConfig, Transformer # noqa: PLC0415
				141	from experiments.grug.moe_pipeline.pipeline import ( # noqa: PLC0415
				142	TRAIN_LOSS_KEY,
				143	GrugMoePipelineConfig,
				144	automatic_stage_to_mpmd_indices,
				145	initialize_mpmd_automatic_pipeline_state,
				146	make_automatic_pipeline_step,
				147	make_pipeline_mesh,
				148	prepare_automatic_mpmd_step,
				149	)
				150
				151	mp_policy = jmp.get_policy(config.mp_policy_string)
				152	pipeline_config = GrugMoePipelineConfig(
				153	stages=config.stages,
				154	microbatches=config.microbatches,
				155	physical_stages=None if config.physical_stages == config.stages else config.physical_stages,
				156	)
				157
				158	if jax.process_count() != config.physical_stages:
				159	raise ValueError(
				160	f"expected one process per physical stage ({config.physical_stages}), got {jax.process_count()}"
				161	)
				162	_validate_local_mesh(
				163	local_device_count=jax.local_device_count(),
				164	expert_axis_size=config.expert_axis_size,
				165	batch_size=config.batch_size,
				166	microbatches=config.microbatches,
				167	)
				168
	201	model_config: GrugModelConfig,		169	model_config = GrugModelConfig(
	202	) -> tuple[float, dict[str, float]]:		170	vocab_size=config.vocab_size,
				171	hidden_dim=config.hidden_dim,
				172	intermediate_dim=config.intermediate_dim,
				173	shared_expert_intermediate_dim=config.shared_expert_intermediate_dim,
				174	num_layers=config.num_layers,
				175	num_experts=config.num_experts,
				176	num_experts_per_token=config.top_k,
				177	num_heads=config.num_heads,
				178	num_kv_heads=config.num_kv_heads,
				179	max_seq_len=config.seq_len,
				180	sliding_window=config.sliding_window,
				181	qk_mult=config.qk_mult,
				182	router_z_loss_coef=0.0,
				183	attention_implementation=cast(str, config.attention_implementation),
				184	moe_implementation=cast(str, config.moe_implementation),
				185	remat_mode=cast(str, config.remat_mode),
				186	)
				187	mesh, mpmd_mesh = make_pipeline_mesh(
				188	pipeline_config,
				189	expert_axis_size=config.expert_axis_size,
				190	replica_axis_size=1,
				191	)
				192	optimizer = optax.adamw(learning_rate=1e-4, b1=0.9, b2=0.95, weight_decay=0.1)
				193	peak_flops_per_device = device_flops("h100")
				194	if peak_flops_per_device is None:
				195	raise ValueError("Fray does not define H100 BF16 peak FLOP/s")
				196
				197	_log(
				198	"PIPELINE_CONFIG",
				199	process_index=jax.process_index(),
				200	process_count=jax.process_count(),
				201	global_devices=jax.device_count(),
				202	local_devices=jax.local_device_count(),
				203	mesh_shape=dict(mesh.shape),
				204	stages=config.stages,
				205	physical_stages=config.physical_stages,
				206	microbatches=config.microbatches,
				207	batch_size=config.batch_size,
				208	seq_len=config.seq_len,
				209	hidden_dim=config.hidden_dim,
				210	intermediate_dim=model_config.intermediate_dim,
				211	shared_expert_intermediate_dim=model_config.shared_expert_intermediate_dim,
				212	num_layers=config.num_layers,
				213	layers_per_stage=config.layer_counts,
				214	num_experts=config.num_experts,
				215	top_k=config.top_k,
				216	vocab_size=model_config.vocab_size,
				217	num_heads=model_config.num_heads,
				218	num_kv_heads=model_config.num_kv_heads,
				219	sliding_window=model_config.sliding_window,
				220	qk_mult=model_config.qk_mult,
				221	steps=config.steps,
				222	profile_start_step=config.profile_start_step,
				223	profile_steps=config.profile_steps,
				224	mp_policy=config.mp_policy_string,
				225	remat_mode=config.remat_mode,
				226	schedule=config.schedule,
				227	jax_version=jax.__version__,
				228	)
				229
				230	init_started = time.monotonic()
				231	with jax.set_mesh(mesh):
				232	model = mp_policy.cast_to_param(Transformer.init(model_config, key=jax.random.PRNGKey(0)))
				233	batch_sharding = NamedSharding(mesh, P(BATCH_AXES, None))
				234	token_row = np.arange(config.seq_len, dtype=np.int32) % model_config.vocab_size
				235	host_tokens = np.broadcast_to(token_row, (config.batch_size, config.seq_len)).copy()
				236	host_loss_weight = np.ones((config.batch_size, config.seq_len), dtype=np.float32)
				237	host_loss_weight[:, -1] = 0
				238	batch = GrugLmExample(
				239	tokens=jax.device_put(host_tokens, batch_sharding),
				240	loss_weight=jax.device_put(host_loss_weight, batch_sharding),
				241	)
				242	loss_denominator = jnp.sum(batch.loss_weight.astype(jnp.float32))
				243	batches = reshape_batch_into_microbatches(batch, config.microbatches)
				244	if config.schedule == PipelineSchedule.DUALPIPE_V:
				245	stage_to_mpmd_index = automatic_stage_to_mpmd_indices(pipeline_config, config.schedule.automatic_schedule)
				246	else:
				247	stage_to_mpmd_index = None
				248	state, static_stages = initialize_mpmd_automatic_pipeline_state(
				249	model,
				250	optimizer,
				251	mpmd_mesh,
				252	num_stages=config.stages,
				253	layer_counts=config.layer_counts,
				254	stage_to_mpmd_index=stage_to_mpmd_index,
				255	memory_threshold=config.memory_threshold,
				256	)
				257	jax.block_until_ready((state, batches, loss_denominator))
				258	_log(
				259	"PIPELINE_INIT",
				260	process_index=jax.process_index(),
				261	elapsed_seconds=time.monotonic() - init_started,
				262	)
				263
				264	build_started = time.monotonic()
				265	step = make_automatic_pipeline_step(
				266	optimizer,
				267	mp_policy,
				268	static_stages,
				269	state,
				270	batches,
				271	config=pipeline_config,
				272	mpmd_mesh=mpmd_mesh,
				273	schedule_name=config.schedule.automatic_schedule,
				274	)
				275	_log(
				276	"PIPELINE_BUILD",
				277	process_index=jax.process_index(),
				278	elapsed_seconds=time.monotonic() - build_started,
				279	)
				280	lower_started = time.monotonic()
				281	prepared = prepare_automatic_mpmd_step(
				282	step,
				283	state,
				284	batches,
				285	loss_denominator,
				286	mpmd_mesh,
				287	memory_threshold=config.memory_threshold,
				288	)
				289	step = prepared.step
				290	state = prepared.state
				291	batches = prepared.batches
				292	loss_denominator = prepared.loss_denominator
				293	_log(
				294	"PIPELINE_LOWER",
				295	process_index=jax.process_index(),
				296	elapsed_seconds=time.monotonic() - lower_started,
				297	)
				298
				299	profiler_callback = None
				300	if config.profile_steps > 0:
				301	from levanter.callbacks.profiler import ProfileOptionsConfig, ProfilerConfig # noqa: PLC0415
				302
				303	profiler_callback = ProfilerConfig(
				304	enabled=True,
				305	start_step=config.profile_start_step,
				306	num_steps=config.profile_steps,
				307	perfetto_link=False,
				308	profile_options=ProfileOptionsConfig(
				309	host_tracer_level=1,
				310	python_tracer_level=0,
				311	enable_hlo_proto=True,
				312	),
				313	).build(
				314	"/tmp/grug-moe-pipeline-profiler",
				315	run_id=config.profile_run_id,
				316	)
				317	profiler_callback(SimpleNamespace(step=-1))
				318
				319	step_times = []
				320	loss = None
				321	for step_index in range(config.steps):
				322	started = time.monotonic()
				323	state, metrics = step(state, batches, loss_denominator)
				324	jax.block_until_ready((state, metrics))
				325	elapsed = time.monotonic() - started
				326	step_times.append(elapsed)
				327	metric_loss = metrics[TRAIN_LOSS_KEY]
				328	if profiler_callback is not None:
				329	profiler_callback(SimpleNamespace(step=step_index))
				330	if isinstance(metric_loss, MpmdArray):
				331	if not metric_loss.is_partially_addressable:
				332	continue
				333	local_loss = metric_loss.to_mpmd_local_array
				334	if not isinstance(local_loss, jax.Array):
				335	raise TypeError(f"expected one local loss array, got {type(local_loss).__name__}")
				336	else:
				337	local_loss = metric_loss
				338	loss = float(local_loss)
				339	_log(
				340	"PIPELINE_STEP",
				341	step=step_index,
				342	elapsed_seconds=elapsed,
				343	tokens_per_second=config.batch_size * config.seq_len / elapsed,
				344	loss=loss,
				345	)
				346
				347	measured = step_times[config.warmup_steps :]
				348	if not measured:
				349	measured = step_times
				350	mean_step_seconds = sum(measured) / len(measured)
				351	median_step_seconds = statistics.median(measured)
				352	tokens_per_second = config.batch_size * config.seq_len / mean_step_seconds
				353	median_tokens_per_second = config.batch_size * config.seq_len / median_step_seconds
	203	flops_per_token = lm_flops_per_token(		354	forward_flops_per_token = lm_flops_per_token(
	204	hidden_dim=model_config.hidden_dim,		355	hidden_dim=model_config.hidden_dim,
	205	intermediate_dim=model_config.intermediate_dim,		356	intermediate_dim=model_config.intermediate_dim,
	206	shared_intermediate_dim=model_config.shared_expert_intermediate_dim,		357	shared_intermediate_dim=model_config.shared_expert_intermediate_dim,
	214	num_shared_experts=1 if model_config.shared_expert_intermediate_dim > 0 else 0,		365	num_shared_experts=1 if model_config.shared_expert_intermediate_dim > 0 else 0,
	215	num_experts_per_tok=model_config.num_experts_per_token,		366	num_experts_per_tok=model_config.num_experts_per_token,
	216	)		367	)
n	217	flops_per_example = 3 * flops_per_token * model_config.max_seq_len	n	368	analytic_mfu = 3 * forward_flops_per_token * tokens_per_second / (jax.device_count() * peak_flops_per_device)
	218			369	median_analytic_mfu = (
	219	flops_summary: dict[str, float] = {		370	3 * forward_flops_per_token * median_tokens_per_second / (jax.device_count() * peak_flops_per_device)
	220	"throughput/flops_per_token_analytic": flops_per_token,
	221	"throughput/flops_per_example_analytic": flops_per_example,
	222	}		371	)
	223
	224	return flops_per_example, flops_summary
	225
	226
	227	def _make_mixture_stage_callback(train_dataset: MixtureDataset, batch_schedule: BatchSchedule):
	228	last_mixture_stage = -1
	229
	230	def log_mixture_stage(step_info):
	231	nonlocal last_mixture_stage
	232	seq_index = batch_schedule.global_data_offset_by_step(step_info.step)
	233	block_id = seq_index // train_dataset.block_size
	234	stage = train_dataset._get_stage_for_block(block_id)
	235	if stage == last_mixture_stage:
	236	return
	237
	238	weights = train_dataset.weight_stages[stage][1]
	239	mixture_log = {f"mixture/weight/{name}": weight for name, weight in weights.items()}
	240	mixture_log["mixture/stage"] = stage
	241	levanter.tracker.log(mixture_log, step=step_info.step)
	242	last_mixture_stage = stage
	243
	244	return log_mixture_stage
	245
	246
	247	@register_dataclass
	248	@dataclass(frozen=True)
	249	class GrugTrainState:
	250	step: jax.Array
	251	params: Transformer
	252	opt_state: optax.OptState
	253	ema_params: Transformer \| None
	254	pending_qb_betas: jax.Array
	255
	256
	257	def _apply_qb_betas(model: Transformer, qb_betas: jax.Array) -> Transformer:
	258	"""Set router biases from QB betas (computed on previous step)."""
	259	new_blocks = list(model.blocks)
	260	moe_idx = 0
	261	for i, block in enumerate(model.blocks):
	262	if block.mlp is None:
	263	continue
	264	new_bias = -qb_betas[moe_idx]
	265	new_bias = new_bias - jnp.mean(new_bias)
	266	new_mlp = eqx.tree_at(lambda m: m.router_bias, block.mlp, new_bias)
	267	new_blocks[i] = eqx.tree_at(lambda b: b.mlp, block, new_mlp)
	268	moe_idx += 1
	269	return eqx.tree_at(lambda t: t.blocks, model, tuple(new_blocks))
	270
	271
	272	def initial_state(
	273	model_config: GrugModelConfig,
	274	*,
	275	optimizer: optax.GradientTransformation,
	276	mp: jmp.Policy,
	277	key: PRNGKeyArray,
	278	ema_beta: float \| None,
	279	) -> GrugTrainState:
	280	params = mp.cast_to_param(Transformer.init(model_config, key=key))
	281	num_moe_layers = sum(1 for b in params.blocks if b.mlp is not None)
	282	return GrugTrainState(
	283	step=jnp.array(0, dtype=jnp.int32),
	284	params=params,
	285	opt_state=optimizer.init(params),
	286	ema_params=params if ema_beta is not None else None,
	287	pending_qb_betas=jnp.zeros((num_moe_layers, model_config.num_experts)),
	288	)
	289
	290
	291	def _make_train_step(
	292	optimizer: optax.GradientTransformation,
	293	mp: jmp.Policy,
	294	*,
	295	z_loss_weight: float,
	296	ema_beta: float \| None,
	297	watch_config: WatchConfig \| None = None,
	298	):
	299	one = jnp.array(1, dtype=jnp.int32)
	300	z_loss = z_loss_weight if z_loss_weight > 0 else None
	301	if watch_config is not None:		372	if loss is not None:
	302	if isinstance(watch_config.watch_targets, str):		373	_log(
	303	watch_targets = tuple(t.strip() for t in watch_config.watch_targets.split(","))		374	"PIPELINE_SUMMARY",
	304	else:		375	mean_step_seconds=mean_step_seconds,
	305	watch_targets = tuple(watch_config.watch_targets)		376	median_step_seconds=median_step_seconds,
	306	else:		377	tokens_per_second=tokens_per_second,
	307	watch_targets = ()		378	median_tokens_per_second=median_tokens_per_second,
				379	analytic_mfu=analytic_mfu,
				380	median_analytic_mfu=median_analytic_mfu,
				381	measured_step_seconds=measured,
				382	loss=loss,
				383	measured_steps=len(measured),
				384	)
	308			385
n	309	@functools.partial(jax.jit, donate_argnums=(0,), static_argnames=("compute_watch",))	n	386	for device in jax.local_devices():
	310	def train_step(state: GrugTrainState, batch, *, compute_watch: bool = False):		387	stats = device.memory_stats()
	311	# Apply pending QB betas to router biases inside JIT (avoids eager
	312	# host-side TPU kernel launches that can cause SPMD sync issues).
	313	qb_params = _apply_qb_betas(state.params, state.pending_qb_betas)
	314	if ema_beta is not None:		388	if stats is not None:
	315	qb_ema_params = _apply_qb_betas(state.ema_params, state.pending_qb_betas)		389	_log(
	316	else:		390	"PIPELINE_MEMORY",
	317	qb_ema_params = None		391	process_index=jax.process_index(),
	318			392	device_id=device.id,
	319	def loss_fn(params):		393	peak_bytes_in_use=stats.get("peak_bytes_in_use"),
	320	compute_params = mp.cast_to_compute(params)		394	bytes_limit=stats.get("bytes_limit"),
	321	return compute_params.next_token_loss(
	322	batch.tokens,
	323	batch.loss_weight,
	324	mask=batch.attn_mask,
	325	reduction="mean",
	326	logsumexp_weight=z_loss,
	327	return_router_metrics=True,
	328	)		395	)
	329			396
n	330	(loss, summarized_metrics), grads = jax.value_and_grad(loss_fn, has_aux=True)(qb_params)	n
	331	metrics = {"train/loss": loss, **summarized_metrics}
	332	updates, opt_state = optimizer.update(grads, state.opt_state, qb_params)
	333	params = optax.apply_updates(qb_params, updates)
	334
	335	if ema_beta is None:
	336	ema_params = None
	337	else:
	338	if qb_ema_params is None:
	339	raise ValueError("ema_params must be initialized when ema_beta is set.")
	340	ema_params = jax.tree_util.tree_map(
	341	lambda old, new: ema_beta * old + (1.0 - ema_beta) * new,
	342	qb_ema_params,
	343	params,
	344	)
	345
	346	watch_stats = None
	347	if watch_config is not None and compute_watch:
	348	watch_stats = compute_watch_stats(
	349	watch_targets=watch_targets,
	350	include_norms=watch_config.include_norms,
	351	include_per_parameter_norms=watch_config.include_per_parameter_norms,
	352	include_histogram=watch_config.include_histograms,
	353	split_scan_layers=watch_config.split_scan_layers,
	354	params=qb_params,
	355	grads=grads,
	356	updates=updates,
	357	opt_state=state.opt_state,
	358	model_tree_type=type(state.params),
	359	)
	360
	361	next_state = dataclasses.replace(
	362	state,
	363	step=state.step + one,
	364	params=params,
	365	opt_state=opt_state,
	366	ema_params=ema_params,
	367	pending_qb_betas=metrics["qb_beta_per_layer"],
	368	)
	369
	370	return next_state, metrics, watch_stats
	371
	372	return train_step
	373
	374
	375	def _run_grug_local(config: GrugRunConfig) -> None:
	376	"""Entry point for the grug template training loop."""
	377	trainer = config.trainer.trainer
	378	trainer.initialize()
	379	levanter.tracker.log_configuration(config)
	380
	381	run_id = trainer.id
	382	if run_id is None:
	383	raise ValueError("trainer.id was not initialized")
	384
	385	optimizer = config.optimizer.build(trainer.num_train_steps)
	386	watch_config = trainer.watch
	387	train_step = _make_train_step(
	388	optimizer,
	389	trainer.mp,
	390	z_loss_weight=config.trainer.z_loss_weight,
	391	ema_beta=config.trainer.ema_beta,
	392	watch_config=watch_config if watch_config.is_enabled else None,
	393	)
	394
	395	data_key, model_key = jax.random.split(jax.random.PRNGKey(trainer.seed), 2)
	396	if config.trainer.data_seed is not None:
	397	data_key = jax.random.PRNGKey(config.trainer.data_seed)
	398
	399	# Grug uses raw PartitionSpecs rather than Trainer's logical axis mapping.
	400	# Keep the mesh compact so the batch pspec derived by `_batch_spec(mesh)` spans slices directly.
	401	# replica_axis_size=None lets compact_grug_mesh default to jax.process_count() (full
	402	# cross-slice replication); set it to 1 on GrugTrainerConfig for cross-slice FSDP.
	403	mesh = compact_grug_mesh(
	404	expert_axis_size=config.trainer.expert_axis_size,
	405	replica_axis_size=config.trainer.replica_axis_size,
	406	)
	407	checkpointer = trainer.checkpointer.create(run_id)
	408	with set_mesh(mesh), TrainingDashboard(config, checkpointer.request_checkpoint, run_id):
	409	batch_schedule = trainer.batch_schedule
	410
	411	train_dataset = build_train_dataset(
	412	config.data,
	413	max_seq_len=config.model.max_seq_len,
	414	batch_schedule=batch_schedule,
	415	key=data_key,
	416	)
	417	train_loader = build_train_loader(
	418	train_dataset,
	419	batch_schedule=batch_schedule,
	420	mesh=mesh,
	421	)
	422
	423	@jax.jit
	424	def _init_state(model_rng):
	425	return initial_state(
	426	config.model,
	427	optimizer=optimizer,
	428	mp=trainer.mp,
	429	key=model_rng,
	430	ema_beta=config.trainer.ema_beta,
	431	)
	432
	433	state = _init_state(model_key)
	434
	435	state = restore_grug_state_from_checkpoint(
	436	state,
	437	checkpoint_search_paths=trainer.checkpoint_search_paths(run_id),
	438	load_checkpoint_setting=trainer.load_checkpoint,
	439	mesh=mesh,
	440	allow_partial=trainer.allow_partial_checkpoint,
	441	)
	442	if trainer.initialize_from is not None:
	443	state = init_weights_only_from_checkpoint(
	444	state,
	445	trainer.initialize_from,
	446	mesh=mesh,
	447	allow_partial=trainer.allow_partial_checkpoint,
	448	additional_weight_fields=("pending_qb_betas",),
	449	)
	450	dump_grug_state_sharding_run_artifact(
	451	state,
	452	log_dir=trainer.log_dir,
	453	run_id=run_id,
	454	path_override=config.trainer.sharding_dump_path,
	455	)
	456
	457	levanter.tracker.log_summary({"parameter_count": parameter_count(state.params)})
	458
	459	flops_per_example, flops_summary = _compute_flops(model_config=config.model)
	460	levanter.tracker.log_summary(flops_summary)
	461
	462	eval_cfg = config.eval
	463	evaluator = None
	464	if eval_cfg is not None:
	465	evaluator = build_tagged_evaluator(
	466	data_config=config.data,
	467	max_seq_len=config.model.max_seq_len,
	468	mesh=mesh,
	469	eval_cfg=eval_cfg,
	470	)
	471
	472	profiler_cfg = trainer.profiler
	473	profiler_num_steps = profiler_cfg.resolve_num_profile_steps(num_train_steps=trainer.num_train_steps)
	474	profiler_enabled = profiler_cfg.is_enabled and profiler_num_steps > 0
	475
	476	log_every = max(1, config.trainer.log_every)
	477	iterator = LoadingTimeTrackerIterator(train_loader.iter_from_step(int(state.step)))
	478
	479	state_callbacks = StateCallbackRunner[GrugTrainState](
	480	step_getter=lambda s: s.step,
	481	model_getter=lambda s: s.params,
	482	eval_model_getter=lambda s: s.ema_params if s.ema_params is not None else s.params,
	483	opt_state_getter=lambda s: s.opt_state,
	484	)
	485	state_callbacks.add_hook(
	486	callbacks.log_performance_stats(config.model.max_seq_len, batch_schedule, flops_per_example),
	487	every=log_every,
	488	)
	489	state_callbacks.add_hook(callbacks.pbar_logger(total=trainer.num_train_steps), every=log_every)
	490	state_callbacks.add_hook(callbacks.log_step_info(trainer.num_train_steps), every=log_every)
	491	if profiler_enabled:
	492	state_callbacks.add_hook(
	493	profiler_cfg.build(
	494	str(trainer.log_dir / run_id / "profiler"),
	495	run_id=run_id,
	496	num_steps=profiler_num_steps,
	497	),
	498	every=1,
	499	)
	500	state_callbacks.add_hook(_make_mixture_stage_callback(train_dataset, batch_schedule), every=1)
	501	if evaluator is not None and eval_cfg is not None:
	502	interval = eval_cfg.steps_per_eval
	503	eval_ema = eval_cfg.eval_ema and config.trainer.ema_beta is not None
	504	if interval is not None and interval > 0 and (eval_cfg.eval_current or eval_ema):
	505	state_callbacks.add_hook(
	506	cb_tagged_evaluate(
	507	evaluator,
	508	prefix=eval_cfg.prefix,
	509	eval_current=eval_cfg.eval_current,
	510	eval_ema=eval_ema,
	511	),
	512	every=interval,
	513	)
	514
	515	last_loss: float \| jax.Array = 0.0
	516	last_step_duration = 0.0
	517
	518	# Main optimization loop.
	519	try:
	520	while int(state.step) < trainer.num_train_steps:
	521	with jax.profiler.TraceAnnotation("load_batch"):
	522	batch = next(iterator)
	523	step_start = time.perf_counter()
	524	current_step = int(state.step)
	525	# grad_watch runs only on its configured interval.
	526	compute_watch = (
	527	watch_config.is_enabled and watch_config.interval > 0 and current_step % watch_config.interval == 0
	528	)
	529	state, metrics, watch_stats = train_step(state, batch, compute_watch=compute_watch)
	530	step = int(state.step) - 1
	531
	532	jax.block_until_ready(metrics["train/loss"])
	533
	534	if jnp.isnan(metrics["train/loss"]):
	535	logger.error(f"NaN loss at step {int(state.step)}. Stopping training.")
	536	break
	537	duration = time.perf_counter() - step_start
	538	hook_start = time.perf_counter()
	539	with jax.profiler.TraceAnnotation("callbacks"):
	540	state_callbacks.run(state, loss=metrics["train/loss"], step_duration=duration)
	541	last_loss = metrics["train/loss"]
	542	last_step_duration = duration
	543	levanter.tracker.log({"throughput/hook_time": time.perf_counter() - hook_start}, step=step)
	544	levanter.tracker.log({"throughput/loading_time": iterator.this_load_time}, step=step)
	545	router_metrics = {
	546	key: value
	547	for key, value in metrics.items()
	548	if (key.startswith("train/router/") or key.startswith("moe_bias/"))
	549	and key not in ("train/router/routing_counts_per_layer", "qb_beta_per_layer")
	550	}
	551	if router_metrics:
	552	levanter.tracker.log(router_metrics, step=step)
	553	if "train/cross_entropy_loss" in metrics:
	554	levanter.tracker.log(
	555	{"train/cross_entropy_loss": metrics["train/cross_entropy_loss"]},
	556	step=step,
	557	)
	558
	559	if watch_stats is not None:
	560	levanter.tracker.log(watch_stats, step=step)
	561
	562	if checkpointer is not None:
	563	checkpointer.on_step(tree=state, step=int(state.step))
	564	except BaseException:
	565	logger.exception(
	566	"Fatal error in grug training loop; skipping final callbacks/checkpoint to preserve root cause"
	567	)
	568	raise
	569	else:
	570	# Mirror classic trainer behavior: force callbacks on the last completed step.
	571	state_callbacks.run(state, loss=last_loss, step_duration=last_step_duration, force=True)
	572	if checkpointer is not None:
	573	checkpointer.on_step(tree=state, step=int(state.step), force=True)
	574	checkpointer.wait_until_finished()
	575
	576	levanter.tracker.current_tracker().finish()
	577
	578			397
	579	def run_grug(config: GrugRunConfig) -> None:		398	def run_grug(config: GrugRunConfig) -> None:
n	580	"""Dispatch grug training through Fray jobs."""	n	399	"""Dispatch pipeline-parallel Grug training through Fray."""
	581	trainer = config.trainer.trainer		400	if config.resources.replicas != config.train.physical_stages:
	582	if trainer.id is None:		401	raise ValueError(
	583	raise ValueError("trainer.id must be set before dispatching grug training.")		402	f"training resources must have {config.train.physical_stages} replicas, got {config.resources.replicas}"
	584			403	)
	585	dispatch_grug_training_run(		404	dispatch_grug_training_run(
n	586	run_id=trainer.id,	n	405	run_id=config.run_id,
	587	config=config,		406	config=config.train,
	588	local_entrypoint=_run_grug_local,		407	local_entrypoint=_run_grug_local,
	589	resources=config.resources,		408	resources=config.resources,
n	590	processes_per_task=config.processes_per_task,	n	409	extras=["pipeline"],
	591	)		410	)
t	592		t
	593
	594	__all__ = [
	595	"GrugEvalConfig",
	596	"GrugRunConfig",
	597	"GrugTrainState",
	598	"GrugTrainerConfig",
	599	"initial_state",
	600	"run_grug",
	601	]

benchmark.py

added

+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+"""Environment-configured benchmark for the Grug MoE pipeline trainer."""
+import math
+import os
+from collections.abc import Mapping
+from experiments.grug.moe_pipeline.train import GrugPipelineTrainConfig, PipelineSchedule, _run_grug_local
+_ATTENTION_HEAD_DIM = 128
+def _env_int(environ: Mapping[str, str], name: str, default: int) -> int:
+    value = environ.get(name)
+    return default if not value else int(value)
+def _env_float(environ: Mapping[str, str], name: str, default: float) -> float:
+    value = environ.get(name)
+    return default if not value else float(value)
+def _default_num_kv_heads(num_heads: int) -> int:
+    for candidate in range(num_heads // 4, 0, -1):
+        if num_heads % candidate == 0:
+            return candidate
+    return 1
+def _resolve_benchmark_config(environ: Mapping[str, str]) -> GrugPipelineTrainConfig:
+    stages = _env_int(environ, "PIPELINE_STAGES", 4)
+    physical_stages = _env_int(environ, "PIPELINE_PHYSICAL_STAGES", stages)
+    seq_len = _env_int(environ, "PIPELINE_SEQ_LEN", 4096)
+    hidden_dim = _env_int(environ, "PIPELINE_HIDDEN_DIM", 2560)
+    warmup_steps = _env_int(environ, "PIPELINE_WARMUP_STEPS", 1)
+    profile_start_step = _env_int(environ, "PIPELINE_PROFILE_START_STEP", warmup_steps)
+    profile_steps = _env_int(environ, "PIPELINE_PROFILE_STEPS", 0)
+    profile_run_id = environ.get("PIPELINE_PROFILE_RUN_ID")
+    steps = _env_int(environ, "PIPELINE_STEPS", 4)
+    if profile_steps > 0:
+        if not profile_run_id:
+            raise ValueError("PIPELINE_PROFILE_RUN_ID is required when profiling")
+        if profile_start_step + profile_steps > steps:
+            raise ValueError("pipeline profile window must fit within PIPELINE_STEPS")
+    layer_counts_value = environ.get("PIPELINE_LAYERS_PER_STAGE")
+    layer_counts = None if not layer_counts_value else tuple(int(value) for value in layer_counts_value.split(","))
+    memory_threshold_value = environ.get("PIPELINE_RESHARD_THRESHOLD_BYTES")
+    memory_threshold = None if not memory_threshold_value else int(memory_threshold_value)
+    if hidden_dim % _ATTENTION_HEAD_DIM != 0:
+        raise ValueError(f"PIPELINE_HIDDEN_DIM must be divisible by {_ATTENTION_HEAD_DIM}, got {hidden_dim}")
+    num_heads = hidden_dim // _ATTENTION_HEAD_DIM
+    num_kv_heads = _default_num_kv_heads(num_heads)
+    return GrugPipelineTrainConfig(
+        stages=stages,
+        physical_stages=physical_stages,
+        microbatches=_env_int(environ, "PIPELINE_MICROBATCHES", 8),
+        batch_size=_env_int(environ, "PIPELINE_BATCH", 256),
+        seq_len=seq_len,
+        hidden_dim=hidden_dim,
+        intermediate_dim=_env_int(environ, "PIPELINE_INTERMEDIATE_DIM", math.ceil(hidden_dim / 256) * 128),
+        shared_expert_intermediate_dim=_env_int(
+            environ,
+            "PIPELINE_SHARED_EXPERT_INTERMEDIATE_DIM",
+            hidden_dim,
+        ),
+        num_layers=_env_int(environ, "PIPELINE_LAYERS", 24),
+        num_experts=_env_int(environ, "PIPELINE_EXPERTS", 256),
+        top_k=_env_int(environ, "PIPELINE_TOP_K", 4),
+        expert_axis_size=_env_int(environ, "PIPELINE_EXPERT_AXIS", 8),
+        vocab_size=_env_int(environ, "PIPELINE_VOCAB_SIZE", 128_256),
+        num_heads=_env_int(environ, "PIPELINE_HEADS", num_heads),
+        num_kv_heads=_env_int(environ, "PIPELINE_KV_HEADS", num_kv_heads),
+        sliding_window=_env_int(environ, "PIPELINE_SLIDING_WINDOW", 2048),
+        qk_mult=_env_float(environ, "PIPELINE_QK_MULT", 1.3),
+        steps=steps,
+        warmup_steps=warmup_steps,
+        profile_start_step=profile_start_step,
+        profile_steps=profile_steps,
+        profile_run_id=profile_run_id,
+        layer_counts=layer_counts,
+        memory_threshold=memory_threshold,
+        mp_policy_string=environ.get("PIPELINE_MP", "params=bfloat16,compute=bfloat16,output=bfloat16"),
+        remat_mode=environ.get("PIPELINE_REMAT", "recompute_all"),
+        attention_implementation=environ.get("PIPELINE_ATTENTION", "gpu_fa4_cute"),
+        moe_implementation=environ.get("PIPELINE_MOE", "ring"),
+        schedule=PipelineSchedule(environ.get("PIPELINE_SCHEDULE", PipelineSchedule.ZERO_BUBBLE)),
+def main() -> None:
+    config = _resolve_benchmark_config(os.environ)
+    _run_grug_local(config)
+if __name__ == "__main__":
+    main()

pipeline.py

added

+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+"""Automatic pipeline-parallel training for the Grug MoE pipeline variant."""
+from __future__ import annotations
+import dataclasses
+import itertools
+from collections.abc import Callable
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import TypeGuard
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jmp
+import numpy as np
+import optax
+from haliax.jax_utils import named_call
+from jax.sharding import AxisType, Mesh, NamedSharding
+from jax.sharding import PartitionSpec as P
+from jax.tree_util import register_dataclass
+from jaxtyping import Array, Float, Int
+from levanter.data.text.examples import GrugLmExample
+from levanter.grug.attention import AttentionMask
+from levanter.grug.grug_moe import MOE_REMAT_SAVE_NAMES
+from levanter.grug.loss import fused_linear_softmax_cross_entropy_loss
+from levanter.pipeline import evenly_partition_layers
+from experiments.grug.moe_pipeline.model import (
+    BATCH_AXES,
+    Block,
+    GatedNorm,
+    GrugModelConfig,
+    LayerAttentionMode,
+    RMSNorm,
+    Transformer,
+try:
+    import jaxpp.api as jaxpp
+except ModuleNotFoundError as error:
+    if error.name != "jaxpp":
+        raise
+    jaxpp = None
+if jaxpp is None:
+    mpmd = None
+else:
+    from jaxpp.experimental import mpmd
+TRAIN_LOSS_KEY = "train/loss"
+_QB_BETA_PER_LAYER_KEY = "qb_beta_per_layer"
+_PIPELINE_AXIS = "pipeline"
+type _ArrayValue = jax.Array | jax.ShapeDtypeStruct | jaxpp.MpmdArray
+class AutomaticPipelineSchedule(StrEnum):
+    ZERO_BUBBLE = "zero_bubble"
+    DUALPIPE_V = "dualpipe_v"
+@dataclass(frozen=True)
+class GrugMoePipelineConfig:
+    stages: int
+    microbatches: int
+    physical_stages: int | None = None
+    def __post_init__(self) -> None:
+        if self.stages < 2:
+            raise ValueError(f"pipeline parallelism requires at least 2 stages, got {self.stages}")
+        if self.microbatches <= 0:
+            raise ValueError(f"microbatches must be positive, got {self.microbatches}")
+        if self.physical_stages is not None:
+            if self.physical_stages < 2:
+                raise ValueError(f"pipeline parallelism requires at least 2 physical stages, got {self.physical_stages}"
+>
+            if self.stages != 2 * self.physical_stages:
+                raise ValueError(
+                    "virtual pipeline parallelism requires exactly two logical stages per physical stage; "
+                    f"got {self.stages} logical and {self.physical_stages} physical stages"
+    @property
+    def mpmd_stages(self) -> int:
+        return self.stages if self.physical_stages is None else self.physical_stages
+def make_pipeline_mesh(
+    config: GrugMoePipelineConfig,
+    *,
+    expert_axis_size: int,
+    replica_axis_size: int | None = None,
+):
+    """Build the concrete Grug mesh and wrap it as a JaxPP MPMD mesh."""
+    pp, _ = _jaxpp_modules()
+    if replica_axis_size is None:
+        replica_axis_size = max(1, jax.process_count() // config.mpmd_stages)
+    fixed_axes = config.mpmd_stages * replica_axis_size * expert_axis_size
+    if jax.device_count() % fixed_axes != 0:
+        raise ValueError(
+            f"device count {jax.device_count()} must be divisible by stages ({config.mpmd_stages}) * "
+            f"replicas ({replica_axis_size}) * experts ({expert_axis_size})"
+    data_axis_size = jax.device_count() // fixed_axes
+    shape = (config.mpmd_stages, replica_axis_size, data_axis_size, expert_axis_size, 1)
+    axis_names = (_PIPELINE_AXIS, "replica_dcn", "data", "expert", "model")
+    devices = np.asarray(jax.devices(), dtype=object).reshape(shape)
+    mesh = Mesh(devices, axis_names, axis_types=(AxisType.Explicit,) * len(axis_names))
+    if mesh.is_multi_process:
+        local_stages = {int(np.argwhere(devices == device)[0][0]) for device in jax.local_devices()}
+        if len(local_stages) != 1:
+            raise ValueError(f"each JAX process must own exactly one pipeline stage; got {sorted(local_stages)}")
+    return mesh, pp.MpmdMesh(mesh, _PIPELINE_AXIS)
+class GrugMoePipelineStage(eqx.Module):
+    """The parameters and layer range owned by one pipeline stage."""
+    token_embed: jax.Array | None
+    embed_norm: RMSNorm | None
+    embed_gated_norm: GatedNorm | None
+    output_proj: jax.Array | None
+    blocks: tuple[Block, ...]
+    final_norm: RMSNorm | None
+    final_gated_norm: GatedNorm | None
+    config: GrugModelConfig = eqx.field(static=True)
+    start_layer: int = eqx.field(static=True)
+    end_layer: int = eqx.field(static=True)
+    @named_call
+    def embed(self, token_ids: Int[Array, "B S"]) -> Float[Array, "B S D"]:
+        if self.token_embed is None or self.embed_norm is None or self.embed_gated_norm is None:
+            raise ValueError("only stage 0 owns the token embedding")
+        hidden = self.token_embed.at[token_ids].get(out_sharding=P(BATCH_AXES))
+        return self.embed_gated_norm(self.embed_norm(hidden))
+    @named_call
+    def run_blocks(
+        self,
+        hidden: Float[Array, "B S D"],
+        mask: AttentionMask | jax.Array | None,
+    ) -> tuple[Float[Array, "B S D"], dict[str, jax.Array]]:
+        if mask is None:
+            mask = AttentionMask.causal()
+        cfg = self.config
+        segment_ids = mask.segment_ids if isinstance(mask, AttentionMask) else None
+        short_mask = AttentionMask(is_causal=True, sliding_window=cfg.sliding_window, segment_ids=segment_ids)
+        long_mask = AttentionMask(is_causal=True, sliding_window=None, segment_ids=segment_ids)
+        remat_policy = (
+            jax.checkpoint_policies.save_only_these_names(*MOE_REMAT_SAVE_NAMES)
+            if cfg.remat_mode == "save_moe"
+            else None
+        block_metrics = []
+        for local_index, block in enumerate(self.blocks):
+            layer_index = self.start_layer + local_index
+            is_long = layer_index % 4 == 3 or layer_index == cfg.num_layers - 1
+            layer_mask = long_mask if is_long else short_mask
+            hidden, metrics = eqx.filter_checkpoint(block, policy=remat_policy)(
+                hidden,
+                layer_mask,
+                LayerAttentionMode.LONG if is_long else LayerAttentionMode.SHORT,
+            block_metrics.append(metrics)
+        return hidden, _stack_router_metrics(block_metrics)
+    @named_call
+    def finish(self, hidden: Float[Array, "B S D"]) -> Float[Array, "B S D"]:
+        if self.final_norm is None or self.final_gated_norm is None:
+            raise ValueError("only the final stage owns the final norms")
+        return self.final_gated_norm(self.final_norm(hidden))
+    @named_call
+    def cross_entropy_loss(
+        self,
+        hidden: Float[Array, "B S D"],
+        token_ids: Int[Array, "B S"],
+        loss_weight: Float[Array, "B S"],
+        *,
+        logsumexp_weight: float | None,
+        reduction: str = "mean",
+    ) -> jax.Array:
+        if self.output_proj is None:
+            raise ValueError("only the final stage owns the output projection")
+        labels = jnp.pad(token_ids[:, 1:], ((0, 0), (0, 1))).astype(jnp.int32)
+        return fused_linear_softmax_cross_entropy_loss(
+            hidden,
+            self.output_proj,
+            labels,
+            weight=loss_weight.astype(jnp.float32),
+            reduction=reduction,
+            logsumexp_weight=logsumexp_weight,
+            dtype=jnp.float32,
+    def local_router_loss(self, router_metrics: dict[str, jax.Array]) -> jax.Array:
+        coefficient = jnp.asarray(self.config.router_z_loss_coef / self.config.num_layers, dtype=jnp.float32)
+        return coefficient * jnp.sum(router_metrics["router_z_loss_per_layer"])
+@register_dataclass
+@dataclass(frozen=True)
+class GrugMoeAutomaticPipelineState:
+    """Array state for JaxPP's automatic pipeline transform."""
+    step: jax.Array
+    trainable_params: tuple[GrugMoePipelineStage, ...]
+    opt_state: tuple[optax.OptState, ...]
+    pending_qb_betas: tuple[jax.Array, ...]
+@dataclass(frozen=True)
+class _InitializedMpmdStageState:
+    step: jax.Array
+    opt_state: tuple[optax.OptState, ...]
+    pending_qb_betas: tuple[jax.Array, ...]
+@dataclass(frozen=True)
+class PreparedAutomaticMpmdStep:
+    step: Callable[..., tuple[GrugMoeAutomaticPipelineState, dict[str, jax.Array]]]
+    state: GrugMoeAutomaticPipelineState
+    batches: GrugLmExample
+    loss_denominator: jax.Array
+def split_transformer(
+    model: Transformer,
+    num_stages: int,
+    *,
+    layer_counts: tuple[int, ...] | None = None,
+) -> tuple[GrugMoePipelineStage, ...]:
+    """Split a Grug MoE transformer into contiguous stage pytrees."""
+    if layer_counts is None:
+        ranges = evenly_partition_layers(len(model.blocks), num_stages)
+    else:
+        if len(layer_counts) != num_stages:
+            raise ValueError(f"expected {num_stages} layer counts, got {len(layer_counts)}")
+        if any(count <= 0 for count in layer_counts):
+            raise ValueError(f"layer counts must be positive, got {layer_counts}")
+        if sum(layer_counts) != len(model.blocks):
+            raise ValueError(f"layer counts sum to {sum(layer_counts)}, but model has {len(model.blocks)} layers")
+        boundaries = np.cumsum((0, *layer_counts))
+        ranges = tuple((int(start), int(end)) for start, end in itertools.pairwise(boundaries))
+    stages = []
+    for stage_index, (start_layer, end_layer) in enumerate(ranges):
+        is_first = stage_index == 0
+        is_last = stage_index == num_stages - 1
+        stages.append(
+            GrugMoePipelineStage(
+                token_embed=model.token_embed if is_first else None,
+                embed_norm=model.embed_norm if is_first else None,
+                embed_gated_norm=model.embed_gated_norm if is_first else None,
+                output_proj=model.output_proj if is_last else None,
+                blocks=model.blocks[start_layer:end_layer],
+                final_norm=model.final_norm if is_last else None,
+                final_gated_norm=model.final_gated_norm if is_last else None,
+                config=model.config,
+                start_layer=start_layer,
+                end_layer=end_layer,
+    return tuple(stages)
+def split_automatic_stages(
+    model: Transformer,
+    *,
+    num_stages: int,
+    layer_counts: tuple[int, ...] | None = None,
+) -> tuple[tuple[GrugMoePipelineStage, ...], tuple[GrugMoePipelineStage, ...]]:
+    """Partition automatic-schedule stages into trainable and static pytrees.
+    Returns:
+        The trainable and static stage tuples. Router biases remain static
+        because the pending QB update supplies them separately.
+    """
+    stages = split_transformer(model, num_stages, layer_counts=layer_counts)
+    trainable_stages = []
+    static_stages = []
+    for stage in stages:
+        trainable_stage, static_stage = eqx.partition(stage, eqx.is_array)
+        for block_index in range(len(stage.blocks)):
+            trainable_stage = eqx.tree_at(
+                lambda current, index=block_index: current.blocks[index].mlp.router_bias,
+                trainable_stage,
+                None,
+        trainable_stages.append(trainable_stage)
+        static_stages.append(static_stage)
+    return tuple(trainable_stages), tuple(static_stages)
+def _process_has_sharding(sharding: NamedSharding) -> bool:
+    process_index = jax.process_index()
+    return any(device.process_index == process_index for device in sharding.mesh.devices.flat)
+def _empty_sharded_array(shape: tuple[int, ...], dtype, sharding: NamedSharding) -> jax.Array:
+    return jax.make_array_from_single_device_arrays(shape, sharding, [], dtype=dtype)
+def _stage_local_scalar(value: jax.Array, sharding: NamedSharding) -> jax.Array:
+    if not _process_has_sharding(sharding):
+        return _empty_sharded_array((), value.dtype, sharding)
+    return jax.device_put(np.asarray(value), sharding)
+def _localize_optimizer_scalars(mpmd_mesh, stage_index: int, opt_state):
+    stage_mesh = mpmd_mesh.unstack[stage_index]
+    def localize(value):
+        if _is_array(value) and value.shape == ():
+            return _stage_local_scalar(value, NamedSharding(stage_mesh, P()))
+        return value
+    return jax.tree.map(localize, opt_state)
+def _initialize_mpmd_stage_state(
+    stages: tuple[GrugMoePipelineStage, ...],
+    optimizer: optax.GradientTransformation,
+    mpmd_mesh,
+    stage_to_mpmd_index: tuple[int, ...],
+    memory_threshold: int | None,
+) -> _InitializedMpmdStageState:
+    pp, _ = _jaxpp_modules()
+    qb_targets = tuple(
+        pp.MpmdSharding(mpmd_mesh, mesh_ids={mpmd_index}, spec=P(None, None)) for mpmd_index in stage_to_mpmd_index
+    pending_qb_betas = pp.spmd_to_mpmd_reshard(
+        mpmd_mesh,
+        tuple(jnp.zeros((len(stage.blocks), stage.config.num_experts), dtype=jnp.float32) for stage in stages),
+        qb_targets,
+        threshold=memory_threshold,
+    opt_state = tuple(
+        _localize_optimizer_scalars(mpmd_mesh, mpmd_index, optimizer.init(stage))
+        for mpmd_index, stage in zip(stage_to_mpmd_index, stages, strict=True)
+    step = _stage_local_scalar(jnp.array(0, dtype=jnp.int32), NamedSharding(mpmd_mesh.unstack[0], P()))
+    return _InitializedMpmdStageState(
+        step=step,
+        opt_state=opt_state,
+        pending_qb_betas=pending_qb_betas,
+def initialize_mpmd_automatic_pipeline_state(
+    model: Transformer,
+    optimizer: optax.GradientTransformation,
+    mpmd_mesh,
+    *,
+    num_stages: int,
+    layer_counts: tuple[int, ...] | None = None,
+    stage_to_mpmd_index: tuple[int, ...] | None = None,
+    memory_threshold: int | None = None,
+) -> tuple[GrugMoeAutomaticPipelineState, tuple[GrugMoePipelineStage, ...]]:
+    """Return automatic state and static stages using the requested placement."""
+    pp, _ = _jaxpp_modules()
+    if stage_to_mpmd_index is None:
+        stage_to_mpmd_index = tuple(range(num_stages))
+    if len(stage_to_mpmd_index) != num_stages:
+        raise ValueError(f"expected {num_stages} stage placements, got {len(stage_to_mpmd_index)}")
+    trainable_stages, static_stages = split_automatic_stages(
+        model,
+        num_stages=num_stages,
+        layer_counts=layer_counts,
+    trainable_targets = tuple(
+        _mpmd_sharding_tree(mpmd_mesh, mpmd_index, stage)
+        for mpmd_index, stage in zip(stage_to_mpmd_index, trainable_stages, strict=True)
+    trainable_stages = pp.spmd_to_mpmd_reshard(
+        mpmd_mesh,
+        trainable_stages,
+        trainable_targets,
+        threshold=memory_threshold,
+    initialized = _initialize_mpmd_stage_state(
+        trainable_stages,
+        optimizer,
+        mpmd_mesh,
+        stage_to_mpmd_index,
+        memory_threshold,
+    return (
+        GrugMoeAutomaticPipelineState(
+            step=initialized.step,
+            trainable_params=trainable_stages,
+            opt_state=initialized.opt_state,
+            pending_qb_betas=initialized.pending_qb_betas,
+        ),
+        static_stages,
+def _stack_router_metrics(block_metrics: list[dict[str, jax.Array]]) -> dict[str, jax.Array]:
+    keys = (
+        "routing_entropy",
+        "routing_counts",
+        "load_balancing_loss",
+        "router_z_loss",
+        "qb_beta",
+        "capacity_overflow",
+    return {f"{key}_per_layer": jnp.stack([metrics[key] for metrics in block_metrics]) for key in keys}
+def _apply_qb_betas(stage: GrugMoePipelineStage, qb_betas: jax.Array) -> GrugMoePipelineStage:
+    blocks = list(stage.blocks)
+    for index, block in enumerate(blocks):
+        new_bias = -qb_betas[index]
+        new_bias = new_bias - jnp.mean(new_bias)
+        new_mlp = eqx.tree_at(lambda mlp: mlp.router_bias, block.mlp, new_bias)
+        blocks[index] = eqx.tree_at(lambda current: current.mlp, block, new_mlp)
+    return eqx.tree_at(lambda current: current.blocks, stage, tuple(blocks))
+def _jaxpp_modules():
+    if jaxpp is None or mpmd is None:
+        raise ModuleNotFoundError("The canonical Grug pipeline requires `uv sync --extra pipeline`.")
+    return jaxpp, mpmd
+def automatic_stage_to_mpmd_indices(
+    config: GrugMoePipelineConfig,
+    schedule_name: AutomaticPipelineSchedule,
+) -> tuple[int, ...]:
+    """Return the physical MPMD rank that owns each logical automatic stage."""
+    schedule = _automatic_schedule(config, schedule_name)
+    return tuple(int(schedule.get_mpmd_idx(stage_index)) for stage_index in range(config.stages))
+def _automatic_schedule(config: GrugMoePipelineConfig, schedule_name: AutomaticPipelineSchedule):
+    pp, _ = _jaxpp_modules()
+    if schedule_name == AutomaticPipelineSchedule.ZERO_BUBBLE:
+        return pp.ZeroBubble(num_stages=config.stages)
+    if schedule_name == AutomaticPipelineSchedule.DUALPIPE_V:
+        return pp.DualPipeV(num_stages=config.stages, mpmd_dim=config.mpmd_stages)
+    raise ValueError(f"unknown automatic pipeline schedule: {schedule_name}")
+def _is_array(value: object) -> TypeGuard[_ArrayValue]:
+    if isinstance(value, (jax.Array, jax.ShapeDtypeStruct)):
+        return True
+    return jaxpp is not None and isinstance(value, jaxpp.MpmdArray)
+def _partition_spec_tree(tree):
+    def partition_spec(value):
+        if not _is_array(value):
+            return None
+        if isinstance(value.sharding, NamedSharding):
+            return value.sharding.spec
+        if jaxpp is not None and isinstance(value.sharding, jaxpp.MpmdSharding):
+            return value.sharding.spec
+        return P(*([None] * value.ndim))
+    return jax.tree.map(partition_spec, tree)
+def _mpmd_sharding_tree(mpmd_mesh, stage_index: int, tree):
+    pp, _ = _jaxpp_modules()
+    def sharding(value):
+        if not _is_array(value):
+            return None
+        spec = value.sharding.spec if isinstance(value.sharding, NamedSharding) else P(*([None] * value.ndim))
+        return pp.MpmdSharding(mpmd_mesh, mesh_ids={stage_index}, spec=spec)
+    return jax.tree.map(sharding, tree)
+def make_automatic_pipeline_step(
+    optimizer: optax.GradientTransformation,
+    mp_policy: jmp.Policy,
+    static_stages: tuple[GrugMoePipelineStage, ...],
+    sample_state: GrugMoeAutomaticPipelineState,
+    sample_batches: GrugLmExample,
+    *,
+    config: GrugMoePipelineConfig,
+    mpmd_mesh,
+    schedule_name: AutomaticPipelineSchedule = AutomaticPipelineSchedule.ZERO_BUBBLE,
+    logsumexp_weight: float | None = None,
+):
+    """Build a JaxPP automatic ZeroBubble or DualPipeV optimizer step.
+    ``static_stages`` and ``sample_state`` must come from the same automatic
+    state initializer. ``sample_batches`` must have the leading microbatch axis
+    produced by :func:`levanter.pipeline.reshape_batch_into_microbatches`.
+    """
+    pp, _ = _jaxpp_modules()
+    schedule = _automatic_schedule(config, schedule_name)
+    def pipeline_step(
+        state: GrugMoeAutomaticPipelineState,
+        batches: GrugLmExample,
+        loss_denominator: jax.Array,
+    ):
+        def loss_fn(trainable_stages: tuple[GrugMoePipelineStage, ...], batch: GrugLmExample):
+            hidden = None
+            router_loss = jnp.array(0.0, dtype=jnp.float32)
+            next_qb_betas = []
+            last_stage = None
+            for stage_index, (trainable_stage, static_stage) in enumerate(
+                zip(trainable_stages, static_stages, strict=True)
+            ):
+                stage = eqx.combine(trainable_stage, static_stage)
+                stage = mp_policy.cast_to_compute(_apply_qb_betas(stage, state.pending_qb_betas[stage_index]))
+                if stage_index == 0:
+                    hidden = stage.embed(batch.tokens)
+                assert hidden is not None
+                hidden, router_metrics = stage.run_blocks(hidden, batch.attn_mask)
+                router_loss = router_loss + stage.local_router_loss(router_metrics) / config.microbatches
+                next_qb_betas.append(router_metrics[_QB_BETA_PER_LAYER_KEY])
+                if stage_index < config.stages - 1:
+                    hidden = pp.mark_stage_end(hidden)
+                last_stage = stage
+            assert last_stage is not None
+            hidden = last_stage.finish(hidden)
+            cross_entropy_sum = last_stage.cross_entropy_loss(
+                hidden,
+                batch.tokens,
+                batch.loss_weight,
+                logsumexp_weight=logsumexp_weight,
+                reduction="sum",
+            loss = cross_entropy_sum / loss_denominator + router_loss
+            loss = pp.mark_stage_end(loss)
+            return loss, tuple(next_qb_betas)
+        (loss, next_qb_betas), grads = pp.treduce(
+            lambda batch: jax.value_and_grad(loss_fn, has_aux=True)(state.trainable_params, batch),
+            batches,
+            schedule=schedule,
+            operation=((pp.Add, tuple(pp.Add for _ in range(config.stages))), pp.Add),
+        next_params = []
+        next_opt_state = []
+        for params, opt_state, stage_grads in zip(
+            state.trainable_params,
+            state.opt_state,
+            grads,
+            strict=True,
+        ):
+            updates, stage_opt_state = optimizer.update(stage_grads, opt_state, params)
+            next_params.append(eqx.apply_updates(params, updates))
+            next_opt_state.append(stage_opt_state)
+        next_state = dataclasses.replace(
+            state,
+            step=state.step + jnp.array(1, dtype=state.step.dtype),
+            trainable_params=tuple(next_params),
+            opt_state=tuple(next_opt_state),
+            pending_qb_betas=tuple(beta / config.microbatches for beta in next_qb_betas),
+        return next_state, {TRAIN_LOSS_KEY: loss}
+    return pp.mpmd_jit_with_loop(
+        pipeline_step,
+        mpmd_mesh=mpmd_mesh,
+        in_specs=(_partition_spec_tree(sample_state), _partition_spec_tree(sample_batches), P()),
+        out_specs=(_partition_spec_tree(sample_state), {TRAIN_LOSS_KEY: P()}),
+def prepare_automatic_mpmd_step(
+    step,
+    state: GrugMoeAutomaticPipelineState,
+    batches: GrugLmExample,
+    loss_denominator: jax.Array,
+    mpmd_mesh,
+    *,
+    memory_threshold: int | None = None,
+) -> PreparedAutomaticMpmdStep:
+    """Compile with stage-local state and place only the still-SPMD batch inputs."""
+    pp, _ = _jaxpp_modules()
+    compiled = step.compile(state, batches, loss_denominator)
+    args_shardings, kwargs_shardings = compiled.in_shardings
+    if kwargs_shardings:
+        raise ValueError("automatic pipeline step does not accept keyword arguments")
+    def place_initial_scalar(value, target):
+        if not _is_array(value) or isinstance(value, pp.MpmdArray):
+            return value
+        if value.shape != ():
+            return value
+        mesh_ids = target.mesh_ids or {0}
+        local_arrays = []
+        for stage_index in sorted(mesh_ids):
+            sharding = NamedSharding(mpmd_mesh.unstack[stage_index], target.spec)
+            if _process_has_sharding(sharding):
+                local_arrays.append(jax.device_put(np.zeros((), dtype=value.dtype), sharding))
+        return pp.MpmdArray(
+            local_arrays,
+            target,
+            shape=(),
+            dtype=value.dtype,
+    state = jax.tree.map(place_initial_scalar, state, args_shardings[0])
+    if memory_threshold is None and all(device.memory_stats() is None for device in jax.local_devices()):
+        memory_threshold = 0
+    batches, loss_denominator = pp.spmd_to_mpmd_reshard(
+        mpmd_mesh,
+        (batches, loss_denominator),
+        (args_shardings[1], args_shardings[2]),
+        threshold=memory_threshold,
+    return PreparedAutomaticMpmdStep(
+        step=compiled,
+        state=state,
+        batches=batches,
+        loss_denominator=loss_denominator,

adamh.py

removed

	experiments/grug/moe/adamh.py
t	1	# Copyright The Marin Authors	t
	2	# SPDX-License-Identifier: Apache-2.0
	3	#
	4	# Local copy of AdamH for iteration without modifying Levanter.
	5	# Adapted from levanter.optim.adamh.
	6
	7	from typing import Any, NamedTuple
	8
	9	import chex
	10	import jax
	11	import jax.numpy as jnp
	12	import optax
	13	from optax import tree_utils as otu
	14
	15
	16	class ScaleByAdamHState(NamedTuple):
	17	count: chex.Array
	18	mu: optax.Updates
	19	nu: optax.Updates
	20
	21
	22	def scale_by_adamh(
	23	b1: float = 0.9,
	24	b2: float = 0.999,
	25	eps: float = 1e-8,
	26	learning_rate: float = 0.02,
	27	mu_dtype: Any \| None = None,
	28	) -> optax.GradientTransformation:
	29	mu_dtype = jax.dtypes.canonicalize_dtype(mu_dtype)
	30
	31	def init_fn(params):
	32	mu = otu.tree_zeros_like(params, dtype=mu_dtype)
	33	nu = otu.tree_zeros_like(params)
	34	return ScaleByAdamHState(count=jnp.zeros([], jnp.int32), mu=mu, nu=nu)
	35
	36	def update_fn(updates, state, params):
	37	mu = otu.tree_update_moment(updates, state.mu, b1, 1)
	38	nu = otu.tree_update_moment_per_elem_norm(updates, state.nu, b2, 2)
	39	count_inc = optax.safe_increment(state.count)
	40	mu_hat = otu.tree_bias_correction(mu, b1, count_inc)
	41	nu_hat = otu.tree_bias_correction(nu, b2, count_inc)
	42
	43	adam_updates = jax.tree.map(
	44	lambda m, v: None if m is None else m / (jnp.sqrt(v) + eps),
	45	mu_hat,
	46	nu_hat,
	47	is_leaf=lambda x: x is None,
	48	)
	49	mu = otu.tree_cast(mu, mu_dtype)
	50
	51	def _scale_invariant_2d(p, u):
	52	"""Core update for a 2-D (matrix) parameter."""
	53	p_norm = jnp.linalg.norm(p)
	54	u_norm = jnp.linalg.norm(u)
	55	new_p = p - learning_rate * u * p_norm / jnp.maximum(u_norm, 1e-10)
	56	return new_p / jnp.linalg.norm(new_p) * p_norm - p
	57
	58	def scale_invariant_update(p, u):
	59	if p is None:
	60	return None
	61	if p.ndim <= 2:
	62	return _scale_invariant_2d(p, u)
	63	# For higher-rank tensors, vmap the 2-D logic over the leading axis.
	64	return jax.vmap(_scale_invariant_2d)(p, u)
	65
	66	adamh_updates = jax.tree_util.tree_map(
	67	scale_invariant_update,
	68	params,
	69	adam_updates,
	70	is_leaf=lambda x: x is None,
	71	)
	72
	73	return adamh_updates, ScaleByAdamHState(count=count_inc, mu=mu, nu=nu)
	74
	75	return optax.GradientTransformation(init_fn, update_fn)
	76
	77
	78	__all__ = ["ScaleByAdamHState", "scale_by_adamh"]

agent.md

removed

	experiments/grug/moe/agent.md
t	1	# Agent Guide: experiments/grug/moe	t
	2
	3	## Autonomy
	4
	5	This workflow is designed to run end-to-end without human confirmation. The
	6	agent is authorized to:
	7
	8	- Create branches, commit, and push without asking
	9	- Create GitHub experiment issues and post comments
	10	- Submit Iris jobs and kill only jobs submitted by self
	11	- Run experiments through both gates autonomously
	12
	13	Do not stop to ask for confirmation at any step. If something fails, diagnose
	14	and retry or report the failure — do not block waiting for input.
	15
	16	## Objective
	17
	18	Determine whether a proposed change outperforms the baseline. Baseline results
	19	are in `experiments/grug/moe/README.md` — compare against the table there.
	20
	21	Metrics (from wandb):
	22	- `eval/paloma/macro_loss` (final value)
	23	- `throughput/tokens_per_second` (averaged over the last 100 steps)
	24	- `throughput/total_tokens` (final value)
	25	- `run.state` must be `finished` before pulling final metrics
	26
	27	Baseline scaling law (May Recipe, drop-1e18 fit, L∞ pinned at 1.6):
	28
	29	```
	30	loss(C) = 1.6 + 88.32 · C^(-0.0941)
	31	```
	32
	33	### Gate 1: effective speedup at two small scales
	34
	35	Run the variant at `d512` (3.82e17 FLOPs) and `d768` (2.81e18 FLOPs).
	36
	37	For each scale, compute the effective speedup at a fixed macro_loss target
	38	(use the baseline's final macro_loss at that scale as the target). The variant
	39	passes gate 1 if it shows an effective speedup at both scales.
	40
	41	### Gate 2: scaling law projection
	42
	43	Run the variant at the two larger scales: `d1024` (1.16e19) and `d1280`
	44	(3.46e19). Combine with the gate 1 results (d512, d768) for four total points.
	45
	46	The variant passes gate 2 if:
	47	1. It shows an effective speedup at all four scales.
	48	2. Fit a new scaling law `loss(C) = 1.6 + A · C^(-alpha)` (asymptote pinned
	49	at 1.6) on the variant's four optima. Project to 1e21 and 1e23 FLOPs.
	50	The variant's projected loss must be lower than the baseline's at both
	51	budgets (baseline: 2.534 at 1e21, 2.205 at 1e23).
	52
	53	### Effective speedup calculation
	54
	55	The question: how much faster (in wall-clock time) does the variant reach
	56	its final loss compared to how long the baseline would need to reach that
	57	same loss?
	58
	59	The scaling law `loss(C) = 1.6 + A · C^(-0.0941)` has a fixed exponent but
	60	the coefficient A varies. We recenter A to pass through the baseline's
	61	actual (budget, loss) point, then invert to find how much compute the
	62	baseline would need to match the variant's loss. Finally, compare wall-clock
	63	times.
	64
	65	```python
	66	def effective_speedup(baseline_loss, baseline_tps, variant_loss, variant_tps, budget):
	67	"""Wall-clock speedup of the variant over the baseline.
	68
	69	Returns > 1 if the variant reaches its loss faster in real time.
	70	"""
	71	L_inf = 1.6
	72	alpha = 0.0941
	73
	74	# 1. Fit the scaling law coefficient through the baseline's data point.
	75	# loss = L_inf + A_bl * C^(-alpha) => A_bl = (loss - L_inf) * C^alpha
	76	A_bl = (baseline_loss - L_inf) * budget ** alpha
	77
	78	# 2. How much compute would the baseline need to reach variant's loss?
	79	C_needed = (A_bl / (variant_loss - L_inf)) ** (1 / alpha)
	80
	81	# 3. Compare wall-clock times:
	82	# baseline_time = C_needed / baseline_tps
	83	# variant_time = budget / variant_tps
	84	return (C_needed / baseline_tps) / (budget / variant_tps)
	85	```
	86
	87	### Example
	88
	89	Suppose at d512 / 3.82e17 FLOPs (May Recipe d=512 compute-optimal point):
	90	- Baseline: macro_loss = 3.5422, tok/s = 433,986 (v4-32 EP=1 reference run)
	91	- Variant: macro_loss = 3.52, tok/s = 410,000 (better loss, ~5% slower)
	92
	93	```python
	94	# 1. Fit A through baseline point
	95	A_bl = (3.5422 - 1.6) * (3.82e17) ** 0.0941 # ≈ 87.7
	96
	97	# 2. Compute baseline needs to reach 3.52
	98	C_needed = (A_bl / (3.52 - 1.6)) ** (1 / 0.0941) # > 3.82e17
	99
	100	# 3. Wall-clock comparison
	101	baseline_time = C_needed / 433_986
	102	variant_time = 3.82e17 / 410_000
	103	speedup = baseline_time / variant_time
	104	```
	105
	106	If the quality improvement (needing more baseline compute to match) outweighs
	107	the throughput cost (variant taking longer per FLOP), speedup > 1.
	108
	109	## Implementation
	110
	111	Most promotable changes will land in one of three files:
	112
	113	- `model.py` — architecture tweaks (routing, norms, attention, activation functions, expert layout, etc.).
	114	- `heuristic.py` — scaling heuristics (LR formula coefficients, depth/width formula, GQA ratio, per-batch-size epsilon/b
	>	eta2 scaling).
	115	- `optimizer.py` — optimizer internals (AdamH components, parameter-group partitioning, per-group learning rates, weight
	>	decay).
	116
	117	## Documentation & GitHub Issues
	118
	119	Create a new branch for each experiment issue. Branch off `main`.
	120
	121	Follow `.agents/skills/run-research/SKILL.md` for all documentation, logbooks,
	122	W&B tracking, and GitHub experiment issue management tied to work in this
	123	directory. Pay attention to this file carefully.
	124
	125	Experiment issues should be titled `Agent MoE Experiment: [description]`.
	126	Include the exact prompt from the user that initiated the experiment in the
	127	issue body.
	128
	129	After creating the issue, add it as a sub-issue of #4281 (April 2026 MoE
	130	scaling tracking issue) using the GitHub GraphQL API. This is required — do not skip it. First get the node IDs, then
	131	call `addSubIssue`:
	132
	133	```bash
	134	# 1. Get node IDs for the parent and the new issue
	135	gh api graphql -f query='
	136	query {
	137	repository(owner: "marin-community", name: "marin") {
	138	parent: issue(number: 4281) { id }
	139	child: issue(number: <NEW_ISSUE_NUMBER>) { id }
	140	}
	141	}'
	142
	143	# 2. Add the sub-issue relationship
	144	gh api graphql -f query='
	145	mutation {
	146	addSubIssue(input: {issueId: "<PARENT_ID>", subIssueId: "<CHILD_ID>"}) {
	147	issue { number }
	148	subIssue { number }
	149	}
	150	}'
	151	```
	152
	153	## Authentication
	154
	155	Assume the user has already completed these before job submission:
	156	- `WANDB_API_KEY` set in the environment
	157	- `gcloud auth login` and `gcloud auth application-default login`
	158
	159	## Job Submission
	160
	161	Jobs in this directory are submitted to Iris on a v5p-8.
	162
	163	### Submission command
	164
	165	```bash
	166	.venv/bin/iris --cluster=marin job run \
	167	--no-wait \
	168	--reserve v5p-8 \
	169	-e WANDB_API_KEY "$WANDB_API_KEY" \
	170	-- python -m experiments.grug.moe.launch --version dev --run
	171	```
	172
	173	Swap the module path (`experiments.grug.moe.launch`) for whichever launch
	174	script in this directory you are running. Every launch script takes the shared
	175	experiment CLI (`marin.experiment.cli`): `--version` sets the checkpoint version
	176	(required; `dev` to iterate, a calendar `YYYY.MM.DD` to keep), and `--run` builds
	177	the trial — without it the lowered plan is printed and nothing runs.
	178
	179	`--reserve v5p-8` is a hard zone constraint (confines the job to a zone where a
	180	v5p-8 has actually been obtained, waiting otherwise), not a capacity reservation
	181	— it does not hold or guarantee capacity.
	182
	183	### Monitoring
	184
	185	Runs may take time to find a TPU, and 5–10 minutes to start once scheduled.
	186	After confirming the run is progressing on wandb, jobs typically take over an
	187	hour to complete. Sleep at reasonable intervals (e.g. 15 minutes) before
	188	checking status — do not poll in a tight loop.
	189
	190	Reconnect to logs:
	191	```bash
	192	.venv/bin/iris --cluster=marin job logs -f JOB_ID
	193	```
	194
	195	List your jobs:
	196	```bash
	197	.venv/bin/iris --cluster=marin job list \| grep "$(whoami)"
	198	```
	199
	200	Check runs in wandb (match `<PREFIX>` to the run_id pattern in `launch.py`):
	201	```python
	202	import wandb
	203	api = wandb.Api()
	204	runs = api.runs('marin-community/marin_moe', filters={'displayName': {'$regex': '^<PREFIX>'}}, order='-created_at')
	205	for r in runs:
	206	print(f'{r.name:<50} state={r.state:<10} step={r.summary.get("global_step", "n/a")}')
	207	```

heuristic.py

removed

	experiments/grug/moe/heuristic.py
t	1	# Copyright The Marin Authors	t
	2	# SPDX-License-Identifier: Apache-2.0
	3
	4	"""Compute-scaling heuristic for MoE — May Recipe (MuonH, 256 experts).
	5
	6	Refit on the May Recipe LR sweep (issue #5951: 17 cells across
	7	d ∈ {512, 768, 1024, 1280}, MuonH optimizer, R²=0.996), and is the
	8	current default for compute-optimal cells and ablation comparisons.
	9	All empirical fits were measured on runs with seq_len=4096; the
	10	formulas are written in terms of ``tokens_per_batch = batch_size *
	11	seq_len`` so they generalise to other sequence lengths, though the
	12	coefficients are an extrapolation beyond 4096.
	13
	14	For the earlier AdamH-tuned heuristic (64 experts), see the pre-rename
	15	``experiments/grug/moe/heuristic.py`` on main:
	16	https://github.com/marin-community/marin/blob/8586719b524bf7743ec5034403c7e834505fe73e/experiments/grug/moe/heuristic.py
	17
	18	Formulas:
	19	- Adam LR: ``adam_lr = lr_coeff * tokens^lr_tokens_exp * dim^lr_dim_exp * sqrt(B)``
	20	(``lr_coeff = 0.06602``, ``lr_tokens_exp = -0.395``, ``lr_dim_exp = -0.150``)
	21	- MuonH/AdamH LR: ``lr = (13/3) * adam_lr``
	22	- Compute budget convention: ``C = 3 * flops_per_token(no_lm_head) * tokens``
	23	- Epsilon: ``epsilon = epsilon_base * sqrt(r0/r)`` where ``r = (BT0)/(B0T)``
	24	- Beta1: fixed at ``0.9062``
	25	- Beta2: ``beta2 = clip(beta2_base^(B/B0), min_beta2, max_beta2)``
	26	"""
	27
	28	import math
	29	from dataclasses import dataclass
	30
	31	from levanter.utils.flop_utils import lm_flops_per_token
	32
	33	from experiments.grug.moe.model import GrugModelConfig
	34	from experiments.grug.moe.optimizer import GrugMoeMuonHConfig
	35
	36	SEQ_LEN: int = 8192
	37	MIN_BATCH_SIZE: int = 32
	38	DEFAULT_TARGET_STEPS: int = 2**14
	39
	40
	41	def _round_to_power_of_two(x: float) -> int:
	42	if x <= 1:
	43	return 1
	44	return 2 ** math.ceil(math.log2(x))
	45
	46
	47	def compute_flops_per_token(cfg: GrugModelConfig) -> float:
	48	"""Non-embedding FLOPs per token (excludes lm_head)."""
	49	fpt_with_lm_head = lm_flops_per_token(
	50	hidden_dim=cfg.hidden_dim,
	51	intermediate_dim=cfg.intermediate_dim,
	52	num_layers=cfg.num_layers,
	53	num_kv_heads=cfg.num_kv_heads,
	54	num_heads=cfg.num_heads,
	55	seq_len=cfg.max_seq_len,
	56	vocab_size=cfg.vocab_size,
	57	glu=True,
	58	num_experts=cfg.num_experts,
	59	num_shared_experts=1 if cfg.shared_expert_intermediate_dim > 0 else 0,
	60	num_experts_per_tok=cfg.num_experts_per_token,
	61	shared_intermediate_dim=cfg.shared_expert_intermediate_dim,
	62	)
	63	return fpt_with_lm_head - 2 * cfg.hidden_dim * cfg.vocab_size
	64
	65
	66	def compute_tokens_and_batch(
	67	budget: float,
	68	flops_per_token: float,
	69	target_steps: int = DEFAULT_TARGET_STEPS,
	70	min_batch_size: int = MIN_BATCH_SIZE,
	71	seq_len: int = SEQ_LEN,
	72	) -> tuple[float, int, int]:
	73	"""Derive (tokens, batch_size, num_steps) from a compute budget and FLOPs-per-token.
	74
	75	``seq_len`` controls the sequence length used to convert between batch_size
	76	(sequences) and tokens_per_batch (tokens per step). All downstream formulas
	77	that depend on batch size use ``tokens_per_batch = batch_size * seq_len``
	78	so they work correctly at any sequence length.
	79	"""
	80	tokens = budget / (3 * flops_per_token)
	81	batch_exact = tokens / (target_steps * seq_len)
	82	batch_size = max(min_batch_size, _round_to_power_of_two(batch_exact))
	83	train_steps = max(1, round(tokens / (batch_size * seq_len)))
	84	return tokens, batch_size, train_steps
	85
	86
	87	@dataclass(frozen=True)
	88	class MoeHeuristic:
	89	"""Compute-scaling heuristic for MoE models (May Recipe refit).
	90
	91	Returns a ``MuonH`` optimizer config via ``build_optimizer_config``.
	92
	93	``adam_lr = lr_coeff * tokens^lr_tokens_exp * dim^lr_dim_exp * sqrt(tokens_per_batch)``
	94	``muonh_lr = muonh_ratio * adam_lr``
	95	``C = 3 * flops_per_token * tokens`` (flops_per_token excludes lm_head)
	96	"""
	97
	98	# --- LR scaling ---
	99	# May Recipe refit (issue #5951; 17 cells across d{512,768,1024,1280},
	100	# R^2=0.996): muonh_lr = 18.31 * tokens^-0.395 * dim^-0.150 * sqrt(B).
	101	lr_coeff: float = 0.06602
	102	lr_tokens_exp: float = -0.395
	103	lr_dim_exp: float = -0.150
	104	muonh_ratio: float = 13 / 3
	105
	106	# --- Base hyperparameters ---
	107	epsilon_coeff: float = 9.676e-18
	108	beta1: float = 0.9062
	109	beta2_base: float = 0.999
	110	beta2_reference_tpb: int = 131_072 # beta2 = beta2_base^(tpb / beta2_reference_tpb)
	111
	112	# --- Fixed hyperparameters ---
	113	max_grad_norm: float = 1.0
	114
	115	# --- Schedule ---
	116	min_lr_ratio: float = 0.05
	117	lr_schedule: str = "linear"
	118	decay: float \| None = None
	119
	120	# --- Architecture ---
	121	vocab_size: int = 128_256
	122	hidden_head_ratio: int = 128
	123	gqa_ratio: int \| None = 4 # None = MHA, 4 = 4:1 GQA, etc.
	124	base_hidden_layer_ratio: int = 64
	125	layer_scaling_factor: float = 4.0
	126	layer_formula_offset: int = 9
	127
	128	# --- Constraints ---
	129	max_learning_rate: float = 0.05
	130	min_beta2: float = 0.95
	131	max_beta2: float = 0.9999
	132
	133	def _compute_adam_lr(self, tokens_per_batch: int, tokens: float, hidden_dim: int) -> float:
	134	"""adam_lr = lr_coeff * tokens^lr_tokens_exp * dim^lr_dim_exp * sqrt(tokens_per_batch)"""
	135	adam_lr = (
	136	self.lr_coeff * (tokens*self.lr_tokens_exp) (hidden_dim*self.lr_dim_exp) math.sqrt(tokens_per_batch)
	137	)
	138	return min(self.max_learning_rate, adam_lr)
	139
	140	def _compute_learning_rate(self, tokens_per_batch: int, tokens: float, hidden_dim: int) -> float:
	141	"""muonh_lr = muonh_ratio * adam_lr"""
	142	adam_lr = self._compute_adam_lr(tokens_per_batch, tokens, hidden_dim)
	143	return min(self.max_learning_rate, self.muonh_ratio * adam_lr)
	144
	145	def _compute_epsilon(self, tokens_per_batch: int, tokens: float) -> float:
	146	"""epsilon = epsilon_coeff * sqrt(tokens / tokens_per_batch)"""
	147	return self.epsilon_coeff * math.sqrt(tokens / tokens_per_batch)
	148
	149	def _compute_beta2(self, tokens_per_batch: int) -> float:
	150	"""beta2 = clip(beta2_0^(tpb/tpb0), min_beta2, max_beta2). Constant token half-life."""
	151	exponent = tokens_per_batch / self.beta2_reference_tpb
	152	return max(self.min_beta2, min(self.max_beta2, self.beta2_base**exponent))
	153
	154	def build_optimizer_config(
	155	self, batch_size: int, tokens: float, hidden_dim: int, seq_len: int = SEQ_LEN
	156	) -> GrugMoeMuonHConfig:
	157	"""Return a ``GrugMoeMuonHConfig`` with May Recipe 1pct-noclip defaults.
	158
	159	MuonH optimizer (3 LR groups — muonh / adamh / adam) with
	160	``warmup=0.01`` and ``max_grad_norm=None``. LR / beta / epsilon
	161	scaling is the May Recipe refit (see module docstring).
	162	"""
	163	tokens_per_batch = batch_size * seq_len
	164	lr = self._compute_learning_rate(tokens_per_batch, tokens, hidden_dim)
	165	adam_lr = self._compute_adam_lr(tokens_per_batch, tokens, hidden_dim)
	166	epsilon = self._compute_epsilon(tokens_per_batch, tokens)
	167	beta2 = self._compute_beta2(tokens_per_batch)
	168	return GrugMoeMuonHConfig(
	169	learning_rate=lr,
	170	adam_lr=adam_lr,
	171	min_lr_ratio=self.min_lr_ratio,
	172	warmup=0.01, # 1pct-noclip schedule
	173	beta1=self.beta1,
	174	beta2=beta2,
	175	epsilon=epsilon,
	176	max_grad_norm=None, # no clipping
	177	lr_schedule=self.lr_schedule,
	178	decay=self.decay,
	179	)
	180
	181	def _compute_num_layers(self, hidden_size: int) -> int:
	182	hs_pow = math.log2(hidden_size)
	183	return round(
	184	hidden_size
	185	/ (self.base_hidden_layer_ratio + (hs_pow * self.layer_scaling_factor) - self.layer_formula_offset)
	186	)
	187
	188	def _get_step_size(self, budget: float) -> int:
	189	if budget > self.budget_step_threshold:
	190	return self.large_budget_step_size
	191	return self.small_budget_step_size
	192
	193	def _max_params_for_budget(self, budget: float) -> float:
	194	scaling = self.base_max_params * math.sqrt(budget / self.base_max_params_budget)
	195	return min(max(self.base_max_params, scaling), self.global_max_params)
	196
	197	@staticmethod
	198	def _compute_kv_heads(num_heads: int, gqa_ratio: int \| None) -> int:
	199	"""Compute num_kv_heads for a given GQA ratio.
	200
	201	If gqa_ratio is None, returns num_heads (MHA).
	202	Otherwise returns the largest divisor of num_heads <= num_heads // gqa_ratio.
	203	"""
	204	if gqa_ratio is None:
	205	return num_heads
	206	target = num_heads // gqa_ratio
	207	for k in range(target, 0, -1):
	208	if num_heads % k == 0:
	209	return k
	210	return 1
	211
	212	def build_model_config(self, hidden_size: int, seq_len: int = SEQ_LEN) -> GrugModelConfig:
	213	if hidden_size % self.hidden_head_ratio != 0:
	214	raise ValueError(
	215	f"hidden_size ({hidden_size}) must be divisible by hidden_head_ratio ({self.hidden_head_ratio})."
	216	)
	217	num_layers = self._compute_num_layers(hidden_size)
	218	num_heads = max(1, hidden_size // self.hidden_head_ratio)
	219	num_kv_heads = self._compute_kv_heads(num_heads, self.gqa_ratio)
	220
	221	return GrugModelConfig(
	222	vocab_size=self.vocab_size,
	223	hidden_dim=hidden_size,
	224	# Round up to nearest 128 for Pallas TPU MoE kernel compatibility
	225	intermediate_dim=math.ceil(hidden_size / 2 / 128) * 128,
	226	shared_expert_intermediate_dim=hidden_size,
	227	num_layers=num_layers,
	228	num_heads=num_heads,
	229	num_kv_heads=num_kv_heads,
	230	max_seq_len=seq_len,
	231	sliding_window=2048,
	232	initializer_std=0.5 / math.sqrt(hidden_size),
	233	qk_mult=1.3,
	234	)
	235
	236
	237	def build_from_heuristic(
	238	*,
	239	budget: float,
	240	hidden_dim: int,
	241	heuristic: MoeHeuristic \| None = None,
	242	target_steps: int = DEFAULT_TARGET_STEPS,
	243	min_batch_size: int = MIN_BATCH_SIZE,
	244	seq_len: int = SEQ_LEN,
	245	) -> tuple[GrugModelConfig, GrugMoeMuonHConfig, int, int]:
	246	"""Construct (model, optimizer, batch_size, num_steps) for a compute budget.
	247
	248	Uses ``MoeHeuristic`` to size the model (from ``hidden_dim``) and to set
	249	the MuonH hyperparameters (scaled by tokens_per_batch = batch_size * seq_len).
	250	Callers who want manual control should continue passing ``GrugModelConfig`` /
	251	``GrugMoeMuonHConfig`` directly to ``GrugMoeLaunchConfig``.
	252	"""
	253	h = heuristic or MoeHeuristic()
	254	model_cfg = h.build_model_config(hidden_dim, seq_len=seq_len)
	255	fpt = compute_flops_per_token(model_cfg)
	256	tokens, batch_size, num_steps = compute_tokens_and_batch(
	257	budget,
	258	fpt,
	259	target_steps=target_steps,
	260	min_batch_size=min_batch_size,
	261	seq_len=seq_len,
	262	)
	263	optimizer_cfg = h.build_optimizer_config(batch_size, tokens, hidden_dim, seq_len=seq_len)
	264	return model_cfg, optimizer_cfg, batch_size, num_steps

launch.py

removed

	experiments/grug/moe/launch.py
t	1	# Copyright The Marin Authors	t
	2	# SPDX-License-Identifier: Apache-2.0
	3
	4	"""grug-moe trial run, authored as a lazy artifact.
	5
	6	The run is a function that returns a typed :class:`Checkpoint` handle addressed by an
	7	explicit ``name@version``. The model, optimizer, data mixture, token budget, z-loss,
	8	and evals are all stated inline; the output path is ``ctx.output_path`` and the TPU is a
	9	run-arg, so neither bears on the artifact's identity.
	10
	11	The grug-moe training mechanism (``GrugMoeLaunchConfig`` + ``run_grug_moe_trial`` ->
	12	``run_grug``) is grug-specific compute and is kept as-is: the recipe's ``fn`` is
	13	``run_grug_moe_trial``, which builds the Levanter trainer and dispatches the training
	14	job to Fray. Only the data/validation wiring around it is assembled lazily from
	15	dataset handles.
	16	"""
	17
	18	import dataclasses
	19	import os
	20	from dataclasses import dataclass, field
	21	from datetime import timedelta
	22
	23	import jmp
	24	from fray.cluster import ResourceConfig
	25	from levanter.callbacks.profiler import ProfilerConfig
	26	from levanter.checkpoint import CheckpointerConfig, latest_checkpoint_path
	27	from levanter.data.text.datasets import LmDataConfig
	28	from levanter.optim.config import OptimizerConfig
	29	from levanter.tracker import TrackerConfig
	30	from levanter.tracker.wandb import WandbConfig
	31	from levanter.trainer import TrainerConfig
	32	from marin.execution.build_context import resolve_version
	33	from marin.execution.lazy import ArtifactStep, StepContext
	34	from marin.experiment.cli import experiment_main
	35	from marin.experiment.data import mixture, tokenized
	36	from marin.experiment.namespacing import user_namespaced_name
	37	from marin.processing.tokenize.tokenize import TokenizedCache
	38	from marin.training.training import LevanterCheckpoint, resolve_checkpointer_output_path
	39
	40	from experiments.datasets.nemotron import nemotron_datasets
	41	from experiments.datasets.paloma import paloma_datasets
	42	from experiments.datasets.proofpile import proofpile_dataset
	43	from experiments.datasets.starcoder import starcoder_dataset
	44	from experiments.datasets.uncheatable import uncheatable_datasets
	45	from experiments.grug.moe.heuristic import build_from_heuristic
	46	from experiments.grug.moe.model import GrugModelConfig
	47	from experiments.grug.moe.train import GrugEvalConfig, GrugRunConfig, GrugTrainerConfig, run_grug
	48	from experiments.llama import llama3_tokenizer
	49
	50	# SlimPajama-6B tokenization OOMs at the default 10g worker resources.
	51	_SLIMPAJAMA_TOKENIZE_RESOURCES = ResourceConfig(ram="64g", disk="64g")
	52
	53	# The TPU the training job is dispatched onto. A run-arg, not part of the config's
	54	# identity: re-running on a different TPU is the same checkpoint. The launcher step
	55	# runs inline (run_grug dispatches its own Fray job).
	56	_TRAIN_RESOURCES = ResourceConfig.with_tpu("v5p-8")
	57
	58	# Nemotron CC mixture weights: the corpus's TiB proportions, plus starcoder and
	59	# proof-pile at their published weights. Policy lives here, in the experiment.
	60	_NEMOTRON_WEIGHTS = {
	61	"hq_actual": 0.91351,
	62	"hq_synth": 2.72,
	63	"medium_high": 0.82471,
	64	"medium": 3.38,
	65	"medium_low": 1.54,
	66	"low_actual": 0.70123,
	67	"low_synth": 0.62771,
	68	}
	69	_STARCODER_WEIGHT = 0.25
	70	_PROOFPILE_WEIGHT = 0.055
	71
	72
	73	@dataclass(frozen=True)
	74	class GrugMoeLaunchConfig:
	75	"""Last-mile run config for the MoE grug template.
	76
	77	Keep this as the main entry point for day-to-day edits (model/data/optimizer/trainer/eval knobs).
	78	"""
	79
	80	model: GrugModelConfig
	81	data: LmDataConfig
	82	output_path: str
	83	run_id: str
	84	resources: ResourceConfig
	85	steps: int
	86	batch_size: int
	87	seed: int
	88	mp: str # jmp policy string, e.g. "params=float32,compute=bfloat16,output=bfloat16".
	89	tracker: TrackerConfig
	90	optimizer: OptimizerConfig
	91	profiler: ProfilerConfig = field(default_factory=ProfilerConfig)
	92	grug_trainer: GrugTrainerConfig = field(default_factory=GrugTrainerConfig)
	93	eval: GrugEvalConfig \| None = field(default_factory=GrugEvalConfig)
	94	processes_per_task: int = 1
	95	"""GPU processes per task. > 1 fans each node into one JAX process per GPU
	96	(multi-controller) via the iris.hooks.multigpu_main supervisor; 1 keeps the
	97	single-process-per-node model."""
	98	checkpointer: CheckpointerConfig \| None = None
	99	"""Override the checkpointer. None builds the default (periodic + final saves
	100	under output_path). Throughput experiments point this at node-local disk so a
	101	slow object-store commit can't wedge the end-of-run barrier."""
	102	init_from: str \| None = None
	103	"""Checkpoint base directory to initialize weights from (the latest checkpoint
	104	under it is loaded). None trains from scratch. Used to chain training phases —
	105	a midtrain/SFT/RL run points this at the prior phase's ``checkpoints`` directory."""
	106
	107
	108	def env_int(key: str, default: int) -> int:
	109	"""Read an int from ``os.environ[key]``, falling back to ``default`` when unset/empty."""
	110	raw = os.environ.get(key, "")
	111	return int(raw) if raw else default
	112
	113
	114	def slimpajama_6b_dataset() -> ArtifactStep[TokenizedCache]:
	115	"""SlimPajama-6B, llama3-tokenized — a small corpus for GPU smoke/scale runs.
	116
	117	Returns the tokenized :class:`TokenizedCache` handle; the launcher assembles it into an
	118	``LmDataConfig`` with :func:`~marin.experiment.data.mixture`. Tokenization runs as
	119	its own Fray job (a production pretraining mixture would instead pin an already
	120	materialized cache to avoid a cross-region tokenize).
	121	"""
	122	return tokenized(
	123	"slimpajama-6b",
	124	source="DKYoon/SlimPajama-6B",
	125	tokenizer=llama3_tokenizer,
	126	resources=_SLIMPAJAMA_TOKENIZE_RESOURCES,
	127	version="2026.06.28",
	128	)
	129
	130
	131	def _resolve_run_id(default_run_id: str) -> str:
	132	"""Resolve run id and append `FERRY_DATE` when launching from ferry workflows."""
	133	run_id = os.environ.get("GRUG_RUN_ID", default_run_id)
	134	ferry_date = os.environ.get("FERRY_DATE")
	135	if ferry_date:
	136	run_id = f"{run_id}-{ferry_date}"
	137	return run_id
	138
	139
	140	def _resolve_tracker(tracker: TrackerConfig, run_id: str) -> TrackerConfig:
	141	if isinstance(tracker, WandbConfig):
	142	return dataclasses.replace(tracker, name=run_id)
	143	return tracker
	144
	145
	146	def run_grug_moe_trial(config: GrugMoeLaunchConfig) -> None:
	147	"""Map template launch knobs onto a full Levanter trainer and dispatch the run.
	148
	149	Runs inline on the launcher; ``run_grug`` submits the training job to Fray and
	150	blocks until it completes.
	151	"""
	152	initialize_from = latest_checkpoint_path(config.init_from) if config.init_from is not None else None
	153	trainer = TrainerConfig(
	154	id=config.run_id,
	155	seed=config.seed,
	156	train_batch_size=config.batch_size,
	157	num_train_steps=config.steps,
	158	profiler=config.profiler,
	159	mp=jmp.get_policy(config.mp),
	160	tracker=_resolve_tracker(config.tracker, config.run_id),
	161	use_explicit_mesh_axes=True,
	162	require_accelerator=True,
	163	allow_nondivisible_batch_size=False,
	164	initialize_from=initialize_from,
	165	checkpointer=config.checkpointer
	166	or resolve_checkpointer_output_path(
	167	CheckpointerConfig(save_interval=timedelta(minutes=10), keep=None),
	168	config.output_path,
	169	),
	170	)
	171
	172	grug_trainer = dataclasses.replace(config.grug_trainer, trainer=trainer)
	173
	174	run_config = GrugRunConfig(
	175	model=config.model,
	176	data=config.data,
	177	resources=config.resources,
	178	optimizer=config.optimizer,
	179	trainer=grug_trainer,
	180	eval=config.eval,
	181	processes_per_task=config.processes_per_task,
	182	)
	183	run_grug(run_config)
	184
	185
	186	RESOLVED_RUN_ID = _resolve_run_id("4_10_test_moe")
	187
	188
	189	# Baseline: 1e18 compute budget, d1024. Model + optimizer + batch + steps are
	190	# all derived from `MoeAdamHHeuristic`. To override any of these, swap in
	191	# an explicit `GrugModelConfig` / `GrugMoeAdamHConfig` below.
	192	_BASELINE_BUDGET: float = 1e18
	193	_BASELINE_HIDDEN_DIM: int = 1024
	194	_BASELINE_TARGET_STEPS: int = 2**14
	195	_baseline_model, _baseline_optimizer, _baseline_batch, _baseline_steps = build_from_heuristic(
	196	budget=_BASELINE_BUDGET,
	197	hidden_dim=_BASELINE_HIDDEN_DIM,
	198	target_steps=_BASELINE_TARGET_STEPS,
	199	)
	200
	201
	202	def grug_moe_baseline(*, version: str \| None = None) -> ArtifactStep[LevanterCheckpoint]:
	203	"""The baseline grug MoE (QB+GN+XSA+zloss) on the Nemotron mix as a lazy checkpoint.
	204
	205	Every component is a :class:`Dataset` handle, so the whole graph lowers via
	206	:func:`~marin.execution.lazy.lower`. Pinned components never re-tokenize; the
	207	paloma/uncheatable suites are validation (weight 0).
	208	"""
	209	name = "grug/4_10_baseline_moe"
	210	version = resolve_version(name, version)
	211	nem = nemotron_datasets(tokenizer=llama3_tokenizer)
	212	train = {nem[split]: weight for split, weight in _NEMOTRON_WEIGHTS.items()}
	213	train[starcoder_dataset(tokenizer=llama3_tokenizer)] = _STARCODER_WEIGHT
	214	train[proofpile_dataset(tokenizer=llama3_tokenizer)] = _PROOFPILE_WEIGHT
	215	validation = [
	216	*paloma_datasets(tokenizer=llama3_tokenizer).values(),
	217	*uncheatable_datasets(tokenizer=llama3_tokenizer).values(),
	218	]
	219
	220	def build_config(ctx: StepContext) -> GrugMoeLaunchConfig:
	221	return GrugMoeLaunchConfig(
	222	model=_baseline_model,
	223	data=mixture(ctx, train, validation=validation),
	224	output_path=ctx.output_path,
	225	run_id=RESOLVED_RUN_ID,
	226	resources=ctx.runtime_arg("train_resources"),
	227	steps=_baseline_steps,
	228	batch_size=_baseline_batch,
	229	seed=0,
	230	mp="params=float32,compute=bfloat16,output=bfloat16",
	231	tracker=WandbConfig(project="marin_moe", tags=["moe"], group="moe-iter04", name=None),
	232	optimizer=_baseline_optimizer,
	233	grug_trainer=GrugTrainerConfig(z_loss_weight=1e-4, ema_beta=None, log_every=1),
	234	eval=GrugEvalConfig(
	235	eval_batch_size=512,
	236	steps_per_eval=1000,
	237	max_eval_batches=8,
	238	eval_current=True,
	239	eval_ema=False,
	240	),
	241	)
	242
	243	return ArtifactStep(
	244	name=user_namespaced_name(name, version),
	245	version=version,
	246	artifact_type=LevanterCheckpoint,
	247	run=run_grug_moe_trial,
	248	build_config=build_config,
	249	deps=(train, validation),
	250	runtime_args={"train_resources": _TRAIN_RESOURCES},
	251	)
	252
	253
	254	if __name__ == "__main__":
	255	experiment_main(grug_moe_baseline)()

launch_cw_scale.py

removed

	experiments/grug/moe/launch_cw_scale.py
t	1	# Copyright The Marin Authors	t
	2	# SPDX-License-Identifier: Apache-2.0
	3
	4	"""Large sparse-MoE scale run for the CoreWeave cw-us-east-02a H100 cluster.
	5
	6	Launches a ~90B-total / ~5B-active Grug MoE (hidden 3072, 48 layers, 128 experts,
	7	top-4 -> ~17x sparsity) across all 32 nodes / 256 H100s. Parameters are fully
	8	sharded over the cross-node ``data`` axis (FSDP) while the 128 routed experts are
	9	sharded 8-way over the intra-node NVLink ``expert`` axis (expert parallelism).
	10
	11	This is the size class the cluster can train as a single model. The canary
	12	(``experiments/ferries/canary_ferry.py``) replicates parameters per node, so it
	13	caps at the ~9.5B that fits on one node's 8 GPUs; here ``replica_axis_size=1``
	14	shards one model across every device instead.
	15
	16	Env knobs (all optional; defaults give the full 90B run on 256 H100):
	17
	18	SCALE_GPU_REPLICAS number of 8xH100 nodes (default 32 -> 256 GPUs)
	19	SCALE_EXPERT_AXIS expert-parallel axis size, intra-node (default 8)
	20	SCALE_REPLICA_AXIS cross-node replication; 1 = pure FSDP (default 1)
	21	SCALE_PROCESSES_PER_TASK GPU processes per node: 1 = one process per node
	22	(default), 8 = one JAX process per GPU (multi-controller)
	23	SCALE_BATCH global batch in sequences (default 256)
	24	SCALE_SEQ_LEN sequence length (default 2048)
	25	SCALE_STEPS training steps (default 50)
	26	SCALE_HIDDEN_DIM / SCALE_NUM_LAYERS / SCALE_NUM_EXPERTS / SCALE_TOP_K
	27	model-shape overrides (e.g. a smaller FSDP smoke test)
	28	SCALE_REMAT recompute_all (default) \| save_moe -- save_moe keeps the
	29	tagged MoE dispatch tensors for backward so the EP
	30	collectives are not re-run during recompute
	31	SCALE_MP jmp policy (default params=float32,compute=bfloat16,
	32	output=bfloat16); params=bfloat16 halves FSDP gather bytes
	33	SCALE_TRACKER wandb \| json_logger (default json_logger)
	34	SCALE_PROFILER_STEPS >0 enables a jax_profile capture window of N steps
	35	(use SCALE_TRACKER=wandb so the artifact uploads)
	36	SCALE_PROFILER_START profiler start step (default 8, past compile/warmup)
	37	SCALE_CHECKPOINTS s3 (default) \| local. local writes checkpoints to
	38	node-local disk with no periodic saves -- for throughput
	39	experiments where the checkpoint is disposable and a
	40	slow S3 commit must not wedge the end-of-run barrier
	41	RUN_ID unique run identifier
	42	"""
	43
	44	import datetime
	45	import os
	46	from typing import cast
	47
	48	from fray.cluster import ResourceConfig
	49	from levanter.callbacks.profiler import ProfilerConfig
	50	from levanter.checkpoint import CheckpointerConfig
	51	from levanter.data.text.datasets import BlockShuffleConfig
	52	from levanter.optim.config import AdamConfig
	53	from levanter.tracker.json_logger import JsonLoggerConfig
	54	from levanter.tracker.wandb import WandbConfig
	55	from marin.execution.build_context import resolve_version
	56	from marin.execution.lazy import ArtifactStep, StepContext
	57	from marin.experiment.cli import experiment_main
	58	from marin.experiment.data import mixture
	59	from marin.experiment.namespacing import user_namespaced_name
	60	from marin.training.training import LevanterCheckpoint
	61
	62	from experiments.grug.moe.launch import GrugMoeLaunchConfig, env_int, run_grug_moe_trial, slimpajama_6b_dataset
	63	from experiments.grug.moe.model import GrugModelConfig, RematMode
	64	from experiments.grug.moe.train import GrugTrainerConfig
	65	from experiments.llama import llama3_tokenizer_vocab_size
	66
	67	# head_dim is fixed at 128; hidden_dim must be a multiple of it.
	68	HEAD_DIM = 128
	69	VOCAB_SIZE = llama3_tokenizer_vocab_size
	70	GPUS_PER_NODE = 8 # H100s per gd-8xh100ib node; the batch-shard math and with_gpu(count=...) must track
	71	# Default seq for the 90B run. FSDP reshards the [batch, seq, hidden] activation
	72	# through a fully-replicated intermediate (an XLA SPMD limitation, pending Shardy),
	73	# so peak memory scales with batch*seq; at the default 89.7B model, batch 256 x
	74	# seq 2048 fits in 80GB while 512 x 4096 OOMs (~58GiB replicated tile).
	75	DEFAULT_SEQ_LEN = 2048
	76	DEFAULT_BATCH = 256
	77
	78	# Modest, schedule-stable Adam for a short scale/throughput run (not trained to
	79	# convergence). expert weights share the schedule; the goal is to exercise the
	80	# FSDP+expert-parallel mesh at scale, not to produce a checkpoint.
	81	SCALE_OPTIMIZER = AdamConfig(
	82	learning_rate=6e-4,
	83	weight_decay=0.1,
	84	lr_schedule="cosine",
	85	warmup=10,
	86	min_lr_ratio=0.1,
	87	)
	88
	89	SCALE_TRAINER_DEFAULTS = dict(z_loss_weight=1e-4, ema_beta=None, log_every=1)
	90
	91	# Subdirectory of MARIN_PREFIX these scale runs write their per-run output dirs
	92	# into, so they stay grouped instead of cluttering the prefix root.
	93	OUTPUT_SUBDIR = "experiments/grug-moe-cw"
	94
	95	# SlimPajama block-shuffle: a small, R2-local corpus for the scale/throughput run.
	96	_SLIMPAJAMA_SHUFFLE = BlockShuffleConfig(io_block_size=256, window_blocks=256, perm_type="feistel")
	97
	98
	99	def build_scale_model() -> GrugModelConfig:
	100	"""~90B-total / ~5B-active sparse MoE (overridable via SCALE_* env vars)."""
	101	hidden_dim = env_int("SCALE_HIDDEN_DIM", 3072)
	102	if hidden_dim % HEAD_DIM != 0:
	103	raise ValueError(f"SCALE_HIDDEN_DIM={hidden_dim} must be a multiple of head_dim={HEAD_DIM}")
	104	num_heads = hidden_dim // HEAD_DIM
	105	# ~4:1 grouped-query attention; back off to the nearest divisor of num_heads.
	106	num_kv_heads = max(1, num_heads // 4)
	107	while num_heads % num_kv_heads != 0:
	108	num_kv_heads -= 1
	109	intermediate_dim = hidden_dim // 2 # expert FFN inner width (~d/2)
	110	seq_len = env_int("SCALE_SEQ_LEN", DEFAULT_SEQ_LEN)
	111	remat_mode = os.environ.get("SCALE_REMAT", "recompute_all")
	112	if remat_mode not in ("recompute_all", "save_moe"):
	113	raise ValueError(f"SCALE_REMAT={remat_mode!r} must be 'recompute_all' or 'save_moe'")
	114	return GrugModelConfig(
	115	vocab_size=VOCAB_SIZE,
	116	hidden_dim=hidden_dim,
	117	num_layers=env_int("SCALE_NUM_LAYERS", 48),
	118	num_heads=num_heads,
	119	num_kv_heads=num_kv_heads,
	120	head_dim=HEAD_DIM,
	121	intermediate_dim=intermediate_dim,
	122	shared_expert_intermediate_dim=intermediate_dim,
	123	num_experts=env_int("SCALE_NUM_EXPERTS", 128),
	124	num_experts_per_token=env_int("SCALE_TOP_K", 4),
	125	max_seq_len=seq_len,
	126	sliding_window=seq_len,
	127	remat_mode=cast(RematMode, remat_mode),
	128	)
	129
	130
	131	def build_scale_checkpoint(*, version: str \| None = None) -> ArtifactStep[LevanterCheckpoint]:
	132	"""Assemble the CoreWeave scale run as a lazy :class:`LevanterCheckpoint` from SCALE_* env."""
	133	run_id = os.environ.get("RUN_ID") or datetime.datetime.now(datetime.UTC).strftime("%Y%m%d-%H%M%S")
	134
	135	replicas = env_int("SCALE_GPU_REPLICAS", 32)
	136	expert_axis = env_int("SCALE_EXPERT_AXIS", 8)
	137	replica_axis = env_int("SCALE_REPLICA_AXIS", 1)
	138	batch_size = env_int("SCALE_BATCH", DEFAULT_BATCH)
	139	steps = env_int("SCALE_STEPS", 50)
	140	# 1 = one process per node (8 local GPUs). 8 = one JAX process per GPU
	141	# (multi-controller) via the iris.hooks.multigpu_main supervisor.
	142	processes_per_task = env_int("SCALE_PROCESSES_PER_TASK", 1)
	143	# SCALE_PROFILER_STEPS > 0 captures a jax_profile window of that many steps
	144	# (uploaded via the tracker, so pair with SCALE_TRACKER=wandb to retrieve it).
	145	profiler_steps = env_int("SCALE_PROFILER_STEPS", 0)
	146	profiler = ProfilerConfig(
	147	enabled=profiler_steps > 0,
	148	start_step=env_int("SCALE_PROFILER_START", 8),
	149	num_steps=profiler_steps,
	150	)
	151
	152	checkpoint_mode = os.environ.get("SCALE_CHECKPOINTS", "s3").lower()
	153	if checkpoint_mode == "local":
	154	checkpointer = CheckpointerConfig(
	155	base_path=f"/tmp/grug-scale-ckpt/{run_id}",
	156	append_run_id_to_base_path=False,
	157	save_interval=None,
	158	keep=None,
	159	)
	160	elif checkpoint_mode == "s3":
	161	checkpointer = None
	162	else:
	163	raise ValueError(f"SCALE_CHECKPOINTS={checkpoint_mode!r} must be 's3' or 'local'")
	164
	165	model = build_scale_model()
	166	if model.num_experts % expert_axis != 0:
	167	raise ValueError(f"num_experts={model.num_experts} must be divisible by SCALE_EXPERT_AXIS={expert_axis}")
	168
	169	# Batch is sharded over the (replica_dcn, data, expert) axes; data absorbs the
	170	# rest of the 8*replicas devices. Require the global batch to cover every shard.
	171	data_axis = (replicas * GPUS_PER_NODE) // (replica_axis * expert_axis)
	172	batch_shards = replica_axis * data_axis * expert_axis
	173	if batch_size % batch_shards != 0:
	174	raise ValueError(f"SCALE_BATCH={batch_size} must be divisible by batch shards={batch_shards}")
	175
	176	resources = ResourceConfig.with_gpu("H100", count=GPUS_PER_NODE, cpu=32, ram="256g", disk="256g", replicas=replicas)
	177
	178	use_wandb = os.environ.get("SCALE_TRACKER", "json_logger").lower() == "wandb"
	179	json_logger_name = os.environ.get("SCALE_JSON_LOGGER", "grug_moe_scale.metrics")
	180	wandb_entity = os.environ.get("WANDB_ENTITY") or None
	181	wandb_project = os.environ.get("WANDB_PROJECT", "marin_moe")
	182
	183	grug_trainer = GrugTrainerConfig(
	184	expert_axis_size=expert_axis,
	185	replica_axis_size=replica_axis,
	186	**SCALE_TRAINER_DEFAULTS,
	187	)
	188
	189	mp = os.environ.get("SCALE_MP", "params=float32,compute=bfloat16,output=bfloat16")
	190	name = f"grug-moe-cw-d{model.hidden_dim}-L{model.num_layers}-e{model.num_experts}-r{replicas}"
	191	slim = slimpajama_6b_dataset()
	192
	193	def build_config(ctx: StepContext) -> GrugMoeLaunchConfig:
	194	if use_wandb:
	195	tracker = WandbConfig(
	196	entity=wandb_entity,
	197	project=wandb_project,
	198	tags=["grug", "moe", "cw", "h100", "scale"],
	199	group="grug-moe-cw-scale",
	200	name=None,
	201	replicate_path=ctx.output_path,
	202	)
	203	else:
	204	tracker = JsonLoggerConfig(logger_name=json_logger_name)
	205	return GrugMoeLaunchConfig(
	206	model=model,
	207	data=mixture(ctx, {slim: 1.0}, shuffle=_SLIMPAJAMA_SHUFFLE),
	208	output_path=ctx.output_path,
	209	run_id=run_id,
	210	resources=ctx.runtime_arg("train_resources"),
	211	steps=steps,
	212	batch_size=batch_size,
	213	seed=0,
	214	mp=mp,
	215	tracker=tracker,
	216	optimizer=SCALE_OPTIMIZER,
	217	grug_trainer=grug_trainer,
	218	processes_per_task=processes_per_task,
	219	eval=None,
	220	profiler=profiler,
	221	checkpointer=checkpointer,
	222	)
	223
	224	step_name = f"{OUTPUT_SUBDIR}/{name}-{run_id}"
	225	version = resolve_version(step_name, version)
	226	return ArtifactStep(
	227	name=user_namespaced_name(step_name, version),
	228	version=version,
	229	artifact_type=LevanterCheckpoint,
	230	run=run_grug_moe_trial,
	231	build_config=build_config,
	232	deps=(slim,),
	233	runtime_args={"train_resources": resources},
	234	)
	235
	236
	237	if __name__ == "__main__":
	238	experiment_main(build_scale_checkpoint)()

launch_datakit_moe_mix.py

removed

	experiments/grug/moe/launch_datakit_moe_mix.py
t	1	# Copyright The Marin Authors	t
	2	# SPDX-License-Identifier: Apache-2.0
	3
	4	"""Grug-MoE run on the datakit store with the mixture from ``mixture-3.csv``.
	5
	6	The source CSV was provided at ``/Users/held/Downloads/mixture-3.csv`` with
	7	columns ``bucket,phase_0,phase_1,sum``. The two phase columns are independently
	8	normalized below and used as a two-stage schedule.
	9	"""
	10
	11	import math
	12
	13	from fray.cluster import ResourceConfig
	14	from levanter.data.text.datasets import ConcatDatasetComponent, DatasetComponent, LmDataConfig, UrlDatasetSourceConfig
	15	from levanter.data.text.formats import TextLmDatasetFormat
	16	from levanter.tracker.wandb import WandbConfig
	17	from marin.execution.build_context import resolve_version
	18	from marin.execution.lazy import ArtifactStep, StepContext
	19	from marin.experiment.cli import experiment_main
	20	from marin.experiment.namespacing import user_namespaced_name
	21	from marin.training.training import LevanterCheckpoint
	22
	23	from experiments.datasets.paloma import paloma_datasets
	24	from experiments.datasets.uncheatable import uncheatable_datasets
	25	from experiments.grug.moe.heuristic import build_from_heuristic
	26	from experiments.grug.moe.launch import GrugMoeLaunchConfig, run_grug_moe_trial
	27	from experiments.grug.moe.train import GrugEvalConfig, GrugTrainerConfig
	28	from experiments.marin_tokenizer import marin_tokenizer
	29
	30	_STORE_PREFIX = "datakit/store_8ac06c74"
	31	# Large enough that the smallest nonzero CSV weights still receive at least
	32	# one sequence per block, while aligning cleanly with the heuristic batch size.
	33	_MIXTURE_BLOCK_SIZE = 49_152
	34	_PHASE_1_START_FRACTION = 0.8
	35	ENABLE_SIMULATED_EPOCHING = True
	36
	37	# Natural size of ``datakit/store_8ac06c74`` from Will's datakit-moe-mix branch:
	38	# 167 mixable bucket caches plus the 33-cache tail component.
	39	_TARGET_BUDGET_TOKENS = 10_372_343_704_053
	40
	41	# The datakit store lives in us-central2; pin training there to avoid cross-region I/O.
	42	_TRAIN_RESOURCES = ResourceConfig.with_tpu("v4-8", zone="us-central2-b")
	43
	44	_BUCKET_PHASE_WEIGHTS: tuple[tuple[str, float, float], ...] = (
	45	("c01q0", 0.022957, 0.031937),
	46	("c01q1", 0.018953, 0.041846),
	47	("c01q2", 0.069309, 0.053564),
	48	("c01q3", 0.106783, 0.093756),
	49	("c01q4", 0.124669, 0.103804),
	50	("c02q0", 0.005079, 0.013405),
	51	("c02q1", 0.003787, 0.00623),
	52	("c02q2", 0.005353, 0.003842),
	53	("c02q3", 0.015501, 0.013842),
	54	("c02q4", 0.021761, 0.008706),
	55	("c03q0", 0.001046, 0.008815),
	56	("c03q1", 0.003039, 0.001349),
	57	("c03q2", 0.003699, 0.003679),
	58	("c03q3", 0.001283, 0.006298),
	59	("c03q4", 0.001561, 0.002798),
	60	("c05q0", 0.011293, 0.012139),
	61	("c05q1", 0.021915, 0.023056),
	62	("c05q2", 0.035838, 0.017236),
	63	("c05q3", 0.009497, 0.010704),
	64	("c05q4", 0.004878, 0.009904),
	65	("c06q0", 0.008112, 0.010894),
	66	("c06q1", 0.009701, 0.017409),
	67	("c06q2", 0.017897, 0.005937),
	68	("c06q3", 0.001296, 0.001905),
	69	("c06q4", 0.001886, 0.000709),
	70	("c08q0", 0.004234, 0.003704),
	71	("c08q1", 0.000809, 0.000849),
	72	("c08q2", 0.001256, 0.00157),
	73	("c08q3", 0.000469, 0.000629),
	74	("c08q4", 0.000374, 0.000279),
	75	("c10q0", 0.000086, 0.0001),
	76	("c11q0", 0.005798, 0.002363),
	77	("c11q1", 0.003333, 0.0018),
	78	("c11q2", 0.002542, 0.002297),
	79	("c11q3", 0.000779, 0.001259),
	80	("c11q4", 0.000199, 0.000672),
	81	("c12q0", 0.000311, 0.000545),
	82	("c12q1", 0.000259, 0.000419),
	83	("c12q2", 0.000151, 0.000791),
	84	("c12q3", 0.000296, 0.000881),
	85	("c12q4", 0.000247, 0.000687),
	86	("c13q0", 0.000227, 0.000526),
	87	("c13q1", 0.000364, 0.000095),
	88	("c13q2", 0.000116, 0.0005),
	89	("c13q3", 0.000472, 0.00205),
	90	("c13q4", 0.000063, 0.000157),
	91	("c14q0", 0.000633, 0.000512),
	92	("c14q1", 0.000439, 0.000516),
	93	("c14q2", 0.000976, 0.000399),
	94	("c14q3", 0.000585, 0.000462),
	95	("c14q4", 0.000443, 0.000526),
	96	("c15q0", 0.002384, 0.006218),
	97	("c15q1", 0.003226, 0.003075),
	98	("c15q2", 0.002358, 0.002513),
	99	("c15q3", 0.002819, 0.003405),
	100	("c15q4", 0.000486, 0.000773),
	101	("c16q0", 0.004551, 0.00549),
	102	("c16q1", 0.006622, 0.007018),
	103	("c16q2", 0.004254, 0.004357),
	104	("c16q3", 0.005395, 0.002788),
	105	("c16q4", 0.001036, 0.001565),
	106	("c17q0", 0.001453, 0.001401),
	107	("c17q1", 0.003893, 0.001949),
	108	("c17q2", 0.003539, 0.006945),
	109	("c17q3", 0.002154, 0.002518),
	110	("c17q4", 0.00181, 0.000973),
	111	("c18q0", 0.002814, 0.003565),
	112	("c18q1", 0.004576, 0.003527),
	113	("c18q2", 0.003187, 0.004375),
	114	("c18q3", 0.005101, 0.003618),
	115	("c18q4", 0.001469, 0.002143),
	116	("c19q0", 0.004086, 0.002186),
	117	("c19q1", 0.003805, 0.00188),
	118	("c19q2", 0.002327, 0.002511),
	119	("c19q3", 0.002111, 0.001474),
	120	("c19q4", 0.000521, 0.000995),
	121	("c20q0", 0.000992, 0.002133),
	122	("c20q1", 0.001298, 0.001715),
	123	("c20q2", 0.001601, 0.001149),
	124	("c20q3", 0.001476, 0.00099),
	125	("c20q4", 0.000374, 0.00054),
	126	("c21q0", 0.002196, 0.000878),
	127	("c21q1", 0.002377, 0.0012),
	128	("c21q2", 0.001096, 0.000948),
	129	("c21q3", 0.000976, 0.000974),
	130	("c21q4", 0.001686, 0.001902),
	131	("c22q0", 0.002234, 0.003176),
	132	("c22q1", 0.003447, 0.003771),
	133	("c22q2", 0.003587, 0.002235),
	134	("c22q3", 0.002543, 0.003216),
	135	("c22q4", 0.001364, 0.00072),
	136	("c23q0", 0.007766, 0.002833),
	137	("c23q1", 0.003635, 0.002799),
	138	("c23q2", 0.002582, 0.004731),
	139	("c23q3", 0.004336, 0.002363),
	140	("c23q4", 0.001879, 0.002431),
	141	("c24q0", 0.004557, 0.004443),
	142	("c24q1", 0.004783, 0.006302),
	143	("c24q2", 0.001682, 0.002943),
	144	("c24q3", 0.002718, 0.00169),
	145	("c24q4", 0.000578, 0.000684),
	146	("c25q0", 0.005998, 0.005436),
	147	("c25q1", 0.005659, 0.003901),
	148	("c25q2", 0.003611, 0.006288),
	149	("c25q3", 0.003449, 0.002666),
	150	("c25q4", 0.000614, 0.000481),
	151	("c26q0", 0.004179, 0.00445),
	152	("c26q1", 0.005244, 0.004102),
	153	("c26q2", 0.004386, 0.004962),
	154	("c26q3", 0.002379, 0.002656),
	155	("c26q4", 0.00082, 0.001751),
	156	("c27q0", 0.005148, 0.001263),
	157	("c27q1", 0.001803, 0.001539),
	158	("c27q2", 0.003706, 0.002862),
	159	("c27q3", 0.000881, 0.001019),
	160	("c27q4", 0.000684, 0.003425),
	161	("c28q0", 0.006398, 0.015222),
	162	("c28q1", 0.020521, 0.015838),
	163	("c28q2", 0.005379, 0.010677),
	164	("c28q3", 0.007225, 0.002797),
	165	("c28q4", 0.001426, 0.001549),
	166	("c29q0", 0.004778, 0.003881),
	167	("c29q1", 0.003144, 0.003859),
	168	("c29q2", 0.002375, 0.002314),
	169	("c29q3", 0.001843, 0.001895),
	170	("c29q4", 0.000971, 0.002561),
	171	("c30q0", 0.007903, 0.035271),
	172	("c30q1", 0.008205, 0.007997),
	173	("c30q2", 0.013895, 0.009056),
	174	("c30q3", 0.004825, 0.005009),
	175	("c30q4", 0.001323, 0.001261),
	176	("c31q0", 0.007819, 0.003531),
	177	("c31q1", 0.002745, 0.00469),
	178	("c31q2", 0.001324, 0.002866),
	179	("c31q3", 0.002475, 0.001757),
	180	("c31q4", 0.000456, 0.001024),
	181	("c32q0", 0.00486, 0.002465),
	182	("c32q1", 0.001979, 0.008457),
	183	("c32q2", 0.001013, 0.00262),
	184	("c32q3", 0.001109, 0.001447),
	185	("c32q4", 0.000106, 0.000135),
	186	("c33q0", 0.002961, 0.003771),
	187	("c33q1", 0.004153, 0.005535),
	188	("c33q2", 0.006076, 0.002388),
	189	("c33q3", 0.004801, 0.001326),
	190	("c33q4", 0.000316, 0.000683),
	191	("c34q0", 0.000514, 0.000736),
	192	("c34q1", 0.000941, 0.000924),
	193	("c34q2", 0.00867, 0.007028),
	194	("c34q3", 0.003685, 0.010664),
	195	("c34q4", 0.000834, 0.00333),
	196	("c35q0", 0.004318, 0.003901),
	197	("c35q1", 0.007961, 0.002719),
	198	("c35q2", 0.004451, 0.002239),
	199	("c35q3", 0.002711, 0.001726),
	200	("c35q4", 0.003109, 0.000526),
	201	("c36q0", 0.002889, 0.002709),
	202	("c36q1", 0.002673, 0.002149),
	203	("c36q2", 0.002718, 0.002256),
	204	("c36q3", 0.00471, 0.003361),
	205	("c36q4", 0.009675, 0.002205),
	206	("c37q0", 0.002776, 0.003737),
	207	("c37q1", 0.003821, 0.004736),
	208	("c37q2", 0.012379, 0.020839),
	209	("c37q3", 0.036626, 0.031543),
	210	("c37q4", 0.013194, 0.021405),
	211	("c38q0", 0.000022, 0.000179),
	212	("tail", 0.00007, 0.00003),
	213	)
	214
	215	_TAIL_BUCKETS: tuple[str, ...] = (
	216	"c00q0",
	217	"c00q1",
	218	"c00q2",
	219	"c00q3",
	220	"c00q4",
	221	"c04q0",
	222	"c04q1",
	223	"c04q2",
	224	"c04q3",
	225	"c04q4",
	226	"c07q0",
	227	"c07q1",
	228	"c07q2",
	229	"c07q3",
	230	"c07q4",
	231	"c09q0",
	232	"c09q1",
	233	"c09q2",
	234	"c09q3",
	235	"c09q4",
	236	"c10q1",
	237	"c10q2",
	238	"c10q3",
	239	"c10q4",
	240	"c38q1",
	241	"c38q2",
	242	"c38q3",
	243	"c38q4",
	244	"c39q0",
	245	"c39q1",
	246	"c39q2",
	247	"c39q3",
	248	"c39q4",
	249	)
	250
	251
	252	def _bucket_path(bucket: str) -> str:
	253	cluster = int(bucket[1:3])
	254	quality = int(bucket[-1])
	255	return f"{_STORE_PREFIX}/cluster={cluster}/quality={quality}"
	256
	257
	258	def _bucket_component(bucket: str) -> DatasetComponent:
	259	return DatasetComponent(
	260	source=None,
	261	cache_dir=_bucket_path(bucket),
	262	format=TextLmDatasetFormat(),
	263	tags=[bucket],
	264	flat_cache=True,
	265	)
	266
	267
	268	def _val_component(cache_dir: str) -> DatasetComponent:
	269	"""Placeholder validation component from a cache dir path."""
	270	source = UrlDatasetSourceConfig(
	271	tags=[],
	272	train_urls=[],
	273	validation_urls=[],
	274	cache_dir=cache_dir,
	275	format=TextLmDatasetFormat(),
	276	)
	277	return DatasetComponent(source=source, cache_dir=source.cache_dir, format=source.format, tags=source.tags)
	278
	279
	280	def _normalize(weights: dict[str, float]) -> dict[str, float]:
	281	total = sum(weights.values())
	282	if total <= 0:
	283	raise ValueError("mixture weights must sum to a positive value")
	284	return {name: weight / total for name, weight in weights.items() if weight > 0}
	285
	286
	287	def _phase_weights(phase_index: int) -> dict[str, float]:
	288	if phase_index == 0:
	289	return _normalize({bucket: phase_0 for bucket, phase_0, _ in _BUCKET_PHASE_WEIGHTS})
	290	if phase_index == 1:
	291	return _normalize({bucket: phase_1 for bucket, _, phase_1 in _BUCKET_PHASE_WEIGHTS})
	292	raise ValueError(f"unknown phase index {phase_index}")
	293
	294
	295	def _phase_1_start_step(total_steps: int, batch_size: int) -> int:
	296	requested = max(1, int(total_steps * _PHASE_1_START_FRACTION))
	297	step_multiple = _MIXTURE_BLOCK_SIZE // math.gcd(_MIXTURE_BLOCK_SIZE, batch_size)
	298	return max(step_multiple, (requested // step_multiple) * step_multiple)
	299
	300
	301	def _datakit_components() -> dict[str, DatasetComponent \| ConcatDatasetComponent]:
	302	direct = {bucket: _bucket_component(bucket) for bucket, _, _ in _BUCKET_PHASE_WEIGHTS if bucket != "tail"}
	303	return {
	304	**direct,
	305	"tail": ConcatDatasetComponent(
	306	children={bucket: _bucket_component(bucket) for bucket in _TAIL_BUCKETS},
	307	tags=["tail"],
	308	),
	309	}
	310
	311
	312	def _simulated_experiment_budget(*, total_steps: int, batch_size: int, max_seq_len: int) -> int:
	313	return total_steps * batch_size * max_seq_len
	314
	315
	316	def _simulated_epoching_budgets(
	317	*,
	318	total_steps: int,
	319	batch_size: int,
	320	max_seq_len: int,
	321	target_budget: int,
	322	enable_simulated_epoching: bool,
	323	) -> tuple[int \| None, int \| None]:
	324	if not enable_simulated_epoching:
	325	return None, None
	326	experiment_budget = _simulated_experiment_budget(
	327	total_steps=total_steps,
	328	batch_size=batch_size,
	329	max_seq_len=max_seq_len,
	330	)
	331	if experiment_budget > target_budget:
	332	raise ValueError(f"experiment_budget {experiment_budget} exceeds target_budget {target_budget}")
	333	return target_budget, experiment_budget
	334
	335
	336	def _two_phase_data_config(
	337	*,
	338	tokenizer: str,
	339	components: dict[str, DatasetComponent \| ConcatDatasetComponent],
	340	phase_weights: tuple[dict[str, float], dict[str, float]],
	341	phase_1_start: int,
	342	val_components: dict[str, DatasetComponent \| ConcatDatasetComponent],
	343	target_budget: int \| None,
	344	experiment_budget: int \| None,
	345	) -> LmDataConfig:
	346	val_zero_weights = {name: 0.0 for name in val_components}
	347	budget_kwargs = {}
	348	if target_budget is not None:
	349	budget_kwargs = {"target_budget": target_budget, "experiment_budget": experiment_budget}
	350
	351	return LmDataConfig(
	352	tokenizer=tokenizer,
	353	cache_dir=None,
	354	components={components, val_components},
	355	train_weights=[
	356	(0, {phase_weights[0], val_zero_weights}),
	357	(phase_1_start, {phase_weights[1], val_zero_weights}),
	358	],
	359	auto_build_caches=False,
	360	mixture_block_size=_MIXTURE_BLOCK_SIZE,
	361	**budget_kwargs,
	362	)
	363
	364
	365	def _datakit_data_config(
	366	*,
	367	total_steps: int,
	368	batch_size: int,
	369	max_seq_len: int,
	370	enable_simulated_epoching: bool,
	371	val_components: dict[str, DatasetComponent \| ConcatDatasetComponent],
	372	) -> LmDataConfig:
	373	phase_1_start = _phase_1_start_step(total_steps, batch_size)
	374	target_budget, experiment_budget = _simulated_epoching_budgets(
	375	total_steps=total_steps,
	376	batch_size=batch_size,
	377	max_seq_len=max_seq_len,
	378	target_budget=_TARGET_BUDGET_TOKENS,
	379	enable_simulated_epoching=enable_simulated_epoching,
	380	)
	381
	382	return _two_phase_data_config(
	383	tokenizer=marin_tokenizer,
	384	components=_datakit_components(),
	385	phase_weights=(_phase_weights(0), _phase_weights(1)),
	386	phase_1_start=phase_1_start,
	387	val_components=val_components,
	388	target_budget=target_budget,
	389	experiment_budget=experiment_budget,
	390	)
	391
	392
	393	_BUDGET: float = 2.19e17
	394	_HIDDEN_DIM: int = 512
	395	_TARGET_STEPS: int = 2**14
	396	_model, _optimizer, _batch_size, _steps = build_from_heuristic(
	397	budget=_BUDGET,
	398	hidden_dim=_HIDDEN_DIM,
	399	target_steps=_TARGET_STEPS,
	400	)
	401
	402	_SLUG = f"d{_HIDDEN_DIM}-{_BUDGET:.2e}"
	403
	404	_VALIDATION = [
	405	*paloma_datasets(tokenizer=marin_tokenizer).values(),
	406	*uncheatable_datasets(tokenizer=marin_tokenizer).values(),
	407	]
	408
	409
	410	def build(*, version: str \| None = None) -> ArtifactStep[LevanterCheckpoint]:
	411	"""Grug MoE on the us-central2 datakit store with the mixture-3 two-phase bucket schedule."""
	412	name = f"grug/datakit_moe_mix_{_SLUG}"
	413	version = resolve_version(name, version)
	414
	415	def build_config(ctx: StepContext) -> GrugMoeLaunchConfig:
	416	if ctx.is_fingerprint:
	417	val_components = {v.name: _val_component(ctx.artifact_path(v)) for v in _VALIDATION}
	418	else:
	419	val_components = {v.name: ctx.resolved(v).as_component() for v in _VALIDATION}
	420	data = _datakit_data_config(
	421	total_steps=_steps,
	422	batch_size=_batch_size,
	423	max_seq_len=_model.max_seq_len,
	424	enable_simulated_epoching=ENABLE_SIMULATED_EPOCHING,
	425	val_components=val_components,
	426	)
	427	return GrugMoeLaunchConfig(
	428	model=_model,
	429	data=data,
	430	output_path=ctx.output_path,
	431	run_id="datakit-moe-mix",
	432	resources=ctx.runtime_arg("train_resources"),
	433	steps=_steps,
	434	batch_size=_batch_size,
	435	seed=0,
	436	mp="params=float32,compute=bfloat16,output=bfloat16",
	437	tracker=WandbConfig(
	438	project="marin_moe",
	439	tags=["moe", "datakit_store_mix", _SLUG],
	440	group="datakit-moe-mix",
	441	name=None,
	442	),
	443	optimizer=_optimizer,
	444	grug_trainer=GrugTrainerConfig(z_loss_weight=1e-4, ema_beta=None, log_every=1),
	445	eval=GrugEvalConfig(
	446	eval_batch_size=512,
	447	steps_per_eval=1000,
	448	max_eval_batches=8,
	449	eval_current=True,
	450	eval_ema=False,
	451	),
	452	)
	453
	454	return ArtifactStep(
	455	name=user_namespaced_name(name, version),
	456	version=version,
	457	artifact_type=LevanterCheckpoint,
	458	run=run_grug_moe_trial,
	459	build_config=build_config,
	460	deps=tuple(_VALIDATION),
	461	runtime_args={"train_resources": _TRAIN_RESOURCES},
	462	)
	463
	464
	465	if __name__ == "__main__":
	466	experiment_main(build)()

optimizer.py

removed

	experiments/grug/moe/optimizer.py
t	1	# Copyright The Marin Authors	t
	2	# SPDX-License-Identifier: Apache-2.0
	3
	4	from dataclasses import dataclass
	5
	6	import jax
	7	import jax.numpy as jnp
	8	import optax
	9	from levanter.optim.config import OptimizerConfig
	10	from levanter.optim.grugmuon import _grug_scale_with_muon
	11	from levanter.optim.util import CoefficientType
	12	from levanter.utils.jax_utils import leaf_key_paths
	13
	14	from experiments.grug.moe.adamh import scale_by_adamh
	15
	16
	17	def _target_named_sharding(array) -> jax.sharding.NamedSharding \| None:
	18	if array is None or not hasattr(array, "shape"):
	19	return None
	20	sharding = getattr(array, "sharding", None)
	21	if sharding is None:
	22	aval = jax.typeof(array)
	23	sharding = getattr(aval, "sharding", None)
	24	if isinstance(sharding, jax.sharding.NamedSharding):
	25	return sharding
	26	return None
	27
	28
	29	def _match_named_update_sharding() -> optax.GradientTransformation:
	30	"""Restore named mesh sharding without touching single-device arrays."""
	31
	32	def init_fn(params):
	33	del params
	34	return optax.EmptyState()
	35
	36	def update_fn(updates, state, params=None):
	37	if params is None:
	38	return updates, state
	39
	40	def match_sharding(update, param):
	41	if update is None:
	42	return None
	43	target_sharding = _target_named_sharding(param)
	44	if target_sharding is None:
	45	return update
	46	return jax.sharding.reshard(update, target_sharding)
	47
	48	updates = jax.tree.map(match_sharding, updates, params, is_leaf=lambda x: x is None)
	49	return updates, state
	50
	51	return optax.GradientTransformation(init_fn, update_fn)
	52
	53
	54	def _match_named_sharding_to_params(updates, params):
	55	def match_sharding(update, param):
	56	if update is None:
	57	return None
	58	target_sharding = _target_named_sharding(param)
	59	if target_sharding is None:
	60	return update
	61	return jax.sharding.reshard(update, target_sharding)
	62
	63	return jax.tree.map(match_sharding, updates, params, is_leaf=lambda x: x is None)
	64
	65
	66	def _scale_invariant_hyperball_updates(params, direction_updates, learning_rate: float):
	67	direction_updates = _match_named_sharding_to_params(direction_updates, params)
	68
	69	def scale_invariant_update(param, update):
	70	if update is None:
	71	return None
	72	if not hasattr(param, "ndim"):
	73	return update
	74	if param.ndim == 2:
	75	param_norm = jnp.linalg.norm(param)
	76	update_norm = jnp.linalg.norm(update)
	77	new_param = param - learning_rate * update * param_norm / jnp.maximum(update_norm, 1e-10)
	78	new_param_norm = jnp.linalg.norm(new_param)
	79	return new_param / jnp.maximum(new_param_norm, 1e-10) * param_norm - param
	80
	81	axes = tuple(range(1, param.ndim))
	82	param_norm = jnp.sqrt(jnp.sum(jnp.square(param), axis=axes, keepdims=True))
	83	update_norm = jnp.sqrt(jnp.sum(jnp.square(update), axis=axes, keepdims=True))
	84	new_param = param - learning_rate * update * param_norm / jnp.maximum(update_norm, 1e-10)
	85	new_param_norm = jnp.sqrt(jnp.sum(jnp.square(new_param), axis=axes, keepdims=True))
	86	return new_param / jnp.maximum(new_param_norm, 1e-10) * param_norm - param
	87
	88	return jax.tree.map(
	89	scale_invariant_update,
	90	params,
	91	direction_updates,
	92	is_leaf=lambda x: x is None,
	93	)
	94
	95
	96	def scale_with_grug_muonh(
	97	momentum: float = 0.95,
	98	nesterov: bool = True,
	99	steps: int = 5,
	100	muon_eps: float = 1e-8,
	101	learning_rate: float = 0.02,
	102	coefficient_type: CoefficientType = "quintic",
	103	) -> optax.GradientTransformation:
	104	"""MuonH transform for raw Grug arrays with matrix-shaped trailing dims."""
	105	muon_transform = _grug_scale_with_muon(
	106	momentum=momentum,
	107	nesterov=nesterov,
	108	steps=steps,
	109	muon_eps=muon_eps,
	110	use_kimi_scaling=False,
	111	coefficient_type=coefficient_type,
	112	)
	113
	114	def init_fn(params):
	115	return muon_transform.init(params)
	116
	117	def update_fn(updates, state, params=None):
	118	if params is None:
	119	raise ValueError("scale_with_grug_muonh requires params for norm-preserving updates")
	120
	121	muon_updates, next_state = muon_transform.update(updates, state, params)
	122	muonh_updates = _scale_invariant_hyperball_updates(params, muon_updates, learning_rate)
	123	return muonh_updates, next_state
	124
	125	return optax.GradientTransformation(init_fn, update_fn)
	126
	127
	128	@OptimizerConfig.register_subclass("grug_moe_adamh_v2")
	129	@dataclass(frozen=True)
	130	class GrugMoeAdamHConfig(OptimizerConfig):
	131	"""AdamH for Grug MoE. Four optimizer groups, no flags.
	132
	133	- adamh: attention weights, dense MLP weights (2D matrices)
	134	- adamh_expert: expert MLP weights (mlp.expert_mlp.w_gate,
	135	mlp.expert_mlp.w_up, mlp.expert_mlp.w_down, shared.w_*)
	136	- adam: norms, biases, router, embeddings, attention gates (1D / small params)
	137	"""
	138
	139	beta1: float = 0.9
	140	beta2: float = 0.95
	141	epsilon: float = 1e-8
	142	max_grad_norm: float \| None = 1.0
	143	adam_lr: float = 6e-4
	144	expert_lr: float \| None = None
	145
	146	def build(self, num_train_steps):
	147	learning_rate_schedule = self.lr_scheduler(num_train_steps)
	148	adam_lr_schedule = self.lr_scheduler(num_train_steps, override_lr=self.adam_lr)
	149	expert_lr_val = self.expert_lr if self.expert_lr is not None else self.learning_rate
	150	expert_lr_schedule = self.lr_scheduler(num_train_steps, override_lr=expert_lr_val)
	151
	152	def optimizer(learning_rate, adam_lr, expert_lr):
	153	def adamh_transform():
	154	components = []
	155	if self.max_grad_norm:
	156	components.append(optax.clip_by_global_norm(self.max_grad_norm))
	157	components.append(scale_by_adamh(self.beta1, self.beta2, self.epsilon, learning_rate))
	158	return optax.chain(*components)
	159
	160	def adamh_expert_transform():
	161	components = []
	162	if self.max_grad_norm:
	163	components.append(optax.clip_by_global_norm(self.max_grad_norm))
	164	components.append(scale_by_adamh(self.beta1, self.beta2, self.epsilon, expert_lr))
	165	return optax.chain(*components)
	166
	167	def adam_transform():
	168	components = []
	169	if self.max_grad_norm:
	170	components.append(optax.clip_by_global_norm(self.max_grad_norm))
	171	components.append(optax.scale_by_adam(self.beta1, self.beta2, self.epsilon))
	172	components.append(optax.scale(-adam_lr))
	173	return optax.chain(*components)
	174
	175	return optax.multi_transform(
	176	{
	177	"adamh": adamh_transform(),
	178	"adamh_expert": adamh_expert_transform(),
	179	"adam": adam_transform(),
	180	},
	181	self.create_mask,
	182	)
	183
	184	return optax.inject_hyperparams(optimizer)(
	185	learning_rate=learning_rate_schedule,
	186	adam_lr=adam_lr_schedule,
	187	expert_lr=expert_lr_schedule,
	188	)
	189
	190	def create_mask(self, params):
	191	paths = leaf_key_paths(params)
	192
	193	def mask_fn(param, path):
	194	path_str = ".".join(path) if isinstance(path, (list, tuple)) else str(path)
	195	path_lower = path_str.lower()
	196	if "token_embed" in path_lower:
	197	return "adam"
	198	if "router_bias" in path_lower or "attn_gate" in path_lower or ".router" in path_lower:
	199	return "adam"
	200	if ".mlp.expert_mlp.w_" in path_lower or ".mlp.w_" in path_lower or ".shared.w_" in path_lower:
	201	return "adamh_expert"
	202	if hasattr(param, "ndim") and param.ndim >= 2:
	203	return "adamh"
	204	return "adam"
	205
	206	return jax.tree.map(mask_fn, params, paths)
	207
	208
	209	@OptimizerConfig.register_subclass("grug_moe_muonh_v1")
	210	@dataclass(frozen=True)
	211	class GrugMoeMuonHConfig(OptimizerConfig):
	212	"""May Recipe MuonH optimizer: 3 LR groups (muonh / adamh / adam).
	213
	214	Three LR groups:
	215	- ``muonh``: matrices (attn, MoE MLP, shared) and all GatedNorms.
	216	Newton-Schulz orthogonalisation + Frobenius hyperball scale-invariant step.
	217	- ``adamh``: ``lm_head`` / ``output_proj``.
	218	- ``adam``: ``token_embed`` / ``router`` / ``router_bias`` / ``attn_gate``
	219	/ 1-D norm weights.
	220
	221	``max_grad_norm`` defaults to ``None`` here (no clipping) for the 1pct-noclip
	222	schedule used by the May Recipe baseline.
	223	"""
	224
	225	adam_lr: float = 6e-4
	226	momentum: float = 0.95
	227	nesterov: bool = True
	228	backend_steps: int = 5
	229	beta1: float = 0.9
	230	beta2: float = 0.95
	231	epsilon: float = 1e-8
	232	muon_epsilon: float = 1e-8
	233	max_grad_norm: float \| None = None
	234	coefficient_type: CoefficientType = "quintic"
	235
	236	def build(self, num_train_steps):
	237	learning_rate_schedule = self.lr_scheduler(num_train_steps)
	238	adam_lr_schedule = self.lr_scheduler(num_train_steps, override_lr=self.adam_lr)
	239
	240	def optimizer(learning_rate, adam_lr):
	241	def muonh_transform():
	242	components = []
	243	if self.max_grad_norm:
	244	components.append(optax.clip_by_global_norm(self.max_grad_norm))
	245	components.append(
	246	scale_with_grug_muonh(
	247	momentum=self.momentum,
	248	nesterov=self.nesterov,
	249	steps=self.backend_steps,
	250	muon_eps=self.muon_epsilon,
	251	learning_rate=learning_rate,
	252	coefficient_type=self.coefficient_type,
	253	)
	254	)
	255	components.append(_match_named_update_sharding())
	256	return optax.chain(*components)
	257
	258	def adamh_transform_at(lr):
	259	components = []
	260	if self.max_grad_norm:
	261	components.append(optax.clip_by_global_norm(self.max_grad_norm))
	262	components.append(scale_by_adamh(self.beta1, self.beta2, self.epsilon, lr))
	263	return optax.chain(*components)
	264
	265	def adam_transform_at(lr):
	266	components = []
	267	if self.max_grad_norm:
	268	components.append(optax.clip_by_global_norm(self.max_grad_norm))
	269	components.append(optax.scale_by_adam(self.beta1, self.beta2, self.epsilon))
	270	components.append(optax.scale(-lr))
	271	return optax.chain(*components)
	272
	273	transforms = {
	274	"muonh": muonh_transform(),
	275	"adamh": adamh_transform_at(learning_rate),
	276	"adam": adam_transform_at(adam_lr),
	277	}
	278	return optax.multi_transform(transforms, self.create_mask)
	279
	280	return optax.inject_hyperparams(optimizer)(
	281	learning_rate=learning_rate_schedule,
	282	adam_lr=adam_lr_schedule,
	283	)
	284
	285	def create_mask(self, params):
	286	paths = leaf_key_paths(params)
	287
	288	def mask_fn(param, path):
	289	path_str = ".".join(path) if isinstance(path, (list, tuple)) else str(path)
	290	path_lower = path_str.lower()
	291	if (
	292	"token_embed" in path_lower
	293	or "router_bias" in path_lower
	294	or path_lower.endswith(".attn_gate")
	295	or ".router" in path_lower
	296	):
	297	return "adam"
	298	if "output_proj" in path_lower or "lm_head" in path_lower:
	299	return "adamh"
	300	# GatedNorms route to muonh (NS + Frobenius hyperball), same as matrices.
	301	if "gated_norm" in path_lower:
	302	return "muonh"
	303	if hasattr(param, "ndim") and param.ndim in (2, 3):
	304	return "muonh"
	305	return "adam"
	306
	307	return jax.tree.map(mask_fn, params, paths)
	308
	309
	310	__all__ = [
	311	"GrugMoeAdamHConfig",
	312	"GrugMoeMuonHConfig",
	313	"scale_with_grug_muonh",
	314	]

test_optimizer.py

removed

	experiments/grug/moe/test_optimizer.py
t	1	# Copyright The Marin Authors	t
	2	# SPDX-License-Identifier: Apache-2.0
	3
	4	import jax.numpy as jnp
	5
	6	from experiments.grug.moe.optimizer import GrugMoeAdamHConfig
	7
	8
	9	def test_grug_moe_adamh_mask_routes_expert_mlp_weights_to_expert_group():
	10	params = {
	11	"blocks": {
	12	"0": {
	13	"mlp": {
	14	"router": jnp.ones((8, 4), dtype=jnp.float32),
	15	"expert_mlp": {
	16	"w_gate": jnp.ones((4, 8, 16), dtype=jnp.float32),
	17	"w_up": jnp.ones((4, 8, 16), dtype=jnp.float32),
	18	"w_down": jnp.ones((4, 16, 8), dtype=jnp.float32),
	19	},
	20	},
	21	"shared": {
	22	"w_gate": jnp.ones((8, 16), dtype=jnp.float32),
	23	},
	24	},
	25	},
	26	"token_embed": jnp.ones((128, 8), dtype=jnp.float32),
	27	}
	28
	29	mask = GrugMoeAdamHConfig().create_mask(params)
	30
	31	block_mask = mask["blocks"]["0"]
	32	assert block_mask["mlp"]["router"] == "adam"
	33	assert block_mask["mlp"]["expert_mlp"]["w_gate"] == "adamh_expert"
	34	assert block_mask["mlp"]["expert_mlp"]["w_up"] == "adamh_expert"
	35	assert block_mask["mlp"]["expert_mlp"]["w_down"] == "adamh_expert"
	36	assert block_mask["shared"]["w_gate"] == "adamh_expert"
	37	assert mask["token_embed"] == "adam"