Directory Diff Report

Left: experiments/grug/base

Right: experiments/grug/moe_pipeline

Filter: .py, .md

Generated: 2026-09-08 20:33:26Z

Inline Diffs: 6 files

Changed

Added

Removed

Unchanged

File	Status	Lines Added	Lines Deleted
model.py	changed	+272	-163
train.py	changed	+377	-633
README.md	added	+61	-0
benchmark.py	added	+99	-0
pipeline.py	added	+624	-0
launch.py	removed	+0	-263
__init__.py	unchanged	+0	-0

model.py

changed

	experiments/grug/base/model.py			experiments/grug/moe_pipeline/model.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			n	4	"""MoE grug variant model.
				5
				6	Architecture: QB-routed MoE with GatedNorm, XSA, sigmoid combine weights.
				7	No load-balancing loss; router z-loss only. All layers are MoE (no dense layers).
				8	"""
				9
	4	import dataclasses		10	import dataclasses
	5	from dataclasses import dataclass		11	from dataclasses import dataclass
n			n	12	from typing import Any, Literal
	6			13
	7	import equinox as eqx		14	import equinox as eqx
	8	import jax		15	import jax
	9	import jax.numpy as jnp		16	import jax.numpy as jnp
n	10	from einops import rearrange	n	17	from haliax import Axis
	11	from haliax.jax_utils import named_call		18	from haliax.jax_utils import named_call
	12	from jax import random		19	from jax import random
n	13	from jax.sharding import PartitionSpec as P	n
	14	from jax.sharding import reshard		20	from jax.sharding import reshard
	15	from jaxtyping import Array, Float, Int, PRNGKeyArray		21	from jaxtyping import Array, Float, Int, PRNGKeyArray
n	16	from levanter.analysis.backward_flow import (	n	22	from levanter.compat.hf_checkpoints import HFCheckpointConverter
	17	is_backward_flow_active,		23	from levanter.grug.attention import (
	18	log_backward_activation,		24	AttentionMask,
	19	trace_backward_activation,		25	GrugAttentionImplementation,
	20	trace_grads,		26	RotaryConfig,
	21	)		27	)
n	22	from levanter.grug.attention import AttentionMask, RotaryConfig, apply_rotary_embedding, attention	n	28	from levanter.grug.grug_moe import (
				29	MOE_REMAT_SAVE_NAMES,
				30	MoeActivation,
				31	MoeImplementation,
				32	resolve_moe_implementation,
				33	)
	23	from levanter.grug.loss import fused_linear_softmax_cross_entropy_loss		34	from levanter.grug.loss import fused_linear_softmax_cross_entropy_loss
n	24	from levanter.grug.sharding import Pbatch, Pembed_vocab, Plm_head, Plogits, unshard	n	35	from levanter.grug.sharding import Pembed_vocab, Plm_head
				36	from levanter.tracker.histogram import SummaryStats
				37	from transformers import PretrainedConfig as HfConfig
				38
				39	from experiments.grug.moe.model import _BATCH_AXES as BATCH_AXES # noqa: F401
				40	from experiments.grug.moe.model import GRUG_MOE_ARCHITECTURE as GRUG_MOE_ARCHITECTURE
				41	from experiments.grug.moe.model import GRUG_MOE_ARTIFACT_SCHEMA_VERSION as GRUG_MOE_ARTIFACT_SCHEMA_VERSION
				42	from experiments.grug.moe.model import (
				43	GRUG_MOE_ARTIFACT_SCHEMA_VERSION_KEY as GRUG_MOE_ARTIFACT_SCHEMA_VERSION_KEY,
				44	)
				45	from experiments.grug.moe.model import GRUG_MOE_MODEL_TYPE as GRUG_MOE_MODEL_TYPE
				46	from experiments.grug.moe.model import Block as Block
				47	from experiments.grug.moe.model import CausalSelfAttention as CausalSelfAttention
				48	from experiments.grug.moe.model import DenseMLP as DenseMLP
				49	from experiments.grug.moe.model import GatedNorm as GatedNorm
				50	from experiments.grug.moe.model import GrugMoeHfConfig as GrugMoeHfConfig
				51	from experiments.grug.moe.model import MoEMLP as MoEMLP
				52	from experiments.grug.moe.model import RMSNorm as RMSNorm
				53	from experiments.grug.moe.model import _batch_spec as _batch_spec
				54	from experiments.grug.moe.model import _hf_config_attr as _hf_config_attr
				55	from experiments.grug.moe.model import _init_weight as _init_weight
				56	from experiments.grug.moe.model import _summarize_router_metrics as _summarize_router_metrics
				57	from experiments.grug.moe.model import debug_mesh_and_token_pspec as debug_mesh_and_token_pspec
				58	from experiments.grug.moe.model import grugmoe_inference_state_dict as grugmoe_inference_state_dict
				59
				60	RematMode = Literal["recompute_all", "save_moe"]
	25			61
	26			62
	27	@dataclass(frozen=True)		63	@dataclass(frozen=True)
	28	class GrugModelConfig:		64	class GrugModelConfig:
n	29	"""Hyperparameters for the Grug Llama-style transformer."""	n	65	"""Architecture and execution settings for the Grug MoE transformer."""
	30			66
	31	vocab_size: int		67	vocab_size: int
n	32	hidden_dim: int = 2048	n	68	hidden_dim: int = 512
	33	intermediate_dim: int = 5632		69	intermediate_dim: int = 256
				70	shared_expert_intermediate_dim: int = 512
				71	num_experts: int = 256
				72	num_experts_per_token: int = 4
	34	num_layers: int = 24		73	num_layers: int = 6
	35	num_heads: int = 16		74	num_heads: int = 4
	36	num_kv_heads: int = 16		75	num_kv_heads: int = 1
	37	head_dim: int \| None = None		76	head_dim: int \| None = None
n	38	max_seq_len: int = 4096	n	77	max_seq_len: int = 8192
				78	sliding_window: int = 2048
	39	layer_norm_eps: float = 1e-5		79	layer_norm_eps: float = 1e-5
	40	initializer_std: float = 0.02		80	initializer_std: float = 0.02
n			n	81	qk_mult: float = 1.3
				82	router_z_loss_coef: float = 0.0
				83	disable_pko: bool = True
				84	"""When True (default), the every-4th + last 'long' layers skip Partial
				85	Key Offset (no shift of the second half of K, no doc-start zeroing). Short
				86	layers never had PKO. Set to False to re-enable PKO on long layers."""
				87	disable_long_rope: bool = True
				88	"""When True (default), the every-4th + last 'long' layers skip rotary
				89	embedding entirely (Q and K go into attention un-rotated). Short layers
				90	still apply half-RoPE. Set to False to keep RoPE on long layers."""
				91	attention_implementation: GrugAttentionImplementation \| None = None
				92	moe_implementation: MoeImplementation \| None = None
				93	capacity_factor: float = 1.0
				94	remat_mode: RematMode = "recompute_all"
				95	"""Per-block gradient checkpointing. "recompute_all" reruns the whole block in
				96	backward (lowest memory); "save_moe" keeps the tagged MoE dispatch tensors so
				97	backward skips re-running expert dispatch and its EP collectives."""
	41	rope: RotaryConfig = dataclasses.field(default_factory=RotaryConfig)		98	rope: RotaryConfig = dataclasses.field(default_factory=RotaryConfig)
	42			99
	43	def __post_init__(self) -> None:		100	def __post_init__(self) -> None:
	48	raise ValueError("vocab_size must be positive")		105	raise ValueError("vocab_size must be positive")
	49	if self.max_seq_len <= 0:		106	if self.max_seq_len <= 0:
	50	raise ValueError("max_seq_len must be positive")		107	raise ValueError("max_seq_len must be positive")
n			n	108	if self.num_experts <= 0:
				109	raise ValueError("num_experts must be positive")
				110	if self.num_experts_per_token <= 0:
				111	raise ValueError("num_experts_per_token must be positive")
				112	if self.num_experts_per_token > self.num_experts:
				113	raise ValueError("num_experts_per_token must be <= num_experts")
				114	if self.shared_expert_intermediate_dim < 0:
				115	raise ValueError("shared_expert_intermediate_dim must be non-negative")
				116	if self.capacity_factor <= 0:
				117	raise ValueError("capacity_factor must be positive")
				118	resolve_moe_implementation(self.moe_implementation)
				119
				120	@property
				121	def Embed(self) -> Axis:
				122	return Axis("embed", self.hidden_dim)
				123
				124	@property
				125	def model_type(self) -> type["Transformer"]:
				126	return Transformer
	51			127
	52	@property		128	@property
	53	def inferred_head_dim(self) -> int:		129	def inferred_head_dim(self) -> int:
	59	)		135	)
	60	return self.hidden_dim // self.num_heads		136	return self.hidden_dim // self.num_heads
	61			137
n			n	138	def build(self, Vocab: Axis, *, key: PRNGKeyArray) -> "Transformer":
				139	cfg = self if Vocab.size == self.vocab_size else dataclasses.replace(self, vocab_size=Vocab.size)
				140	return Transformer.init(cfg, key=key)
	62			141
n	63	class CausalSelfAttention(eqx.Module):	n	142	def hf_checkpoint_converter(
	64	w_q: jax.Array		143	self,
	65	w_k: jax.Array		144	ref_checkpoint: str \| None = None,
	66	w_v: jax.Array		145	) -> HFCheckpointConverter["GrugModelConfig"]: # type: ignore[type-var]
	67	w_o: jax.Array		146	return HFCheckpointConverter(
	68	cfg: GrugModelConfig = eqx.field(static=True)		147	self.__class__,
				148	reference_checkpoint=ref_checkpoint,
				149	HfConfigClass=GrugMoeHfConfig,
				150	tokenizer=ref_checkpoint,
				151	)
	69			152
n	70	@staticmethod	n	153	@classmethod
	71	def init(cfg: GrugModelConfig, *, key: PRNGKeyArray) -> "CausalSelfAttention":		154	def from_hf_config(cls, hf_config: HfConfig) -> "GrugModelConfig":
	72	k_q, k_k, k_v, k_o = random.split(key, 4)		155	rope = RotaryConfig(theta=float(_hf_config_attr(hf_config, ("rope_theta",), 10000.0)))
	73	d_model, n_heads, n_kv_heads, head_dim = cfg.hidden_dim, cfg.num_heads, cfg.num_kv_heads, cfg.inferred_head_dim		156	return cls(
	74	return CausalSelfAttention(		157	vocab_size=int(_hf_config_attr(hf_config, ("vocab_size",))),
	75	w_q=reshard(_init_weight(k_q, (d_model, n_heads * head_dim), cfg.initializer_std), P("data", "model")),		158	hidden_dim=int(_hf_config_attr(hf_config, ("hidden_dim", "hidden_size"), 2048)),
	76	w_k=reshard(_init_weight(k_k, (d_model, n_kv_heads * head_dim), cfg.initializer_std), P("data", "model")),		159	intermediate_dim=int(
	77	w_v=reshard(_init_weight(k_v, (d_model, n_kv_heads * head_dim), cfg.initializer_std), P("data", "model")),		160	_hf_config_attr(hf_config, ("intermediate_dim", "moe_intermediate_size", "intermediate_size"), 5632)
	78	w_o=reshard(_init_weight(k_o, (n_heads * head_dim, d_model), cfg.initializer_std), P("model", "data")),
	79	cfg=cfg,		161	),
				162	shared_expert_intermediate_dim=int(
				163	_hf_config_attr(
				164	hf_config,
				165	("shared_expert_intermediate_dim", "shared_expert_intermediate_size"),
				166	5632,
				167	)
				168	),
				169	num_experts=int(_hf_config_attr(hf_config, ("num_experts", "num_local_experts"), 8)),
				170	num_experts_per_token=int(_hf_config_attr(hf_config, ("num_experts_per_token", "num_experts_per_tok"), 2)),
				171	num_layers=int(_hf_config_attr(hf_config, ("num_layers", "num_hidden_layers"), 24)),
				172	num_heads=int(_hf_config_attr(hf_config, ("num_heads", "num_attention_heads"), 16)),
				173	num_kv_heads=int(_hf_config_attr(hf_config, ("num_kv_heads", "num_key_value_heads"), 16)),
				174	head_dim=_hf_config_attr(hf_config, ("head_dim", "attention_head_dim")),
				175	max_seq_len=int(_hf_config_attr(hf_config, ("max_seq_len", "max_position_embeddings"), 4096)),
				176	sliding_window=int(_hf_config_attr(hf_config, ("sliding_window",), 4096)),
				177	layer_norm_eps=float(_hf_config_attr(hf_config, ("layer_norm_eps", "rms_norm_eps"), 1e-5)),
				178	initializer_std=float(_hf_config_attr(hf_config, ("initializer_std", "initializer_range"), 0.02)),
				179	qk_mult=float(_hf_config_attr(hf_config, ("qk_mult",), 1.0)),
				180	rope=rope,
				181	)
				182
				183	def to_hf_config(self, vocab_size: int, config_overrides: dict[str, Any] \| None = None) -> GrugMoeHfConfig:
				184	# One name per field: core fields take the universal transformers spelling, MoE fields the
				185	# most common public spelling, and grug-specific extras keep their bare names. from_hf_config
				186	# stays tolerant of the older spellings so existing artifacts keep loading.
				187	config = {
				188	"architectures": [GRUG_MOE_ARCHITECTURE],
				189	"vocab_size": vocab_size,
				190	# core — universal transformers names
				191	"hidden_size": self.hidden_dim,
				192	"num_hidden_layers": self.num_layers,
				193	"num_attention_heads": self.num_heads,
				194	"num_key_value_heads": self.num_kv_heads,
				195	"head_dim": self.inferred_head_dim,
				196	"max_position_embeddings": self.max_seq_len,
				197	"sliding_window": self.sliding_window,
				198	"rms_norm_eps": self.layer_norm_eps,
				199	"initializer_range": self.initializer_std,
				200	"rope_theta": self.rope.theta,
				201	"tie_word_embeddings": False,
				202	# MoE — most common public spelling per field
				203	"num_experts": self.num_experts,
				204	"num_experts_per_tok": self.num_experts_per_token,
				205	"moe_intermediate_size": self.intermediate_dim,
				206	"shared_expert_intermediate_size": self.shared_expert_intermediate_dim,
				207	# grug-specific (no public equivalent)
				208	"qk_mult": self.qk_mult,
				209	"grugmoe_attention_mode": "production",
				210	GRUG_MOE_ARTIFACT_SCHEMA_VERSION_KEY: GRUG_MOE_ARTIFACT_SCHEMA_VERSION,
	80	)		211	}
	81			212	if config_overrides is not None:
	82	@trace_grads		213	config.update(config_overrides)
	83	@named_call		214	return GrugMoeHfConfig(**config)
	84	def __call__(self, x: Float[Array, "B S D"], mask: AttentionMask \| jax.Array) -> Float[Array, "B S D"]:
	85	head_dim = self.cfg.inferred_head_dim
	86	seq_len = x.shape[1]
	87
	88	q = rearrange(jnp.einsum("bsh,hd->bsd", x, self.w_q), "... (n d) -> ... n d", d=head_dim)
	89	k = rearrange(jnp.einsum("bsh,hd->bsd", x, self.w_k), "... (m d) -> ... m d", d=head_dim)
	90	v = rearrange(jnp.einsum("bsh,hd->bsd", x, self.w_v), "... (m d) -> ... m d", d=head_dim)
	91	q, k = apply_rotary_embedding(q, k, seq_len=seq_len, head_dim=head_dim, rope=self.cfg.rope)
	92	attn_out = attention(q, k, v, mask)
	93	attn_out = rearrange(attn_out, "... n d -> ... (n d)")
	94	return jnp.einsum("bsh,hd->bsd", attn_out, self.w_o, out_sharding=Pbatch)
	95
	96
	97	class MLP(eqx.Module):
	98	mlp_up: jax.Array
	99	mlp_down: jax.Array
	100
	101	@staticmethod
	102	def init(cfg: GrugModelConfig, *, key: PRNGKeyArray) -> "MLP":
	103	k_up, k_down = random.split(key, 2)
	104	d_model, d_ff = cfg.hidden_dim, cfg.intermediate_dim
	105	return MLP(
	106	mlp_up=reshard(_init_weight(k_up, (d_model, d_ff), cfg.initializer_std), P("data", "model")),
	107	mlp_down=reshard(_init_weight(k_down, (d_ff, d_model), cfg.initializer_std), P("model", "data")),
	108	)
	109
	110	@trace_grads
	111	@named_call
	112	def __call__(self, x: Float[Array, "B S D"]) -> Float[Array, "B S D"]:
	113	up = jnp.einsum("bsh,hm->bsm", x, self.mlp_up)
	114	activated = jax.nn.relu(up)
	115	return jnp.einsum("bsm,mh->bsh", activated, self.mlp_down, out_sharding=Pbatch)
	116
	117
	118	class RMSNorm(eqx.Module):
	119	weight: jax.Array
	120	eps: float = eqx.field(static=True)
	121
	122	@staticmethod
	123	def init(dim: int, eps: float) -> "RMSNorm":
	124	return RMSNorm(weight=jnp.ones((dim,), dtype=jnp.float32), eps=eps)
	125
	126	@named_call
	127	def __call__(self, x: Float[Array, "... D"]) -> Float[Array, "... D"]:
	128	weight = unshard(self.weight)
	129	dtype = x.dtype
	130	x = x.astype(jnp.float32)
	131	variance = jnp.mean(jnp.square(x), axis=-1, keepdims=True)
	132	normed = x * jax.lax.rsqrt(variance + self.eps)
	133	return (normed * weight).astype(dtype)
	134
	135
	136	class Block(eqx.Module):
	137	rms_attn: RMSNorm
	138	attn: CausalSelfAttention
	139	rms_mlp: RMSNorm
	140	mlp: MLP
	141
	142	@staticmethod
	143	def init(cfg: GrugModelConfig, *, key: PRNGKeyArray) -> "Block":
	144	attn_key, mlp_key = random.split(key, 2)
	145	return Block(
	146	rms_attn=RMSNorm.init(cfg.hidden_dim, cfg.layer_norm_eps),
	147	attn=CausalSelfAttention.init(cfg, key=attn_key),
	148	rms_mlp=RMSNorm.init(cfg.hidden_dim, cfg.layer_norm_eps),
	149	mlp=MLP.init(cfg, key=mlp_key),
	150	)
	151
	152	@named_call
	153	def __call__(self, x: Float[Array, "B S D"], mask: AttentionMask \| jax.Array) -> Float[Array, "B S D"]:
	154	x = trace_backward_activation(x, "resid_in")
	155	x = x + self.attn(self.rms_attn(x), mask)
	156	x = trace_backward_activation(x, "resid_post_attn")
	157	x = x + self.mlp(self.rms_mlp(x))
	158	return trace_backward_activation(x, "resid_out")
	159			215
	160			216
	161	class Transformer(eqx.Module):		217	class Transformer(eqx.Module):
	162	token_embed: jax.Array		218	token_embed: jax.Array
n			n	219	embed_norm: RMSNorm
				220	embed_gated_norm: GatedNorm
	163	output_proj: jax.Array		221	output_proj: jax.Array
	164	blocks: tuple[Block, ...]		222	blocks: tuple[Block, ...]
	165	final_norm: RMSNorm		223	final_norm: RMSNorm
n			n	224	final_gated_norm: GatedNorm
	166	config: GrugModelConfig = eqx.field(static=True)		225	config: GrugModelConfig = eqx.field(static=True)
	167			226
	168	@staticmethod		227	@staticmethod
n	169	def init(cfg: GrugModelConfig, *, key: PRNGKeyArray) -> "Transformer":	n	228	def init(
				229	cfg_or_vocab: GrugModelConfig \| Axis,
				230	config: GrugModelConfig \| None = None,
				231	*,
				232	key: PRNGKeyArray,
				233	) -> "Transformer":
				234	if isinstance(cfg_or_vocab, Axis):
				235	if config is None:
				236	raise ValueError("config must be provided when initializing with a Vocab axis")
				237	cfg = (
				238	config
				239	if cfg_or_vocab.size == config.vocab_size
				240	else dataclasses.replace(config, vocab_size=cfg_or_vocab.size)
				241	)
				242	else:
				243	if config is not None:
				244	raise ValueError("config must not be provided when initializing directly from GrugModelConfig")
				245	cfg = cfg_or_vocab
				246
	170	embed_key, out_key, *block_keys = random.split(key, cfg.num_layers + 2)		247	embed_key, out_key, embed_gn_key, final_gn_key, *block_keys = random.split(key, cfg.num_layers + 4)
	171	token_embed = reshard(		248	token_embed = reshard(
	172	_init_weight(embed_key, (cfg.vocab_size, cfg.hidden_dim), cfg.initializer_std), Pembed_vocab		249	_init_weight(embed_key, (cfg.vocab_size, cfg.hidden_dim), cfg.initializer_std), Pembed_vocab
	173	)		250	)
	174	output_proj = reshard(_init_weight(out_key, (cfg.hidden_dim, cfg.vocab_size), cfg.initializer_std), Plm_head)		251	output_proj = reshard(_init_weight(out_key, (cfg.hidden_dim, cfg.vocab_size), cfg.initializer_std), Plm_head)
n	175	blocks = tuple(Block.init(cfg, key=layer_key) for layer_key in block_keys)	n	252	blocks = tuple(Block.init(cfg, key=block_keys[i]) for i in range(cfg.num_layers))
	176	final_norm = RMSNorm.init(cfg.hidden_dim, cfg.layer_norm_eps)
	177	return Transformer(		253	return Transformer(
	178	token_embed=token_embed,		254	token_embed=token_embed,
n			n	255	embed_norm=RMSNorm.init(cfg.hidden_dim, cfg.layer_norm_eps),
				256	embed_gated_norm=GatedNorm.init(cfg.hidden_dim, cfg.initializer_std, key=embed_gn_key),
	179	output_proj=output_proj,		257	output_proj=output_proj,
	180	blocks=blocks,		258	blocks=blocks,
n	181	final_norm=final_norm,	n	259	final_norm=RMSNorm.init(cfg.hidden_dim, cfg.layer_norm_eps),
				260	final_gated_norm=GatedNorm.init(cfg.hidden_dim, cfg.initializer_std, key=final_gn_key),
	182	config=cfg,		261	config=cfg,
	183	)		262	)
n			n	263
				264	@property
				265	def Vocab(self) -> Axis:
				266	return Axis("vocab", self.config.vocab_size)
	184			267
	185	@named_call		268	@named_call
	186	def __call__(		269	def __call__(
	187	self,		270	self,
	188	token_ids: Int[Array, "B S"],		271	token_ids: Int[Array, "B S"],
	189	mask: AttentionMask \| jax.Array \| None = None,		272	mask: AttentionMask \| jax.Array \| None = None,
n	190	) -> Float[Array, "B S D"]:	n	273	) -> tuple[Float[Array, "B S D"], dict[str, jax.Array]]:
	191	if mask is None:		274	if mask is None:
	192	mask = AttentionMask.causal()		275	mask = AttentionMask.causal()
	193			276
n	194	with jax.named_scope("token_embed"):	n	277	batch_spec = _batch_spec()
				278	cfg = self.config
	195	hidden = self.token_embed.at[token_ids].get(out_sharding=Pbatch)		279	hidden = self.token_embed.at[token_ids].get(out_sharding=batch_spec)
	196	hidden = log_backward_activation(hidden)		280	hidden = self.embed_gated_norm(self.embed_norm(hidden))
				281
				282	# Short layers: sliding window. Long layers (every 4th + last): full causal.
				283	segment_ids = mask.segment_ids if isinstance(mask, AttentionMask) else None
				284	short_mask = AttentionMask(is_causal=True, sliding_window=cfg.sliding_window, segment_ids=segment_ids)
				285	long_mask = AttentionMask(is_causal=True, sliding_window=None, segment_ids=segment_ids)
				286
				287	if cfg.remat_mode == "save_moe":
				288	remat_policy = jax.checkpoint_policies.save_only_these_names(*MOE_REMAT_SAVE_NAMES)
				289	else:
				290	remat_policy = None
				291
				292	num_blocks = len(self.blocks)
				293	moe_router_stats: list[dict[str, jax.Array]] = []
	197	for i, block in enumerate(self.blocks):		294	for i, block in enumerate(self.blocks):
n	198	with jax.named_scope(f"block_{i}"):	n	295	is_last = i == num_blocks - 1
	199	block_fn = block if is_backward_flow_active() else eqx.filter_checkpoint(block)		296	is_long = i % 4 == 3 or is_last
	200	hidden = block_fn(hidden, mask)		297	layer_mask = long_mask if is_long else short_mask
	201	with jax.named_scope("final_norm"):		298	hidden, router_stats = eqx.filter_checkpoint(block, policy=remat_policy)(
	202	hidden = self.final_norm(hidden)		299	hidden,
	203	return log_backward_activation(hidden)		300	layer_mask,
				301	use_pko=is_long and not cfg.disable_pko,
				302	disable_rope=is_long and cfg.disable_long_rope,
				303	)
				304	moe_router_stats.append(router_stats)
				305
				306	router_metrics = {
				307	"routing_entropy_per_layer": jnp.stack([s["routing_entropy"] for s in moe_router_stats], axis=0),
				308	"routing_counts_per_layer": jnp.stack([s["routing_counts"] for s in moe_router_stats], axis=0),
				309	"load_balancing_loss_per_layer": jnp.stack([s["load_balancing_loss"] for s in moe_router_stats], axis=0),
				310	"router_z_loss_per_layer": jnp.stack([s["router_z_loss"] for s in moe_router_stats], axis=0),
				311	"qb_beta_per_layer": jnp.stack([s["qb_beta"] for s in moe_router_stats], axis=0),
				312	"capacity_overflow_per_layer": jnp.stack([s["capacity_overflow"] for s in moe_router_stats], axis=0),
				313	}
				314	hidden = self.final_gated_norm(self.final_norm(hidden))
				315	return hidden, router_metrics
	204			316
	205	@named_call		317	@named_call
	206	def logits(		318	def logits(
	208	token_ids: Int[Array, "B S"],		320	token_ids: Int[Array, "B S"],
	209	mask: AttentionMask \| jax.Array \| None = None,		321	mask: AttentionMask \| jax.Array \| None = None,
	210	) -> Float[Array, "B S V"]:		322	) -> Float[Array, "B S V"]:
n			n	323	batch_spec = _batch_spec()
	211	hidden = self(token_ids, mask=mask)		324	hidden, _ = self(token_ids, mask=mask)
	212	return jnp.einsum("bsh,hd->bsd", hidden, self.output_proj, out_sharding=Plogits)		325	return jnp.einsum("bsh,hd->bsd", hidden, self.output_proj, out_sharding=batch_spec)
				326
				327	def to_state_dict(self, prefix: str \| None = None) -> dict[str, jax.Array]:
				328	return grugmoe_inference_state_dict(self, prefix=prefix)
	213			329
	214	def next_token_loss(		330	def next_token_loss(
	215	self,		331	self,
	220	reduction: str = "mean",		336	reduction: str = "mean",
	221	logsumexp_weight: float \| None = None,		337	logsumexp_weight: float \| None = None,
	222	loss_dtype: jnp.dtype = jnp.float32,		338	loss_dtype: jnp.dtype = jnp.float32,
n	223	loss_implementation: str \| tuple[str, ...] \| None = None,	n	339	return_router_metrics: bool = False,
	224	) -> jax.Array:		340	) -> jax.Array \| tuple[jax.Array, dict[str, jax.Array \| SummaryStats]]:
	225	"""Compute next-token cross-entropy loss for a batch."""
	226	hidden = self(token_ids, mask=mask)		341	hidden, router_metrics = self(token_ids, mask=mask)
	227	labels = jnp.pad(token_ids[:, 1:], ((0, 0), (0, 1))).astype(jnp.int32)		342	labels = jnp.pad(token_ids[:, 1:], ((0, 0), (0, 1))).astype(jnp.int32)
	228	loss_weight = loss_weight.astype(loss_dtype)		343	loss_weight = loss_weight.astype(loss_dtype)
	229			344
n	230	return fused_linear_softmax_cross_entropy_loss(	n	345	cross_entropy_loss = fused_linear_softmax_cross_entropy_loss(
	231	hidden,		346	hidden,
	232	self.output_proj,		347	self.output_proj,
	233	labels,		348	labels,
	235	reduction=reduction,		350	reduction=reduction,
	236	logsumexp_weight=logsumexp_weight,		351	logsumexp_weight=logsumexp_weight,
	237	dtype=loss_dtype,		352	dtype=loss_dtype,
n	238	implementation=loss_implementation,	n
	239	)
	240
	241
	242	def _init_weight(key: PRNGKeyArray, shape: tuple[int, ...], std: float) -> Float[Array, "..."]:
	243	return std * random.truncated_normal(key, -3, 3, shape)
	244
	245
	246	def debug_mesh_and_token_pspec(num_devices: int, model_axis_size: int = 1) -> tuple[jax.sharding.AbstractMesh, P]:
	247	"""Return a small abstract mesh and token sharding for lowering contract tests."""
	248	if num_devices <= 0:
	249	raise ValueError(f"num_devices must be positive, got {num_devices}")
	250	if model_axis_size <= 0:
	251	raise ValueError(f"model_axis_size must be positive, got {model_axis_size}")
	252	if num_devices % model_axis_size != 0:
	253	raise ValueError(f"num_devices ({num_devices}) must be divisible by model_axis_size ({model_axis_size})")
	254	data_axis_size = num_devices // model_axis_size
	255	mesh = jax.sharding.AbstractMesh(
	256	axis_sizes=(1, data_axis_size, model_axis_size),
	257	axis_names=("replica_dcn", "data", "model"),
	258	axis_types=(
	259	jax.sharding.AxisType.Explicit,
	260	jax.sharding.AxisType.Explicit,
	261	jax.sharding.AxisType.Explicit,
	262	),		353	)
	263	)		354	# No load-balancing loss; router z-loss only.
	264	return mesh, P(("replica_dcn", "data"), None)		355	num_moe_layers = router_metrics["router_z_loss_per_layer"].shape[0]
				356	router_z_loss = jnp.sum(router_metrics["router_z_loss_per_layer"]) / num_moe_layers
				357	aux_loss = self.config.router_z_loss_coef * router_z_loss
				358	loss = cross_entropy_loss + aux_loss if reduction != "none" else cross_entropy_loss
				359	if return_router_metrics:
				360	summarized_metrics = _summarize_router_metrics(router_metrics)
				361	summarized_metrics["train/cross_entropy_loss"] = cross_entropy_loss
				362	summarized_metrics["train/router/aux_loss_weighted"] = aux_loss
				363	return loss, summarized_metrics
				364	return loss
	265			365
	266			366
	267	__all__ = [		367	__all__ = [
n	268	"MLP",	n	368	"GRUG_MOE_ARCHITECTURE",
				369	"GRUG_MOE_ARTIFACT_SCHEMA_VERSION",
				370	"GRUG_MOE_ARTIFACT_SCHEMA_VERSION_KEY",
				371	"GRUG_MOE_MODEL_TYPE",
	269	"Block",		372	"Block",
	270	"CausalSelfAttention",		373	"CausalSelfAttention",
n			n	374	"DenseMLP",
				375	"GatedNorm",
	271	"GrugModelConfig",		376	"GrugModelConfig",
n			n	377	"GrugMoeHfConfig",
				378	"MoEMLP",
				379	"MoeActivation",
	272	"RMSNorm",		380	"RMSNorm",
	273	"Transformer",		381	"Transformer",
	274	"debug_mesh_and_token_pspec",		382	"debug_mesh_and_token_pspec",
t			t	383	"grugmoe_inference_state_dict",
	275	]		384	]

train.py

changed

	experiments/grug/base/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		17	import jax
	12	import jax.numpy as jnp		18	import jax.numpy as jnp
	13	import jmp		19	import jmp
n	14	import levanter.callbacks as callbacks	n	20	import numpy as np
	15	import levanter.tracker
	16	import optax		21	import optax
	17	from fray.cluster import ResourceConfig		22	from fray.cluster import ResourceConfig
n	18	from haliax import Axis	n	23	from fray.device_flops import device_flops
	19	from haliax.partitioning import set_mesh		24	from iris.runtime.jax_init import initialize_jax
	20	from jax.sharding import Mesh, NamedSharding		25	from jax.sharding import NamedSharding
	21	from jax.sharding import PartitionSpec as P		26	from jax.sharding import PartitionSpec as P
n	22	from jax.tree_util import register_dataclass	n	27	from levanter.callbacks.profiler import ProfileOptionsConfig, ProfilerConfig
	23	from jaxtyping import PRNGKeyArray		28	from levanter.data.text.examples import GrugLmExample
	24	from levanter.analysis.backward_flow import (		29	from levanter.pipeline import reshape_batch_into_microbatches
	25	BackwardFlowConfig,		30	from levanter.utils.flop_utils import lm_flops_per_token
	26	BackwardFlowGraph,		31
	27	backward_flow_graph_from_jaxpr,		32	from experiments.grug.dispatch import dispatch_grug_training_run
	28	backward_flow_node_stats,		33	from experiments.grug.moe_pipeline.model import BATCH_AXES, GrugModelConfig, Transformer
	29	capture_backward_flow,		34	from experiments.grug.moe_pipeline.pipeline import (
	30	collapse_backward_flow_graph,		35	TRAIN_LOSS_KEY,
	31	infer_backward_flow_render_hints,		36	AutomaticPipelineSchedule,
	32	render_backward_flow_html,		37	GrugMoePipelineConfig,
				38	automatic_stage_to_mpmd_indices,
				39	initialize_mpmd_automatic_pipeline_state,
				40	make_automatic_pipeline_step,
				41	make_pipeline_mesh,
				42	prepare_automatic_mpmd_step,
	33	)		43	)
n	34	from levanter.callbacks.state_adapter import StateCallbackRunner	n
	35	from levanter.callbacks.watch import WatchConfig, compute_watch_stats
	36	from levanter.data.dataset import AsyncDataset
	37	from levanter.data.loader import DataLoader
	38	from levanter.data.mixture import MixtureDataset, rescale_mixture_schedule_for_batch_schedule
	39	from levanter.data.text.datasets import LmDataConfig
	40	from levanter.data.text.examples import GrugLmExample, grug_lm_example_from_named
	41	from levanter.eval import TaggedEvaluator, cb_tagged_evaluate
	42	from levanter.grug.sharding import compact_grug_mesh
	43	from levanter.models.lm_model import LmExample
	44	from levanter.optim.config import AdamConfig, OptimizerConfig
	45	from levanter.schedule import BatchSchedule
	46	from levanter.trainer import TrainerConfig
	47	from levanter.training_control import TrainingDashboard
	48	from levanter.utils.flop_utils import lm_flops_per_token
	49	from levanter.utils.jax_utils import parameter_count
	50	from levanter.utils.logging import LoadingTimeTrackerIterator
	51			44
n	52	from experiments.grug.base.model import GrugModelConfig, Transformer	n
	53	from experiments.grug.checkpointing import init_weights_only_from_checkpoint, restore_grug_state_from_checkpoint
	54	from experiments.grug.dispatch import dispatch_grug_training_run
	55	from experiments.grug.sharding_dump import dump_grug_state_sharding_run_artifact
	56			45
n	57	logger = logging.getLogger(__name__)	n	46	class PipelineSchedule(StrEnum):
	58	_BACKWARD_FLOW_METRICS_KEY = "_backward_flow"		47	ZERO_BUBBLE = "automatic_zero_bubble"
	59	_BACKWARD_FLOW_DEFAULT_INTERVAL = 50		48	DUALPIPE_V = "automatic_dualpipe_v"
				49
				50	@property
				51	def automatic_schedule(self) -> AutomaticPipelineSchedule:
				52	if self == self.DUALPIPE_V:
				53	return AutomaticPipelineSchedule.DUALPIPE_V
				54	if self == self.ZERO_BUBBLE:
				55	return AutomaticPipelineSchedule.ZERO_BUBBLE
				56	raise ValueError(f"unknown pipeline schedule: {self.value}")
	60			57
	61			58
	62	@dataclass(frozen=True)		59	@dataclass(frozen=True)
n	63	class GrugTrainerConfig:	n	60	class GrugPipelineTrainConfig:
	64	"""Runtime knobs for grug training."""		61	"""Model, topology, and loop settings for pipeline-parallel training."""
	65			62
n	66	trainer: TrainerConfig = field(default_factory=lambda: TrainerConfig(use_explicit_mesh_axes=True))	n	63	stages: int
	67	train_batch_pspec: P = field(default_factory=lambda: P(("replica_dcn", "data")))		64	physical_stages: int
	68	data_seed: int \| None = None		65	microbatches: int
	69	log_every: int = 1
	70	ema_beta: float \| None = None # EMA coefficient for eval/checkpoint model; None disables EMA.
	71	z_loss_weight: float = 0.0 # Weight on logsumexp (z-loss) stabilization term.
	72	backward_flow: BackwardFlowConfig = field(
	73	default_factory=lambda: BackwardFlowConfig(interval=_BACKWARD_FLOW_DEFAULT_INTERVAL)
	74	)
	75	loss_implementation: str \| tuple[str, ...] \| None = None
	76	sharding_dump_path: str \| None = None
	77
	78
	79	@dataclass(frozen=True)
	80	class GrugEvalConfig:
	81	"""Perplexity eval settings for grug training."""
	82
	83	eval_batch_size: int = 512		66	batch_size: int
	84	eval_batch_pspec: P = field(default_factory=lambda: P(("replica_dcn", "data")))		67	seq_len: int
	85	steps_per_eval: int \| None = 1000		68	hidden_dim: int
	86	max_eval_batches: int \| None = None		69	intermediate_dim: int
	87	prefix: str = "eval"		70	shared_expert_intermediate_dim: int
	88	eval_current: bool = True		71	num_layers: int
	89	eval_ema: bool = True		72	num_experts: int
	90	compute_bpb: bool = True		73	top_k: int
				74	expert_axis_size: int
				75	vocab_size: int
				76	num_heads: int
				77	num_kv_heads: int
				78	sliding_window: int
				79	qk_mult: float
				80	steps: int
				81	warmup_steps: int
				82	profile_start_step: int
				83	profile_steps: int
				84	profile_run_id: str \| None
				85	layer_counts: tuple[int, ...] \| None
				86	memory_threshold: int \| None
				87	mp_policy_string: str
				88	remat_mode: str
				89	attention_implementation: str
				90	moe_implementation: str
				91	schedule: PipelineSchedule
				92
				93	def __post_init__(self) -> None:
				94	if not isinstance(self.schedule, PipelineSchedule):
				95	raise ValueError(f"unknown pipeline schedule: {self.schedule!r}")
				96	if self.schedule == PipelineSchedule.ZERO_BUBBLE and self.stages != self.physical_stages:
				97	raise ValueError(
				98	"automatic_zero_bubble requires one logical stage per physical stage; "
				99	f"got {self.stages} logical and {self.physical_stages} physical stages"
				100	)
				101	if self.schedule == PipelineSchedule.DUALPIPE_V:
				102	if self.stages != 2 * self.physical_stages:
				103	raise ValueError(
				104	"automatic_dualpipe_v requires two logical stages per physical stage; "
				105	f"got {self.stages} logical and {self.physical_stages} physical stages"
				106	)
				107	if self.microbatches < self.stages:
				108	raise ValueError(
				109	f"automatic_dualpipe_v requires at least {self.stages} microbatches, got {self.microbatches}"
				110	)
	91			111
	92			112
	93	@dataclass(frozen=True)		113	@dataclass(frozen=True)
	94	class GrugRunConfig:		114	class GrugRunConfig:
n	95	"""Top-level config for grug training."""	n	115	"""Resources and identity for one pipeline-parallel training run."""
	96			116
n	97	model: GrugModelConfig	n	117	run_id: str
	98	data: LmDataConfig		118	train: GrugPipelineTrainConfig
	99	resources: ResourceConfig		119	resources: ResourceConfig
n	100	output_path: str = ""	n
	101	optimizer: OptimizerConfig = field(default_factory=AdamConfig)
	102	trainer: GrugTrainerConfig = field(default_factory=GrugTrainerConfig)
	103	eval: GrugEvalConfig \| None = field(default_factory=GrugEvalConfig)
	104			120
	105			121
n	106	def build_train_dataset(	n	122	def _log(event: str, **values) -> None:
	107	data_config: LmDataConfig,		123	print(f"{event} {json.dumps(values, sort_keys=True, default=str)}", flush=True)
				124
				125
				126	def _validate_local_mesh(
	108	*,		127	*,
n	109	max_seq_len: int,	n	128	local_device_count: int,
	110	batch_schedule: BatchSchedule,		129	devices_per_stage: int,
	111	key: PRNGKeyArray,		130	expert_axis_size: int,
	112	) -> MixtureDataset[GrugLmExample]:		131	batch_size: int,
	113	pos = Axis("position", max_seq_len)		132	microbatches: int,
	114	mix_key, shuffle_key = jax.random.split(key)		133	) -> None:
	115	weights = data_config.train_weights		134	if local_device_count % expert_axis_size != 0:
	116	if isinstance(weights, list):		135	raise ValueError(f"local device count {local_device_count} must be divisible by expert axis {expert_axis_size}")
	117	weights = rescale_mixture_schedule_for_batch_schedule(weights, batch_schedule)		136	if batch_size % microbatches != 0:
	118			137	raise ValueError(f"batch size {batch_size} must be divisible by {microbatches} microbatches")
	119	initial_batch_size = batch_schedule.batch_size_at_step(0)		138	microbatch_size = batch_size // microbatches
	120	datasets = data_config.train_sets(pos, key=shuffle_key, initial_batch_size=initial_batch_size)		139	if microbatch_size % devices_per_stage != 0:
	121	return MixtureDataset(		140	raise ValueError(
	122	datasets=datasets,		141	f"microbatch size {microbatch_size} must be divisible by the {devices_per_stage} devices in each stage"
	123	weights=weights,
	124	stop_strategy=data_config.stop_strategy,
	125	key=mix_key,
	126	block_size=data_config.mixture_block_size,
	127	)
	128
	129
	130	def build_train_loader(
	131	dataset: AsyncDataset[GrugLmExample],
	132	*,
	133	batch_schedule: BatchSchedule,
	134	mesh: Mesh,
	135	batch_pspec: P = P(("replica_dcn", "data")),
	136	) -> DataLoader[GrugLmExample]:
	137	# DataLoader uses this batch axis mapping to shard batches across the distributed mesh.
	138	axis_resource = batch_pspec[0]
	139	return DataLoader(
	140	dataset,
	141	batch_schedule.schedule,
	142	mesh=mesh,
	143	axis_resources={"__BATCH__": axis_resource},
	144	batch_axis_name="__BATCH__",
	145	allow_nondivisible_batch_size=False,
	146	)
	147
	148
	149	def build_tagged_evaluator(
	150	*,
	151	data_config: LmDataConfig,
	152	max_seq_len: int,
	153	mesh: Mesh,
	154	eval_cfg: GrugEvalConfig,
	155	) -> TaggedEvaluator[LmExample \| GrugLmExample, Transformer] \| None:
	156	pos = Axis("position", max_seq_len)
	157	tagged_eval_sets = data_config.tagged_eval_sets(pos)
	158	if len(tagged_eval_sets) == 0:
	159	logger.warning("No evaluation datasets provided.")
	160	return None
	161
	162	max_examples_per_dataset = None
	163	if eval_cfg.max_eval_batches is not None:
	164	max_examples_per_dataset = eval_cfg.max_eval_batches * eval_cfg.eval_batch_size
	165
	166	tokenizer = data_config.the_tokenizer if eval_cfg.compute_bpb else None
	167	batch_axis_resource = eval_cfg.eval_batch_pspec[0]
	168	eval_axis_mapping = {"batch": batch_axis_resource}
	169	eval_batch = Axis("batch", eval_cfg.eval_batch_size)
	170	eval_array_sharding = NamedSharding(mesh, P(batch_axis_resource, 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	)		142	)
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			143
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	)
	196			144
n			n	145	def _run_grug_local(config: GrugPipelineTrainConfig) -> None:
				146	os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
				147	initialize_jax()
	197			148
n	198	def _compute_flops(	n	149	from jaxpp.array import MpmdArray # noqa: PLC0415 # pyrefly: ignore[missing-import]
	199	*,		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 != 0:
				159	raise ValueError(
				160	f"process count {jax.process_count()} must be divisible by {config.physical_stages} physical stages"
				161	)
				162	replica_axis_size = jax.process_count() // config.physical_stages
				163	_validate_local_mesh(
				164	local_device_count=jax.local_device_count(),
				165	devices_per_stage=jax.device_count() // config.physical_stages,
				166	expert_axis_size=config.expert_axis_size,
				167	batch_size=config.batch_size,
				168	microbatches=config.microbatches,
				169	)
				170
	200	model_config: GrugModelConfig,		171	model_config = GrugModelConfig(
	201	) -> tuple[float, dict[str, float]]:		172	vocab_size=config.vocab_size,
				173	hidden_dim=config.hidden_dim,
				174	intermediate_dim=config.intermediate_dim,
				175	shared_expert_intermediate_dim=config.shared_expert_intermediate_dim,
				176	num_layers=config.num_layers,
				177	num_experts=config.num_experts,
				178	num_experts_per_token=config.top_k,
				179	num_heads=config.num_heads,
				180	num_kv_heads=config.num_kv_heads,
				181	max_seq_len=config.seq_len,
				182	sliding_window=config.sliding_window,
				183	qk_mult=config.qk_mult,
				184	router_z_loss_coef=0.0,
				185	attention_implementation=cast(str, config.attention_implementation),
				186	moe_implementation=cast(str, config.moe_implementation),
				187	remat_mode=cast(str, config.remat_mode),
				188	)
				189	mesh, mpmd_mesh = make_pipeline_mesh(
				190	pipeline_config,
				191	expert_axis_size=config.expert_axis_size,
				192	replica_axis_size=replica_axis_size,
				193	)
				194	optimizer = optax.adamw(learning_rate=1e-4, b1=0.9, b2=0.95, weight_decay=0.1)
				195	peak_flops_per_device = device_flops("h100")
				196	if peak_flops_per_device is None:
				197	raise ValueError("Fray does not define H100 BF16 peak FLOP/s")
				198
				199	_log(
				200	"PIPELINE_CONFIG",
				201	process_index=jax.process_index(),
				202	process_count=jax.process_count(),
				203	global_devices=jax.device_count(),
				204	local_devices=jax.local_device_count(),
				205	mesh_shape=dict(mesh.shape),
				206	stages=config.stages,
				207	physical_stages=config.physical_stages,
				208	microbatches=config.microbatches,
				209	batch_size=config.batch_size,
				210	seq_len=config.seq_len,
				211	hidden_dim=config.hidden_dim,
				212	intermediate_dim=model_config.intermediate_dim,
				213	shared_expert_intermediate_dim=model_config.shared_expert_intermediate_dim,
				214	num_layers=config.num_layers,
				215	layers_per_stage=config.layer_counts,
				216	num_experts=config.num_experts,
				217	top_k=config.top_k,
				218	vocab_size=model_config.vocab_size,
				219	num_heads=model_config.num_heads,
				220	num_kv_heads=model_config.num_kv_heads,
				221	sliding_window=model_config.sliding_window,
				222	qk_mult=model_config.qk_mult,
				223	steps=config.steps,
				224	profile_start_step=config.profile_start_step,
				225	profile_steps=config.profile_steps,
				226	mp_policy=config.mp_policy_string,
				227	remat_mode=config.remat_mode,
				228	schedule=config.schedule,
				229	jax_version=jax.__version__,
				230	)
				231
				232	init_started = time.monotonic()
				233	with jax.set_mesh(mesh):
				234	model = mp_policy.cast_to_param(Transformer.init(model_config, key=jax.random.PRNGKey(0)))
				235	batch_sharding = NamedSharding(mesh, P(BATCH_AXES, None))
				236	token_row = np.arange(config.seq_len, dtype=np.int32) % model_config.vocab_size
				237	host_tokens = np.broadcast_to(token_row, (config.batch_size, config.seq_len)).copy()
				238	host_loss_weight = np.ones((config.batch_size, config.seq_len), dtype=np.float32)
				239	host_loss_weight[:, -1] = 0
				240	batch = GrugLmExample(
				241	tokens=jax.device_put(host_tokens, batch_sharding),
				242	loss_weight=jax.device_put(host_loss_weight, batch_sharding),
				243	)
				244	loss_denominator = jnp.sum(batch.loss_weight.astype(jnp.float32))
				245	batches = reshape_batch_into_microbatches(batch, config.microbatches)
				246	if config.schedule == PipelineSchedule.DUALPIPE_V:
				247	stage_to_mpmd_index = automatic_stage_to_mpmd_indices(pipeline_config, config.schedule.automatic_schedule)
				248	else:
				249	stage_to_mpmd_index = None
				250	state, static_stages = initialize_mpmd_automatic_pipeline_state(
				251	model,
				252	optimizer,
				253	mpmd_mesh,
				254	num_stages=config.stages,
				255	layer_counts=config.layer_counts,
				256	stage_to_mpmd_index=stage_to_mpmd_index,
				257	memory_threshold=config.memory_threshold,
				258	)
				259	jax.block_until_ready((state, batches, loss_denominator))
				260	_log(
				261	"PIPELINE_INIT",
				262	process_index=jax.process_index(),
				263	elapsed_seconds=time.monotonic() - init_started,
				264	)
				265
				266	build_started = time.monotonic()
				267	step = make_automatic_pipeline_step(
				268	optimizer,
				269	mp_policy,
				270	static_stages,
				271	state,
				272	batches,
				273	config=pipeline_config,
				274	mpmd_mesh=mpmd_mesh,
				275	schedule_name=config.schedule.automatic_schedule,
				276	)
				277	_log(
				278	"PIPELINE_BUILD",
				279	process_index=jax.process_index(),
				280	elapsed_seconds=time.monotonic() - build_started,
				281	)
				282	lower_started = time.monotonic()
				283	prepared = prepare_automatic_mpmd_step(
				284	step,
				285	state,
				286	batches,
				287	loss_denominator,
				288	mpmd_mesh,
				289	memory_threshold=config.memory_threshold,
				290	)
				291	step = prepared.step
				292	state = prepared.state
				293	batches = prepared.batches
				294	loss_denominator = prepared.loss_denominator
				295	_log(
				296	"PIPELINE_LOWER",
				297	process_index=jax.process_index(),
				298	elapsed_seconds=time.monotonic() - lower_started,
				299	)
				300
				301	profiler_callback = None
				302	if config.profile_steps > 0:
				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
	202	flops_per_token = lm_flops_per_token(		354	forward_flops_per_token = lm_flops_per_token(
	203	hidden_dim=model_config.hidden_dim,		355	hidden_dim=model_config.hidden_dim,
	204	intermediate_dim=model_config.intermediate_dim,		356	intermediate_dim=model_config.intermediate_dim,
n			n	357	shared_intermediate_dim=model_config.shared_expert_intermediate_dim,
	205	num_layers=model_config.num_layers,		358	num_layers=model_config.num_layers,
	206	num_kv_heads=model_config.num_kv_heads,		359	num_kv_heads=model_config.num_kv_heads,
	207	num_heads=model_config.num_heads,		360	num_heads=model_config.num_heads,
	208	seq_len=model_config.max_seq_len,		361	seq_len=model_config.max_seq_len,
	209	vocab_size=model_config.vocab_size,		362	vocab_size=model_config.vocab_size,
n	210	glu=False,	n	363	glu=True,
	211	)		364	num_experts=model_config.num_experts,
	212	flops_per_example = 3 * flops_per_token * model_config.max_seq_len		365	num_shared_experts=1 if model_config.shared_expert_intermediate_dim > 0 else 0,
	213			366	num_experts_per_tok=model_config.num_experts_per_token,
	214	flops_summary: dict[str, float] = {
	215	"throughput/flops_per_token_analytic": flops_per_token,
	216	"throughput/flops_per_example_analytic": flops_per_example,
	217	}
	218
	219	return flops_per_example, flops_summary
	220
	221
	222	def _make_mixture_stage_callback(train_dataset: MixtureDataset, batch_schedule: BatchSchedule):
	223	last_mixture_stage = -1
	224
	225	def log_mixture_stage(step_info):
	226	nonlocal last_mixture_stage
	227	seq_index = batch_schedule.global_data_offset_by_step(step_info.step)
	228	block_id = seq_index // train_dataset.block_size
	229	stage = train_dataset._get_stage_for_block(block_id)
	230	if stage == last_mixture_stage:
	231	return
	232
	233	weights = train_dataset.weight_stages[stage][1]
	234	mixture_log = {f"mixture/weight/{name}": weight for name, weight in weights.items()}
	235	mixture_log["mixture/stage"] = stage
	236	levanter.tracker.log(mixture_log, step=step_info.step)
	237	last_mixture_stage = stage
	238
	239	return log_mixture_stage
	240
	241
	242	def _build_backward_flow_graph(
	243	*,
	244	params: Transformer,
	245	batch: GrugLmExample,
	246	mp: jmp.Policy,
	247	z_loss_weight: float,
	248	) -> BackwardFlowGraph:
	249	z_loss = z_loss_weight if z_loss_weight > 0 else None
	250
	251	def loss_fn(model: Transformer) -> jax.Array:
	252	compute_params = mp.cast_to_compute(model)
	253	return compute_params.next_token_loss(
	254	batch.tokens,
	255	batch.loss_weight,
	256	mask=batch.attn_mask,
	257	reduction="mean",
	258	logsumexp_weight=z_loss,
	259	)
	260
	261	with capture_backward_flow(BackwardFlowConfig(interval=1)):
	262	closed_jaxpr, _, _ = eqx.filter_make_jaxpr(jax.grad(loss_fn))(params)
	263	return backward_flow_graph_from_jaxpr(closed_jaxpr)
	264
	265
	266	def _write_backward_flow_artifact(
	267	*,
	268	graph: BackwardFlowGraph,
	269	backward_flow_metrics: dict[str, jax.Array],
	270	log_dir,
	271	run_id: str,
	272	step: int,
	273	) -> None:
	274	node_stats = backward_flow_node_stats(backward_flow_metrics)
	275	collapsed_graph = collapse_backward_flow_graph(graph, node_stats)
	276	if not collapsed_graph.nodes:
	277	return
	278
	279	render_hints = infer_backward_flow_render_hints(collapsed_graph)
	280	html = render_backward_flow_html(
	281	collapsed_graph,
	282	node_stats,
	283	node_lanes=render_hints.node_lanes,
	284	display_edges=render_hints.display_edges,
	285	plates=render_hints.plates,
	286	title=f"Backward Flow Step {step}",
	287	)
	288	artifact_dir = log_dir / run_id / "artifacts" / "backward_flow"
	289	artifact_dir.mkdir(parents=True, exist_ok=True)
	290	artifact_path = artifact_dir / f"step_{step:07d}.html"
	291	artifact_path.write_text(html)
	292	tracker = levanter.tracker.current_tracker()
	293	tracker.log_html("backward_flow/dag", artifact_path, step=step, commit=False)
	294	tracker.log_artifact(
	295	artifact_path,
	296	name=f"backward_flow_step_{step:07d}",
	297	type="backward_flow",
	298	)
	299
	300
	301	@register_dataclass
	302	@dataclass(frozen=True)
	303	class GrugTrainState:
	304	step: jax.Array
	305	params: Transformer
	306	opt_state: optax.OptState
	307	ema_params: Transformer \| None
	308
	309
	310	def initial_state(
	311	model_config: GrugModelConfig,
	312	*,
	313	optimizer: optax.GradientTransformation,
	314	mp: jmp.Policy,
	315	key: PRNGKeyArray,
	316	ema_beta: float \| None,
	317	) -> GrugTrainState:
	318	params = mp.cast_to_param(Transformer.init(model_config, key=key))
	319	return GrugTrainState(
	320	step=jnp.array(0, dtype=jnp.int32),
	321	params=params,
	322	opt_state=optimizer.init(params),
	323	ema_params=params if ema_beta is not None else None,
	324	)
	325
	326
	327	def _make_train_step(
	328	optimizer: optax.GradientTransformation,
	329	mp: jmp.Policy,
	330	*,
	331	z_loss_weight: float,
	332	loss_implementation: str \| tuple[str, ...] \| None = None,
	333	ema_beta: float \| None,
	334	watch_config: WatchConfig \| None = None,
	335	backward_flow_config: BackwardFlowConfig \| None = None,
	336	):
	337	one = jnp.array(1, dtype=jnp.int32)
	338	z_loss = z_loss_weight if z_loss_weight > 0 else None
	339	if watch_config is not None:
	340	if isinstance(watch_config.watch_targets, str):
	341	watch_targets = tuple(t.strip() for t in watch_config.watch_targets.split(","))
	342	else:
	343	watch_targets = tuple(watch_config.watch_targets)
	344	else:
	345	watch_targets = ()
	346
	347	@functools.partial(jax.jit, donate_argnums=(0,), static_argnames=("compute_backward_flow", "compute_watch"))
	348	def train_step(
	349	state: GrugTrainState,
	350	batch,
	351	*,
	352	compute_watch: bool = False,
	353	compute_backward_flow: bool = False,
	354	):		367	)
	355	def loss_fn(params):		368	analytic_mfu = 3 * forward_flops_per_token * tokens_per_second / (jax.device_count() * peak_flops_per_device)
	356	compute_params = mp.cast_to_compute(params)		369	median_analytic_mfu = (
	357	return compute_params.next_token_loss(		370	3 * forward_flops_per_token * median_tokens_per_second / (jax.device_count() * peak_flops_per_device)
	358	batch.tokens,		371	)
	359	batch.loss_weight,		372	if loss is not None:
	360	mask=batch.attn_mask,		373	_log(
	361	reduction="mean",		374	"PIPELINE_SUMMARY",
	362	logsumexp_weight=z_loss,		375	mean_step_seconds=mean_step_seconds,
	363	loss_implementation=loss_implementation,		376	median_step_seconds=median_step_seconds,
				377	tokens_per_second=tokens_per_second,
				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	)
				385
				386	for device in jax.local_devices():
				387	stats = device.memory_stats()
				388	if stats is not None:
				389	_log(
				390	"PIPELINE_MEMORY",
				391	process_index=jax.process_index(),
				392	device_id=device.id,
				393	peak_bytes_in_use=stats.get("peak_bytes_in_use"),
				394	bytes_limit=stats.get("bytes_limit"),
	364	)		395	)
	365			396
n	366	backward_flow_stats = None	n
	367	if backward_flow_config is not None and compute_backward_flow:
	368	gradient_scale = jnp.sum(jnp.asarray(batch.loss_weight, dtype=jnp.float32))
	369	with capture_backward_flow(backward_flow_config, gradient_scale=gradient_scale):
	370	with levanter.tracker.defer_tracker_for_jit() as backward_flow_stats:
	371	loss, grads = jax.value_and_grad(loss_fn)(state.params)
	372	else:
	373	loss, grads = jax.value_and_grad(loss_fn)(state.params)
	374	updates, opt_state = optimizer.update(grads, state.opt_state, state.params)
	375	params = optax.apply_updates(state.params, updates)
	376
	377	if ema_beta is None:
	378	ema_params = None
	379	else:
	380	if state.ema_params is None:
	381	raise ValueError("ema_params must be initialized when ema_beta is set.")
	382	ema_params = jax.tree_util.tree_map(
	383	lambda old, new: ema_beta * old + (1.0 - ema_beta) * new,
	384	state.ema_params,
	385	params,
	386	)
	387
	388	watch_stats = None
	389	if watch_config is not None and compute_watch:
	390	watch_stats = compute_watch_stats(
	391	watch_targets=watch_targets,
	392	include_norms=watch_config.include_norms,
	393	include_per_parameter_norms=watch_config.include_per_parameter_norms,
	394	include_histogram=watch_config.include_histograms,
	395	split_scan_layers=watch_config.split_scan_layers,
	396	params=state.params,
	397	grads=grads,
	398	updates=updates,
	399	opt_state=state.opt_state,
	400	model_tree_type=type(state.params),
	401	)
	402
	403	next_state = dataclasses.replace(
	404	state,
	405	step=state.step + one,
	406	params=params,
	407	opt_state=opt_state,
	408	ema_params=ema_params,
	409	)
	410
	411	metrics = {"train/loss": loss}
	412	if backward_flow_stats is not None:
	413	metrics[_BACKWARD_FLOW_METRICS_KEY] = backward_flow_stats
	414
	415	return next_state, metrics, watch_stats
	416
	417	return train_step
	418
	419
	420	def _run_grug_local(config: GrugRunConfig) -> None:
	421	"""Entry point for the grug template training loop."""
	422	trainer = config.trainer.trainer
	423	trainer.initialize()
	424	levanter.tracker.log_configuration(config)
	425
	426	run_id = trainer.id
	427	if run_id is None:
	428	raise ValueError("trainer.id was not initialized")
	429
	430	optimizer = config.optimizer.build(trainer.num_train_steps)
	431	watch_config = trainer.watch
	432	backward_flow_config = config.trainer.backward_flow if config.trainer.backward_flow.is_enabled else None
	433	train_step = _make_train_step(
	434	optimizer,
	435	trainer.mp,
	436	z_loss_weight=config.trainer.z_loss_weight,
	437	loss_implementation=config.trainer.loss_implementation,
	438	ema_beta=config.trainer.ema_beta,
	439	watch_config=watch_config if watch_config.is_enabled else None,
	440	backward_flow_config=backward_flow_config,
	441	)
	442
	443	data_key, model_key = jax.random.split(jax.random.PRNGKey(trainer.seed), 2)
	444	if config.trainer.data_seed is not None:
	445	data_key = jax.random.PRNGKey(config.trainer.data_seed)
	446
	447	# Grug uses raw PartitionSpecs rather than Trainer's logical axis mapping.
	448	# Keep the mesh compact so P(("replica_dcn", "data")) spans slices directly.
	449	mesh = compact_grug_mesh()
	450	checkpointer = trainer.checkpointer.create(run_id)
	451	with set_mesh(mesh), TrainingDashboard(config, checkpointer.request_checkpoint, run_id):
	452	batch_schedule = trainer.batch_schedule
	453
	454	train_dataset = build_train_dataset(
	455	config.data,
	456	max_seq_len=config.model.max_seq_len,
	457	batch_schedule=batch_schedule,
	458	key=data_key,
	459	)
	460	train_loader = build_train_loader(
	461	train_dataset,
	462	batch_schedule=batch_schedule,
	463	mesh=mesh,
	464	batch_pspec=config.trainer.train_batch_pspec,
	465	)
	466
	467	@jax.jit
	468	def _init_state(model_rng):
	469	return initial_state(
	470	config.model,
	471	optimizer=optimizer,
	472	mp=trainer.mp,
	473	key=model_rng,
	474	ema_beta=config.trainer.ema_beta,
	475	)
	476
	477	state = _init_state(model_key)
	478
	479	state = restore_grug_state_from_checkpoint(
	480	state,
	481	checkpoint_search_paths=trainer.checkpoint_search_paths(run_id),
	482	load_checkpoint_setting=trainer.load_checkpoint,
	483	mesh=mesh,
	484	allow_partial=trainer.allow_partial_checkpoint,
	485	)
	486	if trainer.initialize_from is not None:
	487	state = init_weights_only_from_checkpoint(
	488	state,
	489	trainer.initialize_from,
	490	mesh=mesh,
	491	allow_partial=trainer.allow_partial_checkpoint,
	492	)
	493	dump_grug_state_sharding_run_artifact(
	494	state,
	495	log_dir=trainer.log_dir,
	496	run_id=run_id,
	497	path_override=config.trainer.sharding_dump_path,
	498	)
	499
	500	levanter.tracker.log_summary({"parameter_count": parameter_count(state.params)})
	501
	502	flops_per_example, flops_summary = _compute_flops(model_config=config.model)
	503	levanter.tracker.log_summary(flops_summary)
	504
	505	eval_cfg = config.eval
	506	evaluator = None
	507	if eval_cfg is not None:
	508	evaluator = build_tagged_evaluator(
	509	data_config=config.data,
	510	max_seq_len=config.model.max_seq_len,
	511	mesh=mesh,
	512	eval_cfg=eval_cfg,
	513	)
	514
	515	profiler_cfg = trainer.profiler
	516	profiler_num_steps = profiler_cfg.resolve_num_profile_steps(num_train_steps=trainer.num_train_steps)
	517	profiler_enabled = profiler_cfg.is_enabled and profiler_num_steps > 0
	518
	519	log_every = max(1, config.trainer.log_every)
	520	iterator = LoadingTimeTrackerIterator(train_loader.iter_from_step(int(state.step)))
	521
	522	state_callbacks = StateCallbackRunner[GrugTrainState](
	523	step_getter=lambda s: s.step,
	524	model_getter=lambda s: s.params,
	525	eval_model_getter=lambda s: s.ema_params if s.ema_params is not None else s.params,
	526	opt_state_getter=lambda s: s.opt_state,
	527	)
	528	state_callbacks.add_hook(
	529	callbacks.log_performance_stats(config.model.max_seq_len, batch_schedule, flops_per_example),
	530	every=log_every,
	531	)
	532	state_callbacks.add_hook(callbacks.pbar_logger(total=trainer.num_train_steps), every=log_every)
	533	state_callbacks.add_hook(callbacks.log_step_info(trainer.num_train_steps), every=log_every)
	534	if profiler_enabled:
	535	state_callbacks.add_hook(
	536	profiler_cfg.build(
	537	str(trainer.log_dir / run_id / "profiler"),
	538	run_id=run_id,
	539	num_steps=profiler_num_steps,
	540	),
	541	every=1,
	542	)
	543	state_callbacks.add_hook(_make_mixture_stage_callback(train_dataset, batch_schedule), every=1)
	544	if evaluator is not None and eval_cfg is not None:
	545	interval = eval_cfg.steps_per_eval
	546	eval_ema = eval_cfg.eval_ema and config.trainer.ema_beta is not None
	547	if interval is not None and interval > 0 and (eval_cfg.eval_current or eval_ema):
	548	state_callbacks.add_hook(
	549	cb_tagged_evaluate(
	550	evaluator,
	551	prefix=eval_cfg.prefix,
	552	eval_current=eval_cfg.eval_current,
	553	eval_ema=eval_ema,
	554	),
	555	every=interval,
	556	)
	557
	558	last_loss: float \| jax.Array = 0.0
	559	last_step_duration = 0.0
	560	backward_flow_graph: BackwardFlowGraph \| None = None
	561
	562	# Main optimization loop.
	563	try:
	564	while int(state.step) < trainer.num_train_steps:
	565	with jax.profiler.TraceAnnotation("load_batch"):
	566	batch = next(iterator)
	567	step_start = time.perf_counter()
	568	current_step = int(state.step)
	569	# grad_watch runs only on its configured interval.
	570	compute_watch = (
	571	watch_config.is_enabled and watch_config.interval > 0 and current_step % watch_config.interval == 0
	572	)
	573	compute_backward_flow = (
	574	backward_flow_config is not None
	575	and backward_flow_config.interval > 0
	576	and current_step % backward_flow_config.interval == 0
	577	)
	578	state, metrics, watch_stats = train_step(
	579	state,
	580	batch,
	581	compute_watch=compute_watch,
	582	compute_backward_flow=compute_backward_flow,
	583	)
	584	backward_flow_stats = metrics.get(_BACKWARD_FLOW_METRICS_KEY)
	585	if backward_flow_stats is not None:
	586	metrics = {key: value for key, value in metrics.items() if key != _BACKWARD_FLOW_METRICS_KEY}
	587	step = int(state.step) - 1
	588
	589	jax.block_until_ready(metrics["train/loss"])
	590	duration = time.perf_counter() - step_start
	591	backward_flow_timing_metrics = None
	592	if compute_backward_flow:
	593	backward_flow_timing_metrics = {"backward_flow/compute_step_duration": duration}
	594
	595	hook_start = time.perf_counter()
	596	with jax.profiler.TraceAnnotation("callbacks"):
	597	state_callbacks.run(state, loss=metrics["train/loss"], step_duration=duration)
	598	last_loss = metrics["train/loss"]
	599	last_step_duration = duration
	600	levanter.tracker.log({"throughput/hook_time": time.perf_counter() - hook_start}, step=step)
	601	levanter.tracker.log({"throughput/loading_time": iterator.this_load_time}, step=step)
	602
	603	if watch_stats is not None:
	604	levanter.tracker.log(watch_stats, step=step)
	605	if backward_flow_timing_metrics is not None:
	606	levanter.tracker.log(backward_flow_timing_metrics, step=step)
	607	if backward_flow_stats:
	608	levanter.tracker.log(backward_flow_stats, step=step)
	609	if backward_flow_graph is None:
	610	backward_flow_graph = _build_backward_flow_graph(
	611	params=state.params,
	612	batch=batch,
	613	mp=trainer.mp,
	614	z_loss_weight=config.trainer.z_loss_weight,
	615	)
	616	artifact_start = time.perf_counter()
	617	_write_backward_flow_artifact(
	618	graph=backward_flow_graph,
	619	backward_flow_metrics=backward_flow_stats,
	620	log_dir=trainer.log_dir,
	621	run_id=run_id,
	622	step=step,
	623	)
	624	levanter.tracker.log(
	625	{"backward_flow/artifact_write_duration": time.perf_counter() - artifact_start},
	626	step=step,
	627	)
	628
	629	if checkpointer is not None:
	630	checkpointer.on_step(tree=state, step=int(state.step))
	631	except BaseException:
	632	logger.exception(
	633	"Fatal error in grug training loop; skipping final callbacks/checkpoint to preserve root cause"
	634	)
	635	raise
	636	else:
	637	# Mirror classic trainer behavior: force callbacks on the last completed step.
	638	state_callbacks.run(state, loss=last_loss, step_duration=last_step_duration, force=True)
	639	if checkpointer is not None:
	640	checkpointer.on_step(tree=state, step=int(state.step), force=True)
	641	checkpointer.wait_until_finished()
	642
	643	levanter.tracker.current_tracker().finish()
	644
	645			397
	646	def run_grug(config: GrugRunConfig) -> None:		398	def run_grug(config: GrugRunConfig) -> None:
n	647	"""Dispatch grug training through Fray jobs."""	n	399	"""Dispatch pipeline-parallel Grug training through Fray."""
	648	trainer = config.trainer.trainer		400	if config.resources.replicas % config.train.physical_stages != 0:
	649	if trainer.id is None:		401	raise ValueError(
	650	raise ValueError("trainer.id must be set before dispatching grug training.")		402	f"training resource replicas ({config.resources.replicas}) must be divisible by "
	651			403	f"{config.train.physical_stages} physical stages"
				404	)
	652	dispatch_grug_training_run(		405	dispatch_grug_training_run(
n	653	run_id=trainer.id,	n	406	run_id=config.run_id,
	654	config=config,		407	config=config.train,
	655	local_entrypoint=_run_grug_local,		408	local_entrypoint=_run_grug_local,
	656	resources=config.resources,		409	resources=config.resources,
n			n	410	extras=["pipeline"],
	657	)		411	)
t	658		t
	659
	660	__all__ = [
	661	"GrugEvalConfig",
	662	"GrugRunConfig",
	663	"GrugTrainState",
	664	"GrugTrainerConfig",
	665	"initial_state",
	666	"run_grug",
	667	]

README.md

added

+# grug-moe-pipeline
+Canonical pipeline-parallel Grug MoE implementation. This directory is a copy-paste
+variant of [`experiments/grug/moe`](../moe/): its model and execution path are local
+so pipeline-specific changes do not turn the ordinary MoE implementation into a
+shared trainer framework.
+[`pipeline.py`](./pipeline.py) splits the local transformer into stage pytrees and
+builds automatic JaxPP ZeroBubble or DualPipeV optimizer steps. [`train.py`](./train.py)
+owns model initialization, the training loop, and the Fray dispatch entry point.
+[`benchmark.py`](./benchmark.py) maps environment variables onto that trainer for
+repeatable performance runs. One logical stage maps to each physical pipeline rank with
+`PIPELINE_SCHEDULE=automatic_zero_bubble`. Set `PIPELINE_SCHEDULE=automatic_dualpipe_v`,
+`PIPELINE_STAGES=2P`, and `PIPELINE_PHYSICAL_STAGES=P` to fold two logical stages onto
+each physical rank in JaxPP's V-shaped placement. DualPipeV requires at least `2P`
+microbatches. `PIPELINE_LAYERS_PER_STAGE` contains one positive layer count per logical
+stage.
+The canonical loop initializes fresh parameters, uses a fixed AdamW optimizer, and
+trains on synthetic token rows. It does not save or restore checkpoints. A model-specific
+training variant should replace those parts while retaining the stage split and JaxPP
+orchestration that it needs.
+The best validated Snowball 67B-A2B throughput point uses eight H100x8 replicas,
+sixteen logical stages, batch 256, 32 microbatches, sequence length 8192, and layer
+counts `1,2,2,2,2,2,2,2,1,1,1,1,2,2,2,1`:
+```bash
+export XLA_FLAGS='--xla_gpu_executable_terminate_timeout=300 --xla_gpu_enable_command_buffer='
+export PIPELINE_SCHEDULE=automatic_dualpipe_v
+export PIPELINE_PHYSICAL_STAGES=8
+export PIPELINE_STAGES=16
+export PIPELINE_BATCH=256
+export PIPELINE_MICROBATCHES=32
+export PIPELINE_SEQ_LEN=8192
+export PIPELINE_HIDDEN_DIM=2560
+export PIPELINE_INTERMEDIATE_DIM=1280
+export PIPELINE_SHARED_EXPERT_INTERMEDIATE_DIM=2560
+export PIPELINE_LAYERS=26
+export PIPELINE_LAYERS_PER_STAGE=1,2,2,2,2,2,2,2,1,1,1,1,2,2,2,1
+export PIPELINE_EXPERTS=256
+export PIPELINE_TOP_K=4
+export PIPELINE_EXPERT_AXIS=8
+export PIPELINE_HEADS=20
+export PIPELINE_KV_HEADS=5
+export PIPELINE_VOCAB_SIZE=128256
+export PIPELINE_SLIDING_WINDOW=2048
+export PIPELINE_QK_MULT=1.5703
+export PIPELINE_MP=params=bfloat16,compute=bfloat16,output=bfloat16
+export PIPELINE_ATTENTION=gpu_fa4_cute
+export PIPELINE_MOE=ring
+export PIPELINE_REMAT=recompute_all
+uv run --extra pipeline python -m experiments.grug.moe_pipeline.benchmark
+```
+A 14-step run with four warmups measured 3.476 s median step time, 603,327 tokens/s,
+.46% analytic MFU, and 59.153 GB peak memory per device. A matched batch-128,
+-microbatch control measured 543,490 tokens/s and 17.53% analytic MFU. Both used
+eight examples per microbatch. The larger batch improved median throughput by 11.0%
+without increasing peak memory; final loss differed by 4.0e-5 after 14 synthetic-data
+steps.

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,
+    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, *BATCH_AXES, "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,
+                use_pko=is_long and not cfg.disable_pko,
+                disable_rope=is_long and cfg.disable_long_rope,
+            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."""
+    trainable_params: tuple[GrugMoePipelineStage, ...]
+    opt_state: tuple[optax.OptState, ...]
+    pending_qb_betas: tuple[jax.Array, ...]
+@dataclass(frozen=True)
+class _InitializedMpmdStageState:
+    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 are removed from
+        both 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)
+    return _InitializedMpmdStageState(
+        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(
+            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,
+            is_leaf=lambda value: value is None,
+        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,
+            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 the remaining SPMD 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,

launch.py

removed

	experiments/grug/base/launch.py
t	1	# Copyright The Marin Authors	t
	2	# SPDX-License-Identifier: Apache-2.0
	3
	4	"""grug-base 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, and evals
	8	are stated inline; the output path is ``ctx.output_path`` and the TPU is a run-arg, so neither
	9	bears on the artifact's identity.
	10
	11	The grug training mechanism (``GrugBaseLaunchConfig`` + ``build_grug_run_config`` ->
	12	``run_grug``) is grug-specific compute and is kept as-is: the recipe's ``fn`` builds the
	13	Levanter trainer (its checkpointer paths derive from ``output_path``) and dispatches the
	14	training job to Fray. Only the data/validation wiring around it is assembled lazily.
	15	"""
	16
	17	import dataclasses
	18	import os
	19	from dataclasses import dataclass
	20	from datetime import timedelta
	21
	22	import jmp
	23	from fray.cluster import ResourceConfig
	24	from levanter.callbacks.profiler import ProfilerConfig
	25	from levanter.checkpoint import CheckpointerConfig
	26	from levanter.data.text.datasets import LmDataConfig
	27	from levanter.optim.config import AdamConfig, OptimizerConfig
	28	from levanter.tracker import TrackerConfig
	29	from levanter.tracker.wandb import WandbConfig
	30	from levanter.trainer import TrainerConfig
	31	from marin.execution.build_context import resolve_version
	32	from marin.execution.lazy import ArtifactStep, StepContext
	33	from marin.experiment.cli import experiment_main
	34	from marin.experiment.data import mixture
	35	from marin.experiment.namespacing import user_namespaced_name
	36	from marin.training.training import LevanterCheckpoint, resolve_checkpointer_output_path
	37
	38	from experiments.datasets.nemotron import nemotron_datasets
	39	from experiments.datasets.paloma import paloma_datasets
	40	from experiments.datasets.proofpile import proofpile_dataset
	41	from experiments.datasets.starcoder import starcoder_dataset
	42	from experiments.datasets.uncheatable import uncheatable_datasets
	43	from experiments.grug.base.model import GrugModelConfig
	44	from experiments.grug.base.train import GrugEvalConfig, GrugRunConfig, GrugTrainerConfig, run_grug
	45	from experiments.llama import llama3_tokenizer
	46
	47	# The TPU the training job is dispatched onto. A run-arg, not part of the config's
	48	# identity: re-running on a different TPU is the same checkpoint. The launcher step
	49	# runs inline (run_grug dispatches its own Fray job).
	50	_TRAIN_RESOURCES = ResourceConfig.with_tpu("v5p-8")
	51
	52	# Nemotron CC mixture weights: the corpus's TiB proportions, plus starcoder and
	53	# proof-pile at their published weights. Policy lives here, in the experiment.
	54	_NEMOTRON_WEIGHTS = {
	55	"hq_actual": 0.91351,
	56	"hq_synth": 2.72,
	57	"medium_high": 0.82471,
	58	"medium": 3.38,
	59	"medium_low": 1.54,
	60	"low_actual": 0.70123,
	61	"low_synth": 0.62771,
	62	}
	63	_STARCODER_WEIGHT = 0.25
	64	_PROOFPILE_WEIGHT = 0.055
	65
	66
	67	GRUG_130M_MODEL = GrugModelConfig(
	68	vocab_size=128_256,
	69	hidden_dim=512,
	70	intermediate_dim=1792,
	71	num_layers=6,
	72	num_heads=8,
	73	num_kv_heads=8,
	74	max_seq_len=4096,
	75	head_dim=None,
	76	)
	77
	78
	79	@dataclass(frozen=True)
	80	class GrugBaseLaunchConfig:
	81	"""Last-mile run config for the base grug template.
	82
	83	Keep this as the main entry point for day-to-day edits (model/data/optimizer/trainer/eval knobs).
	84
	85	The trainer and eval knobs are flat scalars rather than nested
	86	``GrugTrainerConfig`` / ``GrugEvalConfig`` objects: those carry
	87	``jax.sharding.PartitionSpec`` batch-sharding defaults, which the artifact
	88	fingerprint cannot serialize. ``build_grug_run_config`` reconstitutes them (with
	89	their default pspecs) at run time, keeping the sharding plumbing out of the
	90	artifact's identity.
	91	"""
	92
	93	model: GrugModelConfig
	94	data: LmDataConfig
	95	output_path: str
	96	run_id: str
	97	resources: ResourceConfig
	98	steps: int
	99	batch_size: int
	100	seed: int
	101	mp: str # jmp policy string, e.g. "params=float32,compute=bfloat16,output=bfloat16".
	102	tracker: TrackerConfig
	103	optimizer: OptimizerConfig
	104	z_loss_weight: float = 1e-4
	105	ema_beta: float \| None = None # EMA coefficient for eval/checkpoint model; None disables EMA.
	106	log_every: int = 1
	107	loss_implementation: str \| tuple[str, ...] \| None = None # cross-entropy kernel; None uses the trainer default.
	108	eval_batch_size: int \| None = 512 # None disables perplexity eval.
	109	steps_per_eval: int = 1000
	110	max_eval_batches: int = 8
	111	eval_current: bool = True
	112	eval_ema: bool = False
	113
	114
	115	def _resolve_run_id(default_run_id: str) -> str:
	116	"""Resolve run id and append `FERRY_DATE` when launching from ferry workflows."""
	117	run_id = os.environ.get("GRUG_RUN_ID", default_run_id)
	118	ferry_date = os.environ.get("FERRY_DATE")
	119	if ferry_date:
	120	run_id = f"{run_id}-{ferry_date}"
	121	return run_id
	122
	123
	124	def _resolve_tracker(tracker: TrackerConfig, run_id: str) -> TrackerConfig:
	125	if isinstance(tracker, WandbConfig):
	126	return dataclasses.replace(tracker, name=run_id)
	127	return tracker
	128
	129
	130	def build_grug_run_config(launch: GrugBaseLaunchConfig) -> GrugRunConfig:
	131	"""Map launch knobs onto the trainer's full ``GrugRunConfig``.
	132
	133	The checkpointer's ``base_path`` and ``temporary_base_path`` derive from
	134	``launch.output_path``, so a run resolves its checkpoint locations under the
	135	region its output path lives in.
	136	"""
	137	output_path = launch.output_path
	138	trainer = TrainerConfig(
	139	id=launch.run_id,
	140	seed=launch.seed,
	141	train_batch_size=launch.batch_size,
	142	num_train_steps=launch.steps,
	143	profiler=ProfilerConfig(enabled=False, start_step=5, num_steps=100, perfetto_link=False),
	144	mp=jmp.get_policy(launch.mp),
	145	tracker=_resolve_tracker(launch.tracker, launch.run_id),
	146	use_explicit_mesh_axes=True,
	147	require_accelerator=True,
	148	allow_nondivisible_batch_size=False,
	149	checkpointer=resolve_checkpointer_output_path(
	150	CheckpointerConfig(save_interval=timedelta(minutes=10), keep=None),
	151	output_path,
	152	),
	153	)
	154
	155	grug_trainer = GrugTrainerConfig(
	156	trainer=trainer,
	157	z_loss_weight=launch.z_loss_weight,
	158	ema_beta=launch.ema_beta,
	159	log_every=launch.log_every,
	160	loss_implementation=launch.loss_implementation,
	161	)
	162
	163	eval_config = (
	164	GrugEvalConfig(
	165	eval_batch_size=launch.eval_batch_size,
	166	steps_per_eval=launch.steps_per_eval,
	167	max_eval_batches=launch.max_eval_batches,
	168	eval_current=launch.eval_current,
	169	eval_ema=launch.eval_ema,
	170	)
	171	if launch.eval_batch_size is not None
	172	else None
	173	)
	174
	175	return GrugRunConfig(
	176	model=launch.model,
	177	data=launch.data,
	178	resources=launch.resources,
	179	output_path=output_path,
	180	optimizer=launch.optimizer,
	181	trainer=grug_trainer,
	182	eval=eval_config,
	183	)
	184
	185
	186	def run_grug_base_trial(config: GrugBaseLaunchConfig) -> None:
	187	"""Build the full grug run config and dispatch the training job to Fray.
	188
	189	Runs inline on the launcher; ``run_grug`` submits the job and blocks until it
	190	completes.
	191	"""
	192	run_grug(build_grug_run_config(config))
	193
	194
	195	def grug_base_trial(*, version: str \| None = None) -> ArtifactStep[LevanterCheckpoint]:
	196	"""The base grug trial on the Nemotron mix as a lazy checkpoint.
	197
	198	Every component is a :class:`Dataset` handle, so the whole graph lowers via
	199	:func:`~marin.execution.lazy.lower`. The paloma/uncheatable suites are validation
	200	(weight 0).
	201	"""
	202	name = "grug/base-trial"
	203	version = resolve_version(name, version)
	204	nem = nemotron_datasets(tokenizer=llama3_tokenizer)
	205	train = {nem[split]: weight for split, weight in _NEMOTRON_WEIGHTS.items()}
	206	train[starcoder_dataset(tokenizer=llama3_tokenizer)] = _STARCODER_WEIGHT
	207	train[proofpile_dataset(tokenizer=llama3_tokenizer)] = _PROOFPILE_WEIGHT
	208	validation = [
	209	*paloma_datasets(tokenizer=llama3_tokenizer).values(),
	210	*uncheatable_datasets(tokenizer=llama3_tokenizer).values(),
	211	]
	212
	213	run_id = _resolve_run_id("grug-base-trial")
	214
	215	def build_config(ctx: StepContext) -> GrugBaseLaunchConfig:
	216	return GrugBaseLaunchConfig(
	217	model=GRUG_130M_MODEL,
	218	data=mixture(ctx, train, validation=validation),
	219	output_path=ctx.output_path,
	220	run_id=run_id,
	221	resources=ctx.runtime_arg("train_resources"),
	222	steps=2_000,
	223	batch_size=512,
	224	seed=0,
	225	mp="params=float32,compute=bfloat16,output=bfloat16",
	226	tracker=WandbConfig(
	227	project="marin",
	228	tags=["grug", "template"],
	229	group="grug-base-trial",
	230	name=None, # filled from run_id in _resolve_tracker
	231	replicate_path=ctx.output_path,
	232	),
	233	optimizer=AdamConfig(
	234	learning_rate=3e-3,
	235	weight_decay=0.1,
	236	lr_schedule="cosine",
	237	decay=0.2,
	238	min_lr_ratio=0.1,
	239	warmup=1000,
	240	),
	241	z_loss_weight=1e-4,
	242	ema_beta=None,
	243	log_every=1,
	244	eval_batch_size=512,
	245	steps_per_eval=1000,
	246	max_eval_batches=8,
	247	eval_current=True,
	248	eval_ema=False,
	249	)
	250
	251	return ArtifactStep(
	252	name=user_namespaced_name(name, version),
	253	version=version,
	254	artifact_type=LevanterCheckpoint,
	255	run=run_grug_base_trial,
	256	build_config=build_config,
	257	deps=(train, validation),
	258	runtime_args={"train_resources": _TRAIN_RESOURCES},
	259	)
	260
	261
	262	if __name__ == "__main__":
	263	experiment_main(grug_base_trial)()