From 2ce8143eead19ec8a2fcc0c4288fb4b9a5bb27d9 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 13 Oct 2026 14:22:33 +0200 Subject: [PATCH 001/132] feat(trainer): random_init_mode chooses between skipping and initializing weights model.debug.random_init keeps its meaning (don't download or load weights). The new model.debug.random_init_mode decides what the weights hold instead: 'skip' leaves them uninitialized (the previous behavior, fastest for benchmarks), 'init' runs model.init. 'auto' resolves per entrypoint: skip for sft and rl, init for pretrain. The init draws each parameter's local shard from per-block seeds, so the values depend only on the seed and the canonical parameter name, not on the world size or sharding. --- .../src/prime_rl/configs/sft.py | 7 ++ .../src/prime_rl/configs/trainer.py | 28 ++++- src/prime_rl/trainer/init.py | 102 ++++++++++++++++++ src/prime_rl/trainer/model.py | 9 +- 4 files changed, 144 insertions(+), 2 deletions(-) create mode 100644 src/prime_rl/trainer/init.py diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index b6cb7accbdd..d2171d5580f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -293,6 +293,13 @@ def resolve_freeze_moe_router_auto(self): self.model.freeze_moe_router = False return self + @model_validator(mode="after") + def resolve_random_init_mode_auto(self): + """Resolve ``model.debug.random_init_mode='auto'``: SFT skips weight loading without initializing.""" + if self.model.debug.random_init_mode == "auto": + self.model.debug.random_init_mode = "skip" + return self + @model_validator(mode="after") def auto_setup_run_identity(self): """Auto-generate the run name (``----``) when unset and diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index f7035c4e949..63b2364e8db 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -100,7 +100,10 @@ class DebugModelConfig(BaseConfig): """Override the number of transformer layers (truncates the model).""" random_init: bool = False - """Randomly initialize the model instead of loading weights.""" + """Don't download or load weights from ``model.name``. ``random_init_mode`` sets what the weights hold instead.""" + + random_init_mode: Literal["auto", "skip", "init"] = "auto" + """What ``random_init`` does. ``skip`` leaves the weights uninitialized (arbitrary memory contents) - the fastest option, for throughput benchmarks and debugging. ``init`` initializes the weights as configured in ``model.init``. ``auto`` resolves to ``skip`` for ``sft`` and ``rl`` and to ``init`` for ``pretrain``.""" force_balanced_routing: bool = False """Replace MoE token-choice routing with a round-robin assignment so every expert sees an equal share. Intended for fake-data smoke tests where untrained routing would otherwise OOM under severe imbalance. Gating scores are still gathered from the override indices so the forward pass stays consistent.""" @@ -230,6 +233,19 @@ class MoERuntimeConfig(BaseConfig): dispatch: MoEDispatchConfig = TorchMoEDispatchConfig() +class InitConfig(BaseConfig): + """Random weight initialization, used when ``model.debug.random_init`` resolves to ``init``.""" + + scheme: Literal["normal", "scaled"] = "normal" + """``normal`` draws every weight matrix and the embedding from N(0, std). ``scaled`` additionally divides the std of the residual output projections (``o_proj``, ``down_proj``) by sqrt(2 * num_layers). Norm weights start at 1, biases at 0.""" + + std: float = Field(0.02, gt=0) + """Base standard deviation.""" + + seed: int = 0 + """Seed for the initialization. The values depend only on the seed, not on the number of GPUs.""" + + class ModelConfig(BaseModelConfig): conversion_dir: Path | None = None """Directory for the auto-converted weights (written to a `prime`/`hf` subdirectory). If not set, we write into the model snapshot directory.""" @@ -296,6 +312,9 @@ class ModelConfig(BaseModelConfig): debug: DebugModelConfig = DebugModelConfig() """Debugging knobs for the model and distributed training.""" + init: InitConfig = InitConfig() + """Random weight initialization, used when ``debug.random_init`` resolves to ``init``.""" + fused_lm_head_token_chunk_size: int | Literal["disabled"] = 8192 """Flattened token chunk size for the fused LM head. ``int >= 1`` sets the tokens per LM-head chunk explicitly; ``disabled`` uses the vanilla LM head. In SFT the fused head computes the summed cross-entropy and its gradients chunk by chunk, holding one chunk's full-vocab logits at a time.""" @@ -642,6 +661,13 @@ def resolve_freeze_moe_router_auto(self): self.model.freeze_moe_router = True return self + @model_validator(mode="after") + def resolve_random_init_mode_auto(self): + """Resolve ``model.debug.random_init_mode='auto'``: RL skips weight loading without initializing.""" + if self.model.debug.random_init_mode == "auto": + self.model.debug.random_init_mode = "skip" + return self + @model_validator(mode="after") def deepep_disables_grad_clipping(self): if self.model.ep != 1 and self.model.moe.dispatch.type == "deepep" and self.optim.max_norm is not None: diff --git a/src/prime_rl/trainer/init.py b/src/prime_rl/trainer/init.py new file mode 100644 index 00000000000..3d2d49f4c0c --- /dev/null +++ b/src/prime_rl/trainer/init.py @@ -0,0 +1,102 @@ +import hashlib +import math +from typing import Literal + +import torch +from torch import nn +from torch.distributed.tensor import DTensor +from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + +from prime_rl.configs.trainer import InitConfig +from prime_rl.utils.weights import resolve_fqn + +ParamRole = Literal["embedding", "output", "residual_out", "hidden", "norm", "bias"] + +# Elements per random block. Blocks are seeded independently, so each rank draws only the +# blocks that overlap its shard and the values do not depend on how the tensor is sharded. +BLOCK_NUMEL = 1 << 22 + + +def param_role(name: str, param: torch.Tensor) -> ParamRole: + """Classify a parameter by its canonical name. Init schemes (and later per-role learning + rates) are defined over these roles.""" + if "embed_tokens" in name: + return "embedding" + if "lm_head" in name: + return "output" + if param.ndim == 1: + if name.endswith("bias"): + return "bias" + if "norm" in name: + return "norm" + raise ValueError(f"No init role for 1-D parameter {name!r}") + if name.endswith(("o_proj.weight", "down_proj.weight", "down_proj")): + return "residual_out" + return "hidden" + + +def role_std(role: ParamRole, config: InitConfig, num_layers: int) -> float: + if config.scheme == "scaled" and role == "residual_out": + return config.std / math.sqrt(2 * num_layers) + return config.std + + +def block_seed(seed: int, name: str, block: int) -> int: + digest = hashlib.sha256(f"{seed}:{name}:{block}".encode()).digest() + return int.from_bytes(digest[:8], "little") & ((1 << 63) - 1) + + +def normal_shard( + name: str, + global_shape: tuple[int, ...], + local_shape: tuple[int, ...], + global_offset: tuple[int, ...], + seed: int, + device: torch.device, +) -> torch.Tensor: + """The ``local_shape`` slice at ``global_offset`` of a standard-normal tensor of + ``global_shape`` whose values depend only on ``(seed, name)``. The tensor is drawn in + blocks of rows along dim 0; only the blocks overlapping the slice are generated.""" + row_numel = math.prod(global_shape[1:]) + rows_per_block = max(1, BLOCK_NUMEL // max(row_numel, 1)) + start, stop = global_offset[0], global_offset[0] + local_shape[0] + pieces = [] + for block in range(start // rows_per_block, math.ceil(stop / rows_per_block)): + block_start = block * rows_per_block + block_rows = min(rows_per_block, global_shape[0] - block_start) + generator = torch.Generator(device=device).manual_seed(block_seed(seed, name, block)) + values = torch.randn((block_rows, *global_shape[1:]), generator=generator, device=device) + lo, hi = max(start, block_start) - block_start, min(stop, block_start + block_rows) - block_start + pieces.append(values[lo:hi]) + out = torch.cat(pieces) if pieces else torch.empty((0, *global_shape[1:]), device=device) + for dim in range(1, len(global_shape)): + out = out.narrow(dim, global_offset[dim], local_shape[dim]) + return out + + +@torch.no_grad() +def init_weights(model: nn.Module, config: InitConfig) -> None: + """Initialize every parameter in place. Each rank fills only its local shard, without + communication, and the values do not depend on the world size or the sharding.""" + num_layers = model.config.num_hidden_layers + for name, param in model.named_parameters(): + # Canonical names, so wrappers such as activation checkpointing don't change the values. + name = resolve_fqn(model, name) + role = param_role(name, param) + local = param.to_local() if isinstance(param, DTensor) else param + if role == "norm": + local.fill_(1.0) + continue + if role == "bias": + local.zero_() + continue + if isinstance(param, DTensor): + local_shape, global_offset = compute_local_shape_and_global_offset( + param.shape, param.device_mesh, param.placements + ) + else: + local_shape, global_offset = tuple(param.shape), (0,) * param.ndim + values = normal_shard( + name, tuple(param.shape), tuple(local_shape), tuple(global_offset), config.seed, local.device + ) + local.copy_(values.mul_(role_std(role, config, num_layers))) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 47455fdecf2..821f9373adf 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -38,6 +38,7 @@ from prime_rl.multimodal import ForwardPolicy from prime_rl.trainer.activation_checkpointing import get_activation_checkpoint_wrapper from prime_rl.trainer.distributed.embedding_parallel import EmbeddingParallel +from prime_rl.trainer.init import init_weights from prime_rl.trainer.lora import apply_lora_to_model, freeze_all_except_lora_and_specified, strip_lora_from_state_dict from prime_rl.trainer.models import ( AutoModelForCausalLMPrimeRL, @@ -700,7 +701,13 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral logger = get_logger() if config.debug.random_init: - logger.warning("Randomly initializing model. Skipping loading weights from HF.") + if config.debug.random_init_mode == "init": + logger.info(f"Randomly initializing model weights ({config.init})") + init_weights(model, config.init) + else: + logger.warning( + "Skipping weight loading. Parameters keep arbitrary memory contents (random_init_mode='skip')." + ) return if not Path(config.name).exists(): From 487a39380656ca3164c55f6c041c6a0628c10f44 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 13 Oct 2026 19:32:39 +0200 Subject: [PATCH 002/132] feat(pretrain): uv run pretrain with text data and perplexity evals A pretraining entrypoint with its own config, launcher and training loop, built from the shared model setup, optimizer, scheduler, checkpointing and monitors. The SFT path is unchanged. - [model.arch] overrides the baseline model's HF config; the launcher writes the result to /model and trains from there. - Text data: documents are packed back to back into rows of seq_len tokens and continue into the next row (no truncation or padding), with optional document masking and a holdout. - [val]: loss, perplexity and bits per byte on fixed, pre-packed held-out sets. - Evals run after the optimizer update, so step N matches the step-N checkpoint. --- .../src/prime_rl/configs/pretrain.py | 261 +++++++++++ pyproject.toml | 1 + src/prime_rl/entrypoints/pretrain.py | 120 +++++ src/prime_rl/trainer/pretrain/data.py | 183 ++++++++ src/prime_rl/trainer/pretrain/evals/common.py | 65 +++ .../trainer/pretrain/evals/perplexity.py | 61 +++ src/prime_rl/trainer/pretrain/train.py | 410 ++++++++++++++++++ 7 files changed, 1101 insertions(+) create mode 100644 packages/prime-rl-configs/src/prime_rl/configs/pretrain.py create mode 100644 src/prime_rl/entrypoints/pretrain.py create mode 100644 src/prime_rl/trainer/pretrain/data.py create mode 100644 src/prime_rl/trainer/pretrain/evals/common.py create mode 100644 src/prime_rl/trainer/pretrain/evals/perplexity.py create mode 100644 src/prime_rl/trainer/pretrain/train.py diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py new file mode 100644 index 00000000000..71a896d315e --- /dev/null +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -0,0 +1,261 @@ +import uuid +from pathlib import Path +from typing import Annotated, Any, Literal, TypeAlias + +from pydantic import Field, model_validator + +from prime_rl.configs.monitors import TrainMonitorsConfig +from prime_rl.configs.sft import BaseDataConfig, FakeDataConfig +from prime_rl.configs.shared import EnvVars, HeartbeatConfig, ResumeConfig, RunConfig, TrainerLogConfig +from prime_rl.configs.trainer import ( + AdamWConfig, + CheckpointConfig, + ConstantSchedulerConfig, + GCConfig, + ModelConfig, + OptimizerConfig, + SchedulerConfig, + TokenizerConfig, + validate_scheduler, +) +from prime_rl.utils.config import BaseConfig, default_output_dir + + +class PretrainModelConfig(ModelConfig): + arch: dict[str, Any] = {} + """HF config fields that override the config of the baseline model ``name``, e.g. ``hidden_size`` or ``num_hidden_layers``. The baseline provides the architecture class, the tokenizer and the vocabulary. The launcher writes the resulting config and tokenizer to ``/model`` and trains from there.""" + + +class TextDataConfig(BaseDataConfig): + type: Literal["text"] = "text" + + name: str = "HuggingFaceFW/fineweb-edu" + """HF dataset name or local path.""" + + subsets: list[str] | None = None + """Subsets (HF configs) to load.""" + + splits: list[str] | None = None + """Splits to load. Defaults to ``train``.""" + + revision: str | None = None + """HF dataset revision.""" + + probabilities: list[float] | None = None + """Sampling probabilities when interleaving several subsets/splits.""" + + stopping_strategy: Literal["first_exhausted", "all_exhausted"] = "all_exhausted" + """Stopping strategy when interleaving several subsets/splits.""" + + text_column: str = "text" + """Column holding the raw document text.""" + + shuffle: bool = True + """Shuffle the documents at the start of each epoch.""" + + seed: int = 0 + """Shuffle seed. Each epoch adds its index to it.""" + + seq_len: int = Field(2048, ge=1) + """Tokens per row. Documents are packed back to back and continue across rows; nothing is truncated or padded.""" + + batch_size: int = Field(256, ge=1) + """Rows per optimizer step, so a step trains on ``batch_size * seq_len`` tokens.""" + + document_masking: bool = True + """Restrict attention to the current document and restart positions at every document boundary. If false, a row is one sequence.""" + + holdout_docs: int = Field(0, ge=0) + """Hold the last N documents (in dataset order, before shuffling) out of training. With ``[val]`` set, they are evaluated as the ``holdout`` source.""" + + @model_validator(mode="after") + def validate_subsets_and_splits(self): + if self.subsets is not None and self.splits is not None and len(self.subsets) != len(self.splits): + raise ValueError("data.subsets and data.splits must have the same length") + sources = len(self.subsets or self.splits or [None]) + if self.probabilities is not None and len(self.probabilities) != sources: + raise ValueError("data.probabilities needs one entry per subset/split") + return self + + +PretrainDataConfig: TypeAlias = Annotated[FakeDataConfig | TextDataConfig, Field(discriminator="type")] + + +class ValSourceConfig(BaseConfig): + name: str + """Label for the metrics (``val/{name}/loss``, ``/ppl``, ``/bpb``).""" + + dataset: str + """HF dataset name or local path.""" + + subset: str | None = None + """HF config to load.""" + + split: str = "validation" + """Split to load.""" + + revision: str | None = None + """HF dataset revision.""" + + text_column: str = "text" + """Column holding the raw document text.""" + + max_tokens: int = Field(10_000_000, ge=1) + """Token budget. Documents are taken in dataset order until it is reached.""" + + +class ValConfig(BaseConfig): + """Perplexity evals: loss, perplexity and bits per byte on fixed held-out token sets. + Each set is tokenized and packed once at startup, so every eval sees the same tokens.""" + + interval: int = Field(500, ge=1) + """Evaluate every N steps, after that step's optimizer update.""" + + eval_on_start: bool = True + """Evaluate the initial weights before the first step (logged as step 0).""" + + source: list[ValSourceConfig] = [] + """Held-out sets. With ``data.holdout_docs > 0`` the training holdout is added as ``holdout``.""" + + +class PretrainDeploymentConfig(BaseConfig): + num_train_gpus: int = Field(1, ge=1) + """GPUs for the trainer, on this node.""" + + +class PretrainConfig(BaseConfig): + model: PretrainModelConfig = PretrainModelConfig() + + env_vars: EnvVars = {} + """Extra environment variables for the trainer processes.""" + + tokenizer: TokenizerConfig = TokenizerConfig() + + data: PretrainDataConfig = TextDataConfig() + + val: ValConfig | None = None + """Perplexity evals. If None, none run.""" + + optim: OptimizerConfig = AdamWConfig() + + scheduler: SchedulerConfig = ConstantSchedulerConfig() + + ckpt: CheckpointConfig | None = CheckpointConfig() + """Checkpointing. Defaults to a checkpoint at the end of training.""" + + resume: ResumeConfig | None = None + """Resume the run from a checkpoint. If None, does not resume.""" + + log: TrainerLogConfig = TrainerLogConfig() + + monitors: TrainMonitorsConfig = TrainMonitorsConfig() + """Metric monitors (``monitors.wandb``, ``monitors.file``).""" + + run: RunConfig = Field(default_factory=RunConfig) + """Run metadata. ``run.name`` names the run directory under ``output_dir``.""" + + output_dir: Path = Field(default_factory=default_output_dir) + """Directory that groups related runs. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``.""" + + clean: bool = False + """Delete the run directory before starting. Required to overwrite a previous run's artifacts when not resuming.""" + + matmul_precision: Literal["highest", "high", "medium"] = "high" + """Precision for float32 matrix multiplications (``torch.set_float32_matmul_precision``).""" + + max_steps: int | None = None + """Maximum training steps. If None, runs indefinitely.""" + + gc: GCConfig | None = GCConfig() + """Deterministic garbage collection every N steps. Set to null for Python's default GC.""" + + dist_timeout_seconds: int = 900 + """Timeout in seconds for torch distributed ops.""" + + heartbeat: HeartbeatConfig | None = None + """BetterStack heartbeat configuration.""" + + deployment: PretrainDeploymentConfig = PretrainDeploymentConfig() + + dry_run: bool = False + """Only validate and dump resolved configs, then exit.""" + + @property + def run_dir(self) -> Path: + assert self.run.dir is not None # resolved at construction + return self.output_dir / self.run.dir + + @model_validator(mode="before") + @classmethod + def default_random_init(cls, data): + """Pretraining starts from randomly initialized weights unless ``model.debug.random_init = false`` + (continued pretraining from ``model.name``).""" + if isinstance(data, dict): + model = data.setdefault("model", {}) + if isinstance(model, dict): + debug = model.setdefault("debug", {}) + if isinstance(debug, dict): + debug.setdefault("random_init", True) + return data + + @model_validator(mode="after") + def resolve_autos(self): + """Resolve the shared ``auto`` fields for pretraining: real random init, a trainable fp32 + MoE router, and weight decay 0.1.""" + if self.model.debug.random_init_mode == "auto": + self.model.debug.random_init_mode = "init" + if self.model.moe_router_dtype == "auto": + self.model.moe_router_dtype = "float32" + if self.model.freeze_moe_router == "auto": + self.model.freeze_moe_router = False + if self.optim.weight_decay == "auto": + self.optim.weight_decay = 0.1 + return self + + @model_validator(mode="after") + def auto_setup_run_identity(self): + if self.run.name is None: + dataset = str(getattr(self.data, "name", "fake")).split("/")[-1] + model = self.model.name.split("/")[-1] + self.run.name = "--".join([dataset, model, uuid.uuid4().hex[:8]]).lower() + if self.run.dir is None: + self.run.dir = self.run.name + if self.monitors.wandb is not None and self.monitors.wandb.name is None: + self.monitors.wandb.name = self.run.name + return self + + @model_validator(mode="after") + def auto_setup_tokenizer(self): + if self.tokenizer.name is None: + self.tokenizer.name = self.model.name + if self.tokenizer.trust_remote_code is None: + self.tokenizer.trust_remote_code = self.model.trust_remote_code + return self + + @model_validator(mode="after") + def validate_monitors(self): + if self.monitors.prime is not None: + raise ValueError("monitors.prime is not supported for pretraining yet.") + return self + + @model_validator(mode="after") + def validate_cp(self): + if self.model.cp > 1: + if self.data.seq_len % self.model.cp != 0: + raise ValueError("data.seq_len must be divisible by model.cp") + if self.data.micro_batch_size != 1: + raise ValueError("data.micro_batch_size must be 1 with context parallelism") + if self.val is not None: + raise ValueError("[val] does not support context parallelism yet") + return self + + @model_validator(mode="after") + def validate_holdout(self): + if self.val is not None and self.data.type == "text" and self.data.holdout_docs == 0 and not self.val.source: + raise ValueError("[val] needs at least one source, or data.holdout_docs > 0") + return self + + @model_validator(mode="after") + def validate_scheduler_steps(self): + validate_scheduler(self.scheduler, self.max_steps) + return self diff --git a/pyproject.toml b/pyproject.toml index 4a06615818d..9a56bf17389 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ [project.scripts] rl = "prime_rl.entrypoints.rl:main" sft = "prime_rl.entrypoints.sft:main" +pretrain = "prime_rl.entrypoints.pretrain:main" inference = "prime_rl.entrypoints.inference:main" trainer = "prime_rl.entrypoints.trainer:main" orchestrator = "prime_rl.entrypoints.orchestrator:main" diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py new file mode 100644 index 00000000000..e94ad716917 --- /dev/null +++ b/src/prime_rl/entrypoints/pretrain.py @@ -0,0 +1,120 @@ +import json +import os +import uuid + +from prime_rl.configs.pretrain import PretrainConfig +from prime_rl.utils.config import cli, dump_resolved_config +from prime_rl.utils.logger import setup_logger +from prime_rl.utils.pathing import ( + format_config_message, + format_log_message, + prepare_attempt_dirs, + validate_run_dir, + write_launch_artifacts, +) +from prime_rl.utils.process import ( + DEFAULT_COMMON_ENV_VARS, + DEFAULT_TRAINER_ENV_VARS, + ProcessGroup, + set_proc_title, + torchrun_cmd, +) + +PRETRAIN_CONFIG = "pretrain.json" + + +def write_model_dir(config: PretrainConfig) -> None: + """Write the baseline's config with ``[model.arch]`` applied, plus its tokenizer, to + ``/model``, and point the model and tokenizer at it.""" + from transformers import AutoConfig, AutoTokenizer + + import prime_rl.trainer.models # noqa: F401 - registers the custom model configs with AutoConfig + + model_dir = config.run_dir / "model" + hf_config = AutoConfig.from_pretrained(config.model.name, trust_remote_code=config.model.trust_remote_code) + unknown = [key for key in config.model.arch if not hasattr(hf_config, key)] + if unknown: + raise ValueError(f"[model.arch] fields {unknown} do not exist on the {hf_config.model_type!r} config") + for key, value in config.model.arch.items(): + setattr(hf_config, key, value) + hf_config.save_pretrained(model_dir) + tokenizer = AutoTokenizer.from_pretrained( + config.tokenizer.name, trust_remote_code=config.tokenizer.trust_remote_code + ) + tokenizer.save_pretrained(model_dir) + config.model.name = config.tokenizer.name = model_dir.as_posix() + + +def download_data(config: PretrainConfig) -> None: + """Download and prepare every dataset once, so the trainer ranks load them from the cache.""" + from prime_rl.trainer.pretrain.data import load_eval_documents, load_text_dataset + + if config.data.type == "text": + load_text_dataset(config.data) + if config.val is not None: + for source in config.val.source: + load_eval_documents(source.dataset, source.subset, source.split, source.revision, source.text_column) + + +def pretrain(config: PretrainConfig): + logger = setup_logger(config.log.level or "info", json_logging=config.log.json_logging) + os.environ.setdefault("PRL_RUN_ID", uuid.uuid4().hex) + assert config.run.name is not None + os.environ["PRL_RUN_NAME"] = config.run.name + + resuming = config.resume is not None + ckpt_output_dir = config.ckpt.output_dir if config.ckpt else None + validate_run_dir( + config.run_dir, + output_dir=config.output_dir, + resuming=resuming, + clean=config.clean, + ckpt_output_dir=ckpt_output_dir, + ) + config.run_dir.mkdir(parents=True, exist_ok=True) + + logger.info("Starting pretraining run") + config_dir, log_dir = prepare_attempt_dirs(config.run_dir) + write_launch_artifacts(config_dir, "pretrain") + + if not config.dry_run: + from prime_rl.trainer.model import pre_download_model + + pre_download_model(config.model.name, skip_weights=config.model.debug.random_init) + if config.model.arch: + write_model_dir(config) + download_data(config) + + config_path = config_dir / PRETRAIN_CONFIG + config_path.write_text(json.dumps(dump_resolved_config(config, exclude={"dry_run", "clean"}), indent=2)) + logger.info(f"Configs:\n{format_config_message(config_dir, 'pretrain', [('Pretrain', config_path)])}") + if config.dry_run: + logger.success("Dry run complete. To start the run, remove --dry-run from your command.") + return + + logger.info(format_log_message(log_dir=log_dir, trainer=True)) + with ProcessGroup() as processes: + processes.start( + "trainer", + torchrun_cmd( + "prime_rl.trainer.pretrain.train", + config_path, + config.deployment.num_train_gpus, + log_dir, + config.log.ranks_filter, + ), + env={**os.environ, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, **config.env_vars}, + log_path=log_dir / "trainer.log", + ) + logger.success(f"Started trainer with {config.deployment.num_train_gpus} GPU(s)") + processes.wait("trainer") + logger.success("Pretraining finished!") + + +def main(): + set_proc_title("Pretrain") + pretrain(cli(PretrainConfig)) + + +if __name__ == "__main__": + main() diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py new file mode 100644 index 00000000000..ab0d17b3eb1 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/data.py @@ -0,0 +1,183 @@ +from collections.abc import Iterable, Iterator +from typing import cast + +from datasets import Dataset, load_dataset +from torchdata.stateful_dataloader import StatefulDataLoader +from transformers.tokenization_utils import PreTrainedTokenizer + +from prime_rl.configs.pretrain import TextDataConfig +from prime_rl.trainer.sft.data import Sample, StatefulIterableDataset, cat_collate, setup_and_interleave_datasets +from prime_rl.utils.logger import get_logger + +LIST_KEYS = ("input_ids", "target_ids", "loss_mask", "position_ids", "seq_lens") + + +def tokenize_document(tokenizer: PreTrainedTokenizer, text: str) -> list[int]: + """Token ids of one document, with the tokenizer's BOS (if it adds one) and a closing EOS.""" + token_ids = tokenizer(text)["input_ids"] + if not token_ids or token_ids[-1] != tokenizer.eos_token_id: + token_ids.append(tokenizer.eos_token_id) + return token_ids + + +def cut_row(tokens: list[int], doc_starts: list[int], seq_len: int, document_masking: bool) -> Sample: + """One training row from ``seq_len + 1`` tokens: inputs are the first ``seq_len``, targets the + last ``seq_len``. ``doc_starts`` are the offsets in ``tokens`` where a document begins. With + document masking, every document (and the continued document at offset 0) is its own segment + with positions starting at 0.""" + assert len(tokens) == seq_len + 1 + if document_masking: + boundaries = sorted({0, *(start for start in doc_starts if 0 < start < seq_len)}) + seq_lens = [end - start for start, end in zip(boundaries, [*boundaries[1:], seq_len])] + else: + seq_lens = [seq_len] + return { + "input_ids": tokens[:-1], + "target_ids": tokens[1:], + "loss_mask": [True] * seq_len, + "position_ids": [position for length in seq_lens for position in range(length)], + "seq_lens": seq_lens, + "mm_kwargs": None, + "mm_token_type_ids": None, + } + + +class TextDataset(StatefulIterableDataset): + """Tokenized documents of this data rank, reshuffled every epoch. Documents are addressed by + their global index, so resuming jumps straight to the saved position.""" + + def __init__( + self, + dataset: Dataset, + tokenizer: PreTrainedTokenizer, + text_column: str = "text", + shuffle: bool = True, + seed: int = 0, + non_dp_size: int = 1, + ): + super().__init__(non_dp_size) + self.dataset = dataset + self.num_examples = len(dataset) + self.tokenizer = tokenizer + self.text_column = text_column + self.shuffle = shuffle + self.seed = seed + + def epoch_dataset(self) -> Dataset: + return self.dataset.shuffle(seed=self.seed + self.epoch) if self.shuffle else self.dataset + + def __iter__(self) -> Iterator[list[int]]: + self._setup_world_info() + dataset = self.epoch_dataset() + while True: + self.step += 1 + epoch = (self.step - 1) // self.num_examples + if epoch > self.epoch: + self.epoch = epoch + dataset = self.epoch_dataset() + if (self.step - 1) % self.data_world_size != self.data_rank: + continue + example = cast(dict, dataset[(self.step - 1) % self.num_examples]) + token_ids = tokenize_document(self.tokenizer, example[self.text_column]) + source = example.get("__subset") or example.get("__split") + self.num_samples[source] += 1 + self.num_tokens[source] += len(token_ids) + yield token_ids + + +class StreamPacker(StatefulIterableDataset): + """Cuts a document stream into rows of exactly ``seq_len`` tokens (see ``cut_row``) and + yields ``rows_per_sample`` consecutive rows concatenated as one packed sample. Documents + continue across rows, so nothing is truncated or padded, but attention never crosses a + row. The tokens not yet emitted are part of the checkpoint state. Any iterable of + documents works as the source; checkpointing needs a ``TextDataset``.""" + + def __init__( + self, + dataset: TextDataset | Iterable[list[int]], + seq_len: int, + document_masking: bool, + rows_per_sample: int = 1, + ): + self.dataset = dataset + self.seq_len = seq_len + self.document_masking = document_masking + self.rows_per_sample = rows_per_sample + self.buffer: list[int] = [] + self.doc_starts: list[int] = [] + + def state_dict(self) -> dict: + return { + "dataset": self.dataset.state_dict(), + "progress": {"num_samples": dict(self.dataset.num_samples), "num_tokens": dict(self.dataset.num_tokens)}, + "buffer": self.buffer, + "doc_starts": self.doc_starts, + } + + def load_state_dict(self, state_dict: dict): + self.dataset.load_state_dict(state_dict["dataset"]) + self.dataset.num_samples.update(state_dict["progress"]["num_samples"]) + self.dataset.num_tokens.update(state_dict["progress"]["num_tokens"]) + self.buffer = state_dict["buffer"] + self.doc_starts = state_dict["doc_starts"] + + def next_row(self) -> Sample: + row = cut_row(self.buffer[: self.seq_len + 1], self.doc_starts, self.seq_len, self.document_masking) + self.buffer = self.buffer[self.seq_len :] + self.doc_starts = [start - self.seq_len for start in self.doc_starts if start >= self.seq_len] + return row + + def __iter__(self) -> Iterator[Sample]: + for document in self.dataset: + self.doc_starts.append(len(self.buffer)) + self.buffer.extend(document) + while len(self.buffer) >= self.seq_len * self.rows_per_sample + 1: + rows = [self.next_row() for _ in range(self.rows_per_sample)] + sample = {key: [value for row in rows for value in row[key]] for key in LIST_KEYS} + yield {**sample, "mm_kwargs": None, "mm_token_type_ids": None} + + +def load_text_dataset(config: TextDataConfig) -> Dataset: + """Load (and interleave) the configured subsets/splits. Only their files are downloaded.""" + splits = config.splits or ["train"] * len(config.subsets or [None]) + subsets = config.subsets or [None] * len(splits) + if len(splits) == 1: + dataset = load_dataset(config.name, subsets[0], split=splits[0], revision=config.revision) + return cast(Dataset, dataset).select_columns([config.text_column]) + dataset = setup_and_interleave_datasets( + dataset_name=config.name, + subsets_and_splits=list(zip(subsets, splits)), + probabilities=config.probabilities, + stopping_strategy=config.stopping_strategy, + seed=config.seed, + revision=config.revision, + ) + return dataset.select_columns([config.text_column, "__subset", "__split"]) + + +def split_holdout(dataset: Dataset, holdout_docs: int) -> tuple[Dataset, Dataset | None]: + """Split off the last ``holdout_docs`` documents (in dataset order) as a held-out set.""" + if holdout_docs == 0: + return dataset, None + num_train = len(dataset) - holdout_docs + assert num_train > 0, f"data.holdout_docs ({holdout_docs}) must be smaller than the dataset ({len(dataset)})" + return dataset.select(range(num_train)), dataset.select(range(num_train, len(dataset))) + + +def load_eval_documents( + dataset: str, subset: str | None, split: str, revision: str | None, text_column: str +) -> Dataset: + return cast(Dataset, load_dataset(dataset, subset, split=split, revision=revision)).select_columns([text_column]) + + +def setup_text_dataloader( + dataset: Dataset, tokenizer: PreTrainedTokenizer, config: TextDataConfig, non_dp_size: int +) -> StatefulDataLoader: + get_logger().info(f"Training on {len(dataset):,} documents from {config.name}") + text_dataset = TextDataset( + dataset, tokenizer, config.text_column, shuffle=config.shuffle, seed=config.seed, non_dp_size=non_dp_size + ) + packer = StreamPacker(text_dataset, config.seq_len, config.document_masking, config.micro_batch_size) + return StatefulDataLoader( + packer, batch_size=1, collate_fn=cat_collate, num_workers=config.num_workers, pin_memory=True + ) diff --git a/src/prime_rl/trainer/pretrain/evals/common.py b/src/prime_rl/trainer/pretrain/evals/common.py new file mode 100644 index 00000000000..e6bbeb5a078 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/evals/common.py @@ -0,0 +1,65 @@ +from typing import TypeVar + +import torch +import torch.distributed as dist +from transformers.tokenization_utils import PreTrainedTokenizer + +from prime_rl.trainer.sft.data import Sample +from prime_rl.trainer.world import get_world + +T = TypeVar("T") + + +def bytes_to_unicode() -> dict[int, str]: + """The byte -> character map of byte-level BPE tokenizers (GPT-2, Llama 3, Qwen, DeepSeek).""" + printable = ( + list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1)) + ) + chars = printable[:] + n = 0 + for byte in range(256): + if byte not in printable: + printable.append(byte) + chars.append(256 + n) + n += 1 + return dict(zip(printable, map(chr, chars))) + + +def token_byte_lengths(tokenizer: PreTrainedTokenizer) -> torch.Tensor: + """UTF-8 bytes each token id stands for. Special tokens count as 0 bytes. Exact for + byte-level BPE tokenizers, approximate otherwise.""" + byte_chars = set(bytes_to_unicode().values()) + special_ids = set(tokenizer.all_special_ids) | set(tokenizer.added_tokens_decoder) + lengths = torch.zeros(len(tokenizer), dtype=torch.long) + for token_id, token in enumerate(tokenizer.convert_ids_to_tokens(list(range(len(tokenizer))))): + if token is None or token_id in special_ids: + continue + if all(char in byte_chars for char in token): + lengths[token_id] = len(token) + else: + lengths[token_id] = len(tokenizer.convert_tokens_to_string([token]).encode()) + return lengths + + +def local_rows(rows: list[T]) -> list[T | None]: + """This rank's share of ``rows``, padded with ``None`` so every rank runs the same number of + forward passes (each forward is an FSDP collective).""" + world = get_world() + per_rank = -(-len(rows) // world.world_size) + mine: list[T | None] = list(rows[world.rank :: world.world_size]) + return mine + [None] * (per_rank - len(mine)) + + +def to_device(row: Sample) -> dict[str, torch.Tensor]: + return { + "input_ids": torch.tensor(row["input_ids"], device="cuda").unsqueeze(0), + "position_ids": torch.tensor(row["position_ids"], device="cuda").unsqueeze(0), + "target_ids": torch.tensor(row["target_ids"], device="cuda").unsqueeze(0), + "loss_mask": torch.tensor(row["loss_mask"], device="cuda").unsqueeze(0), + "seq_lens": torch.tensor(row["seq_lens"], device="cuda"), + } + + +def all_reduce_sum(tensor: torch.Tensor) -> torch.Tensor: + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + return tensor diff --git a/src/prime_rl/trainer/pretrain/evals/perplexity.py b/src/prime_rl/trainer/pretrain/evals/perplexity.py new file mode 100644 index 00000000000..cded029475c --- /dev/null +++ b/src/prime_rl/trainer/pretrain/evals/perplexity.py @@ -0,0 +1,61 @@ +import math +from dataclasses import dataclass +from itertools import islice + +import torch +from torch import nn +from transformers.tokenization_utils import PreTrainedTokenizer + +from prime_rl.configs.pretrain import ValSourceConfig +from prime_rl.trainer.model import forward +from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX +from prime_rl.trainer.pretrain.data import StreamPacker, load_eval_documents, tokenize_document +from prime_rl.trainer.pretrain.evals.common import all_reduce_sum, local_rows, to_device +from prime_rl.trainer.sft.data import Sample + + +@dataclass +class PerplexitySet: + """A fixed, pre-packed held-out token set.""" + + name: str + rows: list[Sample] + + +def build_perplexity_set( + name: str, documents, tokenizer: PreTrainedTokenizer, seq_len: int, max_tokens: int, text_column: str = "text" +) -> PerplexitySet: + """Pack documents (in order) into full rows until ``max_tokens`` is reached. Rows use + document masking, so every document is scored from its own start.""" + token_stream = (tokenize_document(tokenizer, example[text_column]) for example in documents) + rows = list(islice(StreamPacker(token_stream, seq_len, document_masking=True), max(1, max_tokens // seq_len))) + assert rows, f"Eval source {name!r} has fewer than {seq_len + 1} tokens" + return PerplexitySet(name=name, rows=rows) + + +def load_perplexity_set(source: ValSourceConfig, tokenizer: PreTrainedTokenizer, seq_len: int) -> PerplexitySet: + documents = load_eval_documents(source.dataset, source.subset, source.split, source.revision, source.text_column) + return build_perplexity_set(source.name, documents, tokenizer, seq_len, source.max_tokens, source.text_column) + + +@torch.no_grad() +def evaluate_perplexity(model: nn.Module, sets: list[PerplexitySet], byte_lengths: torch.Tensor) -> dict[str, float]: + """Loss (nats per token), perplexity and bits per byte for every set.""" + byte_lengths = byte_lengths.to("cuda") + metrics = {} + for eval_set in sets: + totals = torch.zeros(3, dtype=torch.float64, device="cuda") # loss, tokens, bytes + for row in local_rows(eval_set.rows): + batch = to_device(row if row is not None else eval_set.rows[0]) + loss_mask = batch["loss_mask"] if row is not None else torch.zeros_like(batch["loss_mask"]) + labels = batch["target_ids"].masked_fill(~loss_mask, IGNORE_INDEX) + out = forward(model, batch["input_ids"], batch["position_ids"], seq_lens=batch["seq_lens"], labels=labels) + totals[0] += out["loss"].double() + totals[1] += loss_mask.sum() + totals[2] += byte_lengths[batch["target_ids"][loss_mask]].sum() + loss_sum, num_tokens, num_bytes = all_reduce_sum(totals).tolist() + loss = loss_sum / num_tokens + metrics[f"val/{eval_set.name}/loss"] = loss + metrics[f"val/{eval_set.name}/ppl"] = math.exp(min(loss, 20)) + metrics[f"val/{eval_set.name}/bpb"] = loss_sum / (math.log(2) * num_bytes) + return metrics diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py new file mode 100644 index 00000000000..a29bf227c4b --- /dev/null +++ b/src/prime_rl/trainer/pretrain/train.py @@ -0,0 +1,410 @@ +import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before import + +import asyncio +import math +import time +from datetime import timedelta + +# ruff: noqa: I001 + +from prime_rl.utils.act_offloading import maybe_activation_offloading +import torch +import torch.distributed as dist +from torch import nn +from transformers import AutoConfig + +from prime_rl import monitors +from prime_rl.configs.pretrain import PretrainConfig +from prime_rl.configs.trainer import CheckpointConfig +from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager +from prime_rl.trainer.model import ( + forward, + get_expert_load_stats, + get_full_offload_dtype_policy, + get_global_moe_stats, + get_load_balance_stats, + is_tt_moe_model, + reshard_module, + setup_model, + setup_tokenizer, +) +from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX +from prime_rl.trainer.optim import setup_optimizer +from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep +from prime_rl.trainer.perf import get_perf_counter +from prime_rl.trainer.pretrain.data import load_text_dataset, setup_text_dataloader, split_holdout +from prime_rl.trainer.pretrain.evals.common import token_byte_lengths +from prime_rl.trainer.pretrain.evals.perplexity import build_perplexity_set, evaluate_perplexity, load_perplexity_set +from prime_rl.trainer.scheduler import setup_scheduler +from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state, setup_dataloader, setup_dataset +from prime_rl.trainer.utils import ( + GarbageCollection, + begin_backward, + clip_grad_norm_, + finish_backward, + get_ckpt_disk_metrics, + prepare_gradient_offload, + scale_gradients_, + setup_full_cpu_optimizer_offload, + setup_torch_distributed, +) +from prime_rl.trainer.world import get_world +from prime_rl.utils.config import cli +from prime_rl.utils.cp import setup_context_parallel, setup_cp_params, shard_for_cp +from prime_rl.utils.heartbeat import Heartbeat +from prime_rl.utils.logger import setup_logger +from prime_rl.utils.pathing import resolve_latest_ckpt_step +from prime_rl.utils.process import set_proc_title +from prime_rl.utils.utils import clean_exit + + +def check_arch(config: PretrainConfig) -> None: + """``[model.arch]`` is applied by the launcher, which writes the model dir that ``model.name`` + points at. Fail if the trainer was started on a config that doesn't match.""" + hf_config = AutoConfig.from_pretrained(config.model.name, trust_remote_code=config.model.trust_remote_code) + mismatched = {key: value for key, value in config.model.arch.items() if getattr(hf_config, key, None) != value} + if mismatched: + raise ValueError( + f"{config.model.name} does not match [model.arch] {mismatched}. Launch through `uv run pretrain`, " + "which writes the model dir with the overrides applied." + ) + + +def count_parameters(model: nn.Module) -> dict[str, int]: + """Total, embedding (input embedding and LM head) and non-embedding parameter counts.""" + total = embedding = 0 + for name, param in model.named_parameters(): + total += param.numel() + if "embed_tokens" in name or "lm_head" in name: + embedding += param.numel() + return {"total": total, "embedding": embedding, "non_embedding": total - embedding} + + +@clean_exit +def train(config: PretrainConfig): + world = get_world() + logger = setup_logger(config.log.level, json_logging=config.log.json_logging) + logger.info(f"Starting pretrain trainer in {world} (output_dir={config.run_dir})") + + asyncio.run( + monitors.setup( + producer="trainer", + wandb=config.monitors.wandb, + prime=None, + file=config.monitors.file, + output_dir=config.run_dir, + run_config=config, + overview_flavor="sft", + ) + ) + + heart = Heartbeat(config.heartbeat) if config.heartbeat is not None and world.rank == 0 else None + + setup_torch_distributed( + timeout=timedelta(seconds=config.dist_timeout_seconds), enable_gloo=config.model.offload == "full" + ) + if config.model.offload == "full": + setup_full_cpu_optimizer_offload() + torch.set_float32_matmul_precision(config.matmul_precision) + + resolve_ep(config.model) + parallel_dims = get_parallel_dims(config.model) + + total_micro_batches = config.data.batch_size * config.model.cp + micro_batches_per_step = world.world_size * config.data.micro_batch_size + assert total_micro_batches % micro_batches_per_step == 0, ( + f"batch_size * cp ({total_micro_batches}) must be divisible by " + f"world_size * micro_batch_size ({micro_batches_per_step})" + ) + grad_accum_steps = total_micro_batches // micro_batches_per_step + + if config.model.arch: + check_arch(config) + + ckpt_manager = setup_ckpt_manager(config.run_dir, config.ckpt, resume=config.resume) + checkpoint_step = None + if config.resume is not None: + if config.resume.dir is not None: + checkpoint_step = config.resume.dir_step + else: + checkpoint_step = config.resume.step or resolve_latest_ckpt_step(ckpt_manager.ckpt_dir) + + logger.info(f"Initializing model ({config.model})") + model = setup_model(config.model, parallel_dims, loading_from_checkpoint_later=checkpoint_step is not None) + if parallel_dims.cp_enabled: + setup_context_parallel(model, config.model, parallel_dims) + num_params = count_parameters(model) + logger.info( + f"Parameters: {num_params['non_embedding'] / 1e6:,.1f}M non-embedding, " + f"{num_params['embedding'] / 1e6:,.1f}M embedding, {num_params['total'] / 1e6:,.1f}M total" + ) + + tokenizer = setup_tokenizer(config.tokenizer) + + optimizer, gradient_manager = setup_optimizer( + config.optim, + list(model.named_parameters()), + parallel_dims, + offload=config.model.offload, + model=model, + full_offload_dtype_policy=( + get_full_offload_dtype_policy(model, config.model) if config.model.offload == "full" else None + ), + ) + + scheduler_steps = ( + config.max_steps - checkpoint_step + if config.max_steps is not None and (config.ckpt and config.ckpt.skip_scheduler and checkpoint_step is not None) + else config.max_steps + ) + scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) + + holdout = None + if config.data.type == "text": + documents, holdout = split_holdout(load_text_dataset(config.data), config.data.holdout_docs) + dataloader = setup_text_dataloader(documents, tokenizer, config.data, config.model.cp) + else: + dataloader = setup_dataloader( + setup_dataset(tokenizer, config.data, config.model.cp), config.data, config.model.cp + ) + + eval_seq_len = config.data.seq_len + perplexity_sets = [] + if config.val is not None: + if holdout is not None: + perplexity_sets.append( + build_perplexity_set( + "holdout", holdout, tokenizer, eval_seq_len, max_tokens=10**12, text_column=config.data.text_column + ) + ) + perplexity_sets += [load_perplexity_set(source, tokenizer, eval_seq_len) for source in config.val.source] + for eval_set in perplexity_sets: + logger.info(f"Perplexity eval {eval_set.name!r}: {len(eval_set.rows) * eval_seq_len:,} tokens") + byte_lengths = token_byte_lengths(tokenizer) if perplexity_sets else None + + progress = Progress() + if checkpoint_step is not None: + skip = config.ckpt or CheckpointConfig() + ckpt_manager.load( + checkpoint_step, + model, + [optimizer], + scheduler if not skip.skip_scheduler else None, + progress, + dataloader=dataloader if not skip.skip_dataloader else None, + path=config.resume.dir / "trainer" if config.resume.dir is not None else None, + ) + progress.step += 1 + if skip.skip_scheduler: + scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) + logger.info(f"Resuming from step {checkpoint_step} (dataset_state={get_dataset_state(dataloader)})") + + # Create the iterator only after a potential resume: iter() forks the workers with the + # dataset's current state. + dataiter = iter(dataloader) + + cp_enabled = parallel_dims.cp_enabled + cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() if cp_enabled else 0 + cp_group = parallel_dims.world_mesh["cp"].get_group() if cp_enabled else None + cp_size = parallel_dims.cp + dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() + ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None + is_moe_model = is_tt_moe_model(model) + + def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: + """Forward pass returning (loss_sum, token_count) over unmasked tokens.""" + input_ids = micro_batch["input_ids"].to("cuda", non_blocking=True) + position_ids = micro_batch["position_ids"].to("cuda", non_blocking=True) + target_ids = micro_batch["target_ids"].to("cuda", non_blocking=True) + loss_mask = micro_batch["loss_mask"].to("cuda", non_blocking=True) + seq_lens = micro_batch["seq_lens"].to("cuda", non_blocking=True) + if cp_enabled: + input_ids, position_ids = setup_cp_params( + input_ids, position_ids, cp_rank, cp_size, cp_group, seq_lens=seq_lens, cp_style=config.model.cp_style + ) + target_ids = shard_for_cp(target_ids, cp_rank=cp_rank, cp_world_size=cp_size) + loss_mask = shard_for_cp(loss_mask, cp_rank=cp_rank, cp_world_size=cp_size) + labels = target_ids.masked_fill(~loss_mask, IGNORE_INDEX) + with maybe_activation_offloading(config.model.ac_offloading): + out = forward( + model, input_ids, position_ids, seq_lens=seq_lens, labels=labels, seq_lens_are_pre_shard=cp_enabled + ) + return out["loss"], loss_mask.sum(dtype=torch.int64) + + def run_evals(step: int) -> None: + """Perplexity evals on the weights after ``step`` optimizer updates.""" + metrics = evaluate_perplexity(model, perplexity_sets, byte_lengths) + # The no-grad eval forwards leave the [lm_head, norm] FSDP group unsharded, and their + # routing must not count toward the next step's MoE statistics. + reshard_module(model) + if is_moe_model: + get_load_balance_stats(model) + logger.success(f"Evals | Step {step} | " + " | ".join(f"{k} {v:.4f}" for k, v in sorted(metrics.items()))) + asyncio.run(monitors.log({**metrics, "step": step}, step=step)) + + if checkpoint_step is None and config.val is not None and config.val.eval_on_start: + run_evals(step=0) + + gc_handler = GarbageCollection(config.gc.interval) if config.gc else None + max_memory = torch.cuda.mem_get_info()[1] / 1024**3 + logger.info(f"Starting training loop (max_steps={config.max_steps or 'infinite'})") + while True: + torch.cuda.reset_peak_memory_stats() + if gc_handler is not None: + gc_handler.run(progress.step) + is_last_step = config.max_steps is not None and progress.step >= config.max_steps + step_start_time = time.perf_counter() + + step_loss_sum = torch.tensor(0.0, device="cuda") + nan_loss_count = torch.tensor(0, device="cuda") + moe_stats = ( + { + "max_vio/mean": torch.tensor(0.0), + "max_vio/max": torch.tensor(0.0), + "routing_confidence/mean": torch.tensor(0.0), + } + if is_moe_model + else {} + ) + if gradient_manager is None: + micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) + step_local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") + else: + micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] + local_token_count = sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches) + global_step_token_count = torch.tensor(local_token_count, dtype=torch.int64, device="cuda") + dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count_val = global_step_token_count.item() // cp_size + grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val + prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=True) + + step_tokens_per_expert = 0 + step_local_num_tokens = 0 + for micro_step, micro_batch in enumerate(micro_batches): + step_local_num_tokens += micro_batch["input_ids"].shape[1] + local_loss_sum, batch_token_count = compute_loss(micro_batch) + if gradient_manager is None: + step_local_token_count += batch_token_count + if torch.isnan(local_loss_sum.detach()): + nan_loss_count += 1 + logger.warning("Local loss is nan, excluding this micro step from backward") + scaled_loss = torch.nan_to_num(local_loss_sum, nan=0.0) / grad_accum_steps + else: + step_loss_sum += local_loss_sum.detach() + scaled_loss = local_loss_sum / grad_accum_steps + begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) + scaled_loss.backward() + finish_backward(gradient_manager) + + if is_moe_model: + micro_moe_stats, tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) + for name, value in micro_moe_stats.items(): + moe_stats[f"{name}/mean"] += value / grad_accum_steps + if name == "max_vio": + moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) + step_tokens_per_expert += tokens_per_expert + + forward_backward_time = time.perf_counter() - step_start_time + expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) if is_moe_model else {} + + if gradient_manager is None: + global_step_token_count = step_local_token_count.clone() + dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count_val = global_step_token_count.item() + grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val + scale_gradients_(None, model, grad_scale) + + dist.all_reduce(step_loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) + dist.all_reduce(nan_loss_count, op=dist.ReduceOp.SUM) + batch_loss = (step_loss_sum / global_token_count_val).item() + nan_loss_count = nan_loss_count.item() + + grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm) if config.optim.max_norm else None + optimizer.step() + optimizer.zero_grad() + current_lr = optimizer.param_groups[0]["lr"] + scheduler.step() + + # The last step's checkpoint is written once after the loop. + save_ckpt_time = 0.0 + if config.ckpt and config.ckpt.interval and progress.step % config.ckpt.interval == 0 and not is_last_step: + save_ckpt_start_time = time.perf_counter() + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) + ckpt_manager.maybe_clean() + save_ckpt_time = time.perf_counter() - save_ckpt_start_time + + eval_start_time = time.perf_counter() + if config.val is not None and (is_last_step or progress.step % config.val.interval == 0): + run_evals(progress.step) + eval_time = time.perf_counter() - eval_start_time + + global_num_tokens = torch.tensor(step_local_num_tokens, dtype=torch.int64, device="cuda") + dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) + num_tokens = global_num_tokens.item() + progress.total_tokens += num_tokens + dataset_progress = get_dataset_progress(dataloader) + progress.total_samples = dataset_progress["step"] + step_time = time.perf_counter() - step_start_time - eval_time + perf_counter = get_perf_counter(model, config.data.seq_len) + perf_counter.count_tokens(num_tokens) + throughput = perf_counter.get_step_tokens_per_second(num_tokens, step_time) + mfu = perf_counter.get_step_mfu(num_tokens, step_time) + peak_memory = torch.cuda.max_memory_reserved() / 1024**3 + + step_message = f"Step {progress.step} | {step_time:.2f}s | Loss {batch_loss:.4f}" + if grad_norm is not None: + step_message += f" | Grad. Norm {grad_norm:.4f}" + step_message += ( + f" | LR {current_lr:.2e} | Throughput {throughput:.0f} tokens/s | MFU {mfu:.1f}%" + f" | Peak Mem. {peak_memory:.1f}/{max_memory:.1f} GiB" + ) + logger.success(step_message) + + metrics = { + "loss/mean": batch_loss, + "loss/perplexity": math.exp(min(batch_loss, 20)), + "loss/nan_count": nan_loss_count, + "optim/lr": current_lr, + "progress/epoch": dataset_progress["epoch"], + "progress/num_samples": progress.total_samples, + "progress/num_tokens": progress.total_tokens, + "perf/throughput": throughput, + "perf/throughput_per_gpu": throughput / world.world_size, + "perf/peak_memory": peak_memory, + "perf/mfu": mfu, + "time/step": step_time, + "time/forward_backward": forward_backward_time, + "time/save_ckpt": save_ckpt_time, + "time/evals": eval_time, + **get_ckpt_disk_metrics(config.run_dir), + **{name: value.item() for name, value in moe_stats.items()}, + **expert_load_stats, + "step": progress.step, + } + if grad_norm is not None: + metrics["optim/grad_norm"] = grad_norm.item() + asyncio.run(monitors.log(metrics, step=progress.step)) + + if heart is not None: + heart.beat() + if is_last_step: + break + progress.step += 1 + + if config.ckpt is not None: + logger.info(f"Saving final checkpoint at step {progress.step}") + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) + ckpt_manager.maybe_clean() + if gradient_manager is not None: + gradient_manager.close() + logger.success("Pretrain trainer finished") + asyncio.run(monitors.finalize()) + + +def main(): + set_proc_title("PretrainTrainer") + train(cli(PretrainConfig)) + + +if __name__ == "__main__": + main() From 84ca355c80bc01a983b13cd1bdc0e984c21cfb0d Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 13 Oct 2026 19:32:39 +0200 Subject: [PATCH 003/132] docs(pretrain): Llama 3 ladder examples, docs and tests --- configs/debug/fake/pretrain.toml | 17 ++++++ docs/mint.json | 1 + docs/pretraining.md | 71 +++++++++++++++++++++++++ docs/training.md | 1 + examples/pretrain/llama3/150m.toml | 55 +++++++++++++++++++ examples/pretrain/llama3/17m.toml | 55 +++++++++++++++++++ examples/pretrain/llama3/1b.toml | 54 +++++++++++++++++++ examples/pretrain/llama3/500m.toml | 54 +++++++++++++++++++ examples/pretrain/llama3/70m.toml | 55 +++++++++++++++++++ tests/unit/test_configs.py | 2 + tests/unit/train/pretrain/test_data.py | 58 ++++++++++++++++++++ tests/unit/train/pretrain/test_evals.py | 12 +++++ tests/unit/train/pretrain/test_init.py | 21 ++++++++ 13 files changed, 456 insertions(+) create mode 100644 configs/debug/fake/pretrain.toml create mode 100644 docs/pretraining.md create mode 100644 examples/pretrain/llama3/150m.toml create mode 100644 examples/pretrain/llama3/17m.toml create mode 100644 examples/pretrain/llama3/1b.toml create mode 100644 examples/pretrain/llama3/500m.toml create mode 100644 examples/pretrain/llama3/70m.toml create mode 100644 tests/unit/train/pretrain/test_data.py create mode 100644 tests/unit/train/pretrain/test_evals.py create mode 100644 tests/unit/train/pretrain/test_init.py diff --git a/configs/debug/fake/pretrain.toml b/configs/debug/fake/pretrain.toml new file mode 100644 index 00000000000..60ba6332843 --- /dev/null +++ b/configs/debug/fake/pretrain.toml @@ -0,0 +1,17 @@ +max_steps = 5 + +[model] +name = "meta-llama/Llama-3.2-1B" +compile = "None" + +[model.arch] +hidden_size = 256 +num_hidden_layers = 2 +num_attention_heads = 4 +num_key_value_heads = 2 +intermediate_size = 1024 + +[data] +type = "fake" +batch_size = 4 +seq_len = 512 diff --git a/docs/mint.json b/docs/mint.json index d0f201b4e42..b75d50f8137 100644 --- a/docs/mint.json +++ b/docs/mint.json @@ -7,6 +7,7 @@ "overview", "configuration", "training", + "pretraining", "eval", "inference", "scaling", diff --git a/docs/pretraining.md b/docs/pretraining.md new file mode 100644 index 00000000000..8eed93f536b --- /dev/null +++ b/docs/pretraining.md @@ -0,0 +1,71 @@ +# Pretraining + +`uv run pretrain` trains a model from scratch on raw text. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, scheduler, checkpointing and monitors as the SFT trainer. It runs on a single node. + +```bash +uv run pretrain @ examples/pretrain/llama3/17m.toml +``` + +## Model + +`model.name` names a baseline HF model. The baseline provides the architecture class, the tokenizer and the vocabulary. `[model.arch]` overrides fields of its HF config to set the size: + +```toml +[model] +name = "meta-llama/Llama-3.2-1B" + +[model.arch] +hidden_size = 384 +num_hidden_layers = 8 +num_attention_heads = 6 +num_key_value_heads = 2 +head_dim = 64 +intermediate_size = 1536 +``` + +The launcher writes the resulting config and the tokenizer to `/model` and trains from there, so the trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all see an ordinary model directory. The trainer logs the non-embedding, embedding (input embedding and LM head) and total parameter counts at startup. + +The [Llama 3 ladder](https://github.com/PrimeIntellect-ai/prime-rl/tree/main/examples/pretrain/llama3) covers 17M, 70M, 150M, 500M and 1B non-embedding parameters. + +## Initialization + +Pretraining sets `model.debug.random_init = true`: weights are not downloaded or loaded, and `random_init_mode` resolves to `init`. Set `model.debug.random_init = false` to continue pretraining from the baseline's weights. + +| Field | Meaning | +|---|---| +| `model.init.scheme` | `normal`: every weight matrix and the embedding from N(0, std). `scaled`: additionally divides the std of the residual output projections (`o_proj`, `down_proj`) by sqrt(2 * num_layers). Norm weights start at 1, biases at 0. | +| `model.init.std` | Base standard deviation (default 0.02). | +| `model.init.seed` | Seed. Every rank initializes only its own shard, and the values depend only on the seed and the parameter name, not on the number of GPUs. | + +For `sft` and `rl`, `random_init_mode` resolves to `skip`: the weights keep arbitrary memory contents, which is the fastest option for throughput benchmarks. + +## Data + +`data.type = "text"` (the default) reads raw documents from a HF dataset column (`data.text_column`, default `text`). Only the selected subsets are downloaded. Each document is tokenized with the tokenizer's BOS and a closing EOS. Documents are packed back to back into rows of `data.seq_len` tokens and continue into the next row, so no token is truncated or padded. A step trains on exactly `data.batch_size * data.seq_len` tokens. + +| Field | Meaning | +|---|---| +| `data.document_masking` | Attention stays within each document and positions restart at every document start (default). If false, a row is one sequence. Attention never crosses a row. | +| `data.micro_batch_size` | Rows per forward pass. | +| `data.holdout_docs` | Hold the last N documents (in dataset order) out of training. With `[val]` set, they are evaluated as the `holdout` source. | +| `data.subsets`, `data.splits`, `data.probabilities` | Load and interleave several subsets/splits. | + +`data.type = "fake"` trains on generated tokens, for throughput tests. + +## Evals + +Evals run after the optimizer update, so the metrics at step N describe the weights after N updates, the same weights as the step-N checkpoint. With `eval_on_start` (default), they also run before the first step (logged as step 0), and they always run at the last step. They require `model.cp = 1`. + +**Perplexity (`[val]`).** Each `[[val.source]]` is a held-out text set. It is tokenized and packed once at startup (with document masking, up to `max_tokens`), so every eval sees the same tokens. Each source logs `val/{name}/loss` (nats per token), `val/{name}/ppl` and `val/{name}/bpb` (bits per UTF-8 byte, comparable across tokenizers). + +```toml +[val] +interval = 500 + +[[val.source]] +name = "wikitext" +dataset = "Salesforce/wikitext" +subset = "wikitext-103-raw-v1" +split = "test" +``` + diff --git a/docs/training.md b/docs/training.md index fbcddaec4e9..c3867a32e71 100644 --- a/docs/training.md +++ b/docs/training.md @@ -34,6 +34,7 @@ This page covers everything you need to launch, observe, checkpoint, and recover |---|---|---| | `uv run rl` | Wraps the trainer, orchestrator, and inference server in one launch from a merged TOML. | The default for any RL run. Runs locally for single-node experiments; submits to SLURM for single- or multi-node when `[slurm]` is set (see [Scaling § SLURM](scaling.md#slurm)). | | `uv run sft` | Supervised fine-tuning on a HF dataset. | Launches torchrun internally; never call torchrun directly. | +| `uv run pretrain` | Pretraining from scratch on raw text. | Single node; see [Pretraining](pretraining.md). | | `uv run inference` | vLLM server. | Always use this entrypoint over `vllm serve` — it adds `/update_weights`, `/load_lora_adapter`, and `/init_broadcaster`. | | `uv run orchestrator` | Standalone orchestrator process. | Pair with a separately-launched inference server and one `env-server` per source. | | `uv run eval` | Multi-env evals against a live inference server. | One epoch per source, pinned (or adaptive) concurrency, cursor checkpoints + `--resume`, dashboard + optional platform upload; see [Eval](eval.md). | diff --git a/examples/pretrain/llama3/150m.toml b/examples/pretrain/llama3/150m.toml new file mode 100644 index 00000000000..2bbe3c3a3b6 --- /dev/null +++ b/examples/pretrain/llama3/150m.toml @@ -0,0 +1,55 @@ +# Llama 3 architecture with ~147.1M non-embedding parameters, trained from scratch on FineWeb-Edu. +# 5,700 steps x 256 rows x 2048 tokens = 3.0B tokens (~20 tokens per non-embedding parameter). +# The learning rate is a starting point, not a tuned value. +# +# uv run pretrain @ examples/pretrain/llama3/150m.toml + +max_steps = 5700 + +[model] +name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) +offload = "none" +ac = "None" # small enough to keep all activations + +[model.arch] +hidden_size = 768 +num_hidden_layers = 17 +num_attention_heads = 12 +num_key_value_heads = 4 +head_dim = 64 +intermediate_size = 3072 +tie_word_embeddings = false +rope_parameters = { rope_type = "default", rope_theta = 500000.0 } + +[model.init] +scheme = "scaled" + +[data] +name = "HuggingFaceFW/fineweb-edu" +subsets = ["sample-10BT"] +seq_len = 2048 +batch_size = 256 +micro_batch_size = 8 +holdout_docs = 10000 + +[optim] +lr = 1e-3 +betas2 = 0.95 + +[scheduler] +type = "cosine" +warmup_steps = 300 +min_lr = 0.0001 + +[val] +interval = 285 + +[[val.source]] +name = "wikitext" +dataset = "Salesforce/wikitext" +subset = "wikitext-103-raw-v1" +split = "test" + + +[ckpt] +interval = 570 diff --git a/examples/pretrain/llama3/17m.toml b/examples/pretrain/llama3/17m.toml new file mode 100644 index 00000000000..c4694e6e265 --- /dev/null +++ b/examples/pretrain/llama3/17m.toml @@ -0,0 +1,55 @@ +# Llama 3 architecture with ~17.3M non-embedding parameters, trained from scratch on FineWeb-Edu. +# 700 steps x 256 rows x 2048 tokens = 0.4B tokens (~20 tokens per non-embedding parameter). +# The learning rate is a starting point, not a tuned value. +# +# uv run pretrain @ examples/pretrain/llama3/17m.toml + +max_steps = 700 + +[model] +name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) +offload = "none" +ac = "None" # small enough to keep all activations + +[model.arch] +hidden_size = 384 +num_hidden_layers = 8 +num_attention_heads = 6 +num_key_value_heads = 2 +head_dim = 64 +intermediate_size = 1536 +tie_word_embeddings = false +rope_parameters = { rope_type = "default", rope_theta = 500000.0 } + +[model.init] +scheme = "scaled" + +[data] +name = "HuggingFaceFW/fineweb-edu" +subsets = ["sample-10BT"] +seq_len = 2048 +batch_size = 256 +micro_batch_size = 16 +holdout_docs = 10000 + +[optim] +lr = 3e-3 +betas2 = 0.95 + +[scheduler] +type = "cosine" +warmup_steps = 50 +min_lr = 0.0003 + +[val] +interval = 100 + +[[val.source]] +name = "wikitext" +dataset = "Salesforce/wikitext" +subset = "wikitext-103-raw-v1" +split = "test" + + +[ckpt] +interval = 100 diff --git a/examples/pretrain/llama3/1b.toml b/examples/pretrain/llama3/1b.toml new file mode 100644 index 00000000000..9aa25f9e8bb --- /dev/null +++ b/examples/pretrain/llama3/1b.toml @@ -0,0 +1,54 @@ +# Llama 3 architecture with ~973.1M non-embedding parameters, trained from scratch on FineWeb-Edu. +# 37,200 steps x 256 rows x 2048 tokens = 19.5B tokens (~20 tokens per non-embedding parameter). +# The learning rate is a starting point, not a tuned value. +# +# uv run pretrain @ examples/pretrain/llama3/1b.toml + +max_steps = 37200 + +[model] +name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) +offload = "none" + +[model.arch] +hidden_size = 2048 +num_hidden_layers = 16 +num_attention_heads = 32 +num_key_value_heads = 8 +head_dim = 64 +intermediate_size = 8192 +tie_word_embeddings = false +rope_parameters = { rope_type = "default", rope_theta = 500000.0 } + +[model.init] +scheme = "scaled" + +[data] +name = "HuggingFaceFW/fineweb-edu" +subsets = ["sample-100BT"] +seq_len = 2048 +batch_size = 256 +micro_batch_size = 2 +holdout_docs = 10000 + +[optim] +lr = 4e-4 +betas2 = 0.95 + +[scheduler] +type = "cosine" +warmup_steps = 1850 +min_lr = 4e-05 + +[val] +interval = 1860 + +[[val.source]] +name = "wikitext" +dataset = "Salesforce/wikitext" +subset = "wikitext-103-raw-v1" +split = "test" + + +[ckpt] +interval = 3720 diff --git a/examples/pretrain/llama3/500m.toml b/examples/pretrain/llama3/500m.toml new file mode 100644 index 00000000000..c8169205f18 --- /dev/null +++ b/examples/pretrain/llama3/500m.toml @@ -0,0 +1,54 @@ +# Llama 3 architecture with ~495.5M non-embedding parameters, trained from scratch on FineWeb-Edu. +# 19,000 steps x 256 rows x 2048 tokens = 10.0B tokens (~20 tokens per non-embedding parameter). +# The learning rate is a starting point, not a tuned value. +# +# uv run pretrain @ examples/pretrain/llama3/500m.toml + +max_steps = 19000 + +[model] +name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) +offload = "none" + +[model.arch] +hidden_size = 1280 +num_hidden_layers = 21 +num_attention_heads = 20 +num_key_value_heads = 4 +head_dim = 64 +intermediate_size = 5120 +tie_word_embeddings = false +rope_parameters = { rope_type = "default", rope_theta = 500000.0 } + +[model.init] +scheme = "scaled" + +[data] +name = "HuggingFaceFW/fineweb-edu" +subsets = ["sample-100BT"] +seq_len = 2048 +batch_size = 256 +micro_batch_size = 4 +holdout_docs = 10000 + +[optim] +lr = 6e-4 +betas2 = 0.95 + +[scheduler] +type = "cosine" +warmup_steps = 950 +min_lr = 6e-05 + +[val] +interval = 950 + +[[val.source]] +name = "wikitext" +dataset = "Salesforce/wikitext" +subset = "wikitext-103-raw-v1" +split = "test" + + +[ckpt] +interval = 1900 diff --git a/examples/pretrain/llama3/70m.toml b/examples/pretrain/llama3/70m.toml new file mode 100644 index 00000000000..6352b7a95fe --- /dev/null +++ b/examples/pretrain/llama3/70m.toml @@ -0,0 +1,55 @@ +# Llama 3 architecture with ~70.8M non-embedding parameters, trained from scratch on FineWeb-Edu. +# 2,800 steps x 256 rows x 2048 tokens = 1.5B tokens (~20 tokens per non-embedding parameter). +# The learning rate is a starting point, not a tuned value. +# +# uv run pretrain @ examples/pretrain/llama3/70m.toml + +max_steps = 2800 + +[model] +name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) +offload = "none" +ac = "None" # small enough to keep all activations + +[model.arch] +hidden_size = 640 +num_hidden_layers = 12 +num_attention_heads = 10 +num_key_value_heads = 2 +head_dim = 64 +intermediate_size = 2560 +tie_word_embeddings = false +rope_parameters = { rope_type = "default", rope_theta = 500000.0 } + +[model.init] +scheme = "scaled" + +[data] +name = "HuggingFaceFW/fineweb-edu" +subsets = ["sample-10BT"] +seq_len = 2048 +batch_size = 256 +micro_batch_size = 8 +holdout_docs = 10000 + +[optim] +lr = 1.5e-3 +betas2 = 0.95 + +[scheduler] +type = "cosine" +warmup_steps = 150 +min_lr = 0.00015 + +[val] +interval = 140 + +[[val.source]] +name = "wikitext" +dataset = "Salesforce/wikitext" +subset = "wikitext-103-raw-v1" +split = "test" + + +[ckpt] +interval = 280 diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 2195f5f4ec2..aa9380f4c82 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -23,6 +23,7 @@ from prime_rl.configs.orchestrator import ( ModelConfig as OrchestratorModelConfig, ) +from prime_rl.configs.pretrain import PretrainConfig from prime_rl.configs.rl import RLConfig from prime_rl.configs.sft import SFTConfig from prime_rl.configs.trainer import ModelConfig as TrainerModelConfig @@ -34,6 +35,7 @@ RLConfig, TrainerConfig, SFTConfig, + PretrainConfig, OrchestratorConfig, InferenceConfig, EnvServerConfig, diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py new file mode 100644 index 00000000000..8f3f6ab7628 --- /dev/null +++ b/tests/unit/train/pretrain/test_data.py @@ -0,0 +1,58 @@ +from datasets import Dataset +from transformers import AutoTokenizer + +from prime_rl.trainer.pretrain.data import StreamPacker, TextDataset, cut_row, tokenize_document + + +def test_cut_row_segments_documents(): + tokens = list(range(9)) + row = cut_row(tokens, doc_starts=[-2, 3, 5], seq_len=8, document_masking=True) + assert row["input_ids"] == tokens[:8] + assert row["target_ids"] == tokens[1:] + assert row["seq_lens"] == [3, 2, 3] + assert row["position_ids"] == [0, 1, 2, 0, 1, 0, 1, 2] + unmasked = cut_row(tokens, doc_starts=[3, 5], seq_len=8, document_masking=False) + assert unmasked["seq_lens"] == [8] + assert unmasked["position_ids"] == list(range(8)) + + +def test_stream_packer_keeps_every_token(): + documents = [list(range(start, start + length)) for start, length in [(0, 5), (100, 9), (200, 4), (300, 30)]] + rows = list(StreamPacker(documents, seq_len=8, document_masking=True)) + stream = [token for document in documents for token in document] + assert len(rows) == (len(stream) - 1) // 8 + assert [token for row in rows for token in row["input_ids"]] == stream[: len(rows) * 8] + assert all(row["target_ids"][-1] == next_row["input_ids"][0] for row, next_row in zip(rows, rows[1:])) + assert rows[0]["seq_lens"] == [5, 3] + assert rows[1]["seq_lens"] == [6, 2] + + samples = list(StreamPacker(documents, seq_len=8, document_masking=False, rows_per_sample=2)) + assert [sample["seq_lens"] for sample in samples] == [[8, 8]] * len(samples) + assert samples[0]["input_ids"] == rows[0]["input_ids"] + rows[1]["input_ids"] + + +def test_stream_packer_resumes_mid_document(): + tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") + texts = [f"Document {i}: " + "word " * (7 * i + 3) for i in range(40)] + dataset = Dataset.from_list([{"text": text} for text in texts]) + + def packer(): + return StreamPacker(TextDataset(dataset, tokenizer, shuffle=True, seed=1), seq_len=16, document_masking=True) + + reference = packer() + reference_rows = iter(reference) + for _ in range(7): + next(reference_rows) + state = reference.state_dict() + expected = [next(reference_rows) for _ in range(20)] + + resumed = packer() + resumed.load_state_dict(state) + assert [row for row, _ in zip(resumed, range(20))] == expected + + +def test_tokenize_document_appends_eos(): + tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") + token_ids = tokenize_document(tokenizer, "Hello world") + assert token_ids[-1] == tokenizer.eos_token_id + assert token_ids.count(tokenizer.eos_token_id) == 1 diff --git a/tests/unit/train/pretrain/test_evals.py b/tests/unit/train/pretrain/test_evals.py new file mode 100644 index 00000000000..35c5dab7ac5 --- /dev/null +++ b/tests/unit/train/pretrain/test_evals.py @@ -0,0 +1,12 @@ +from transformers import AutoTokenizer + +from prime_rl.trainer.pretrain.evals.common import token_byte_lengths + + +def test_token_byte_lengths_count_utf8_bytes(): + tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") + lengths = token_byte_lengths(tokenizer) + text = "Grüße aus Zürich — 東京 🚀 naïve café" + token_ids = tokenizer(text, add_special_tokens=False)["input_ids"] + assert int(lengths[token_ids].sum()) == len(text.encode()) + assert lengths[tokenizer.eos_token_id] == 0 diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py new file mode 100644 index 00000000000..fd6ef68b9a3 --- /dev/null +++ b/tests/unit/train/pretrain/test_init.py @@ -0,0 +1,21 @@ +import itertools + +import torch + +from prime_rl.trainer.init import normal_shard + + +def test_normal_shard_is_independent_of_sharding(): + shape = (37, 5, 3) + full = normal_shard("w", shape, shape, (0, 0, 0), seed=0, device=torch.device("cpu")) + for parts in (2, 3, 8): + bounds = [len(chunk) for chunk in torch.arange(shape[0]).chunk(parts)] + offsets = [0, *itertools.accumulate(bounds)][:-1] + shards = [ + normal_shard("w", shape, (rows, *shape[1:]), (offset, 0, 0), seed=0, device=torch.device("cpu")) + for rows, offset in zip(bounds, offsets) + ] + assert torch.equal(torch.cat(shards), full) + column_shard = normal_shard("w", shape, (37, 2, 3), (0, 3, 0), seed=0, device=torch.device("cpu")) + assert torch.equal(column_shard, full[:, 3:5]) + assert not torch.equal(normal_shard("w", shape, shape, (0, 0, 0), seed=1, device=torch.device("cpu")), full) From 53bf7e39974259ae68413db445ca2db1dc0a94aa Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 13 Oct 2026 15:23:33 +0200 Subject: [PATCH 004/132] feat(pretrain): CompleteP parametrization and truncated normal init model.init.scheme = "completep" scales init, learning rate, Adam eps and weight decay per parameter role relative to a base shape (muP for width, depth with alpha = 1). The residual-branch and logit multipliers are folded into o_proj, down_proj and the LM head, so the model and its exports stay plain Llama. model.init.truncate draws from the normal truncated at 3 std. AdamW gets an eps field (default unchanged). --- docs/pretraining.md | 17 +++- .../src/prime_rl/configs/pretrain.py | 11 +++ .../src/prime_rl/configs/trainer.py | 16 +++- src/prime_rl/trainer/init.py | 37 +++++++-- src/prime_rl/trainer/optim/__init__.py | 1 + src/prime_rl/trainer/pretrain/completep.py | 32 ++++++++ src/prime_rl/trainer/pretrain/train.py | 27 ++++--- tests/unit/train/pretrain/test_init.py | 77 ++++++++++++++++++- 8 files changed, 198 insertions(+), 20 deletions(-) create mode 100644 src/prime_rl/trainer/pretrain/completep.py diff --git a/docs/pretraining.md b/docs/pretraining.md index 8eed93f536b..92c0b597abc 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -33,10 +33,25 @@ Pretraining sets `model.debug.random_init = true`: weights are not downloaded or | Field | Meaning | |---|---| -| `model.init.scheme` | `normal`: every weight matrix and the embedding from N(0, std). `scaled`: additionally divides the std of the residual output projections (`o_proj`, `down_proj`) by sqrt(2 * num_layers). Norm weights start at 1, biases at 0. | +| `model.init.scheme` | `normal`: every weight matrix and the embedding from N(0, std). `scaled`: additionally divides the std of the residual output projections (`o_proj`, `down_proj`) by sqrt(2 * num_layers). `completep`: CompleteP, see below. Norm weights start at 1, biases at 0. | | `model.init.std` | Base standard deviation (default 0.02). | +| `model.init.truncate` | Draw from the normal truncated at +-3 std. | | `model.init.seed` | Seed. Every rank initializes only its own shard, and the values depend only on the seed and the parameter name, not on the number of GPUs. | +### CompleteP + +`model.init.scheme = "completep"` trains with [CompleteP](https://arxiv.org/abs/2505.01618) (muP extended to depth, alpha = 1), so that hyperparameters tuned on a small base shape carry over to wider and deeper models. `optim.lr`, `optim.eps`, `optim.weight_decay` and `model.init.std` are the base-shape values, and the shape is set by `model.init.base_hidden_size` (default 256) and `model.init.base_num_layers` (default 2). With `m_N = hidden_size / base_hidden_size` and `m_L = num_layers / base_num_layers`: + +| Parameters | Init std | Learning rate | Adam eps | Weight decay | +|---|---|---|---|---| +| Embedding | std | lr | eps / m_N | wd | +| Hidden matrices | std / sqrt(m_N) | lr / m_N | eps / (m_N m_L) | wd m_N | +| `o_proj`, `down_proj` | std / (sqrt(m_N) m_L) | lr / (m_N m_L) | eps / m_N | wd m_N m_L | +| LM head | std / m_N | lr / m_N | eps | wd m_N | +| Norm weights | 1 | lr | eps / m_N | 0 | + +CompleteP scales each residual branch by 1/m_L and the logits by 1/m_N. These multipliers are folded into `o_proj`, `down_proj` and the LM head (the last three rows), so the model, its checkpoints and the exported weights are plain Llama. CompleteP needs AdamW, `model.offload = "none"`, and `scheduler.min_lr = 0` with the cosine scheduler. Keep `head_dim` and the ratios `intermediate_size / hidden_size` and `num_attention_heads / num_key_value_heads` fixed across shapes. + For `sft` and `rl`, `random_init_mode` resolves to `skip`: the weights keep arbitrary memory contents, which is the fastest option for throughput benchmarks. ## Data diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 71a896d315e..53233acc1c5 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -232,6 +232,17 @@ def auto_setup_tokenizer(self): self.tokenizer.trust_remote_code = self.model.trust_remote_code return self + @model_validator(mode="after") + def validate_completep(self): + if self.model.init.scheme != "completep": + return self + if self.optim.type != "adamw" or self.model.offload != "none": + raise ValueError("model.init.scheme = 'completep' needs optim.type = 'adamw' and model.offload = 'none'.") + if self.scheduler.type == "cosine" and self.scheduler.min_lr > 0: + # CosineAnnealingLR decays every parameter group to the same absolute min_lr. + raise ValueError("model.init.scheme = 'completep' needs scheduler.min_lr = 0 with the cosine scheduler.") + return self + @model_validator(mode="after") def validate_monitors(self): if self.monitors.prime is not None: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 63b2364e8db..021c1c93a1d 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -236,12 +236,21 @@ class MoERuntimeConfig(BaseConfig): class InitConfig(BaseConfig): """Random weight initialization, used when ``model.debug.random_init`` resolves to ``init``.""" - scheme: Literal["normal", "scaled"] = "normal" - """``normal`` draws every weight matrix and the embedding from N(0, std). ``scaled`` additionally divides the std of the residual output projections (``o_proj``, ``down_proj``) by sqrt(2 * num_layers). Norm weights start at 1, biases at 0.""" + scheme: Literal["normal", "scaled", "completep"] = "normal" + """``normal`` draws every weight matrix and the embedding from N(0, std). ``scaled`` additionally divides the std of the residual output projections (``o_proj``, ``down_proj``) by sqrt(2 * num_layers). ``completep`` follows CompleteP (muP extended to depth, Dey et al. 2025) relative to the base shape: the embedding keeps ``std``, hidden matrices get ``std / sqrt(m_N)``, residual output projections ``std / (sqrt(m_N) * m_L)`` and the LM head ``std / m_N``, where ``m_N = hidden_size / base_hidden_size`` and ``m_L = num_layers / base_num_layers``. Norm weights start at 1, biases at 0.""" std: float = Field(0.02, gt=0) """Base standard deviation.""" + truncate: bool = False + """Draw from the normal distribution truncated at +-3 std instead of the full normal.""" + + base_hidden_size: int = Field(256, ge=1) + """Hidden size of the base shape. Only used by ``completep``.""" + + base_num_layers: int = Field(2, ge=1) + """Number of layers of the base shape. Only used by ``completep``.""" + seed: int = 0 """Seed for the initialization. The values depend only on the seed, not on the number of GPUs.""" @@ -438,6 +447,9 @@ class AdamWConfig(BaseOptimizerConfig): betas2: float = Field(0.999, ge=0) """Adam second-moment (β2) decay.""" + eps: float = Field(1e-8, gt=0) + """Adam epsilon.""" + class MuonConfig(BaseOptimizerConfig): type: Literal["muon"] = "muon" diff --git a/src/prime_rl/trainer/init.py b/src/prime_rl/trainer/init.py index 3d2d49f4c0c..94f5ace16d1 100644 --- a/src/prime_rl/trainer/init.py +++ b/src/prime_rl/trainer/init.py @@ -35,7 +35,18 @@ def param_role(name: str, param: torch.Tensor) -> ParamRole: return "hidden" -def role_std(role: ParamRole, config: InitConfig, num_layers: int) -> float: +def completep_multipliers(config: InitConfig, hidden_size: int, num_layers: int) -> tuple[float, float]: + """Width and depth multipliers ``(m_N, m_L)`` of the model relative to the CompleteP base shape.""" + return hidden_size / config.base_hidden_size, num_layers / config.base_num_layers + + +def role_std(role: ParamRole, config: InitConfig, hidden_size: int, num_layers: int) -> float: + if config.scheme == "completep": + # CompleteP's 1/m_L residual-branch and 1/m_N logit multipliers are folded into the + # residual output projections and the LM head, so the forward pass stays unchanged. + m_n, m_l = completep_multipliers(config, hidden_size, num_layers) + scale = {"embedding": 1.0, "hidden": m_n**-0.5, "residual_out": m_n**-0.5 / m_l, "output": 1 / m_n} + return config.std * scale[role] if config.scheme == "scaled" and role == "residual_out": return config.std / math.sqrt(2 * num_layers) return config.std @@ -53,10 +64,12 @@ def normal_shard( global_offset: tuple[int, ...], seed: int, device: torch.device, + truncate: bool = False, ) -> torch.Tensor: """The ``local_shape`` slice at ``global_offset`` of a standard-normal tensor of ``global_shape`` whose values depend only on ``(seed, name)``. The tensor is drawn in - blocks of rows along dim 0; only the blocks overlapping the slice are generated.""" + blocks of rows along dim 0; only the blocks overlapping the slice are generated. With + ``truncate``, the distribution is truncated at +-3.""" row_numel = math.prod(global_shape[1:]) rows_per_block = max(1, BLOCK_NUMEL // max(row_numel, 1)) start, stop = global_offset[0], global_offset[0] + local_shape[0] @@ -65,7 +78,13 @@ def normal_shard( block_start = block * rows_per_block block_rows = min(rows_per_block, global_shape[0] - block_start) generator = torch.Generator(device=device).manual_seed(block_seed(seed, name, block)) - values = torch.randn((block_rows, *global_shape[1:]), generator=generator, device=device) + if truncate: + # Inverse-CDF sampling: uniform on (erf(-3/sqrt(2)), erf(3/sqrt(2))) mapped through erfinv. + bound = math.erf(3 / math.sqrt(2)) + uniform = torch.rand((block_rows, *global_shape[1:]), generator=generator, device=device) + values = torch.erfinv(uniform.mul_(2 * bound).sub_(bound)).mul_(math.sqrt(2)) + else: + values = torch.randn((block_rows, *global_shape[1:]), generator=generator, device=device) lo, hi = max(start, block_start) - block_start, min(stop, block_start + block_rows) - block_start pieces.append(values[lo:hi]) out = torch.cat(pieces) if pieces else torch.empty((0, *global_shape[1:]), device=device) @@ -78,7 +97,7 @@ def normal_shard( def init_weights(model: nn.Module, config: InitConfig) -> None: """Initialize every parameter in place. Each rank fills only its local shard, without communication, and the values do not depend on the world size or the sharding.""" - num_layers = model.config.num_hidden_layers + hidden_size, num_layers = model.config.hidden_size, model.config.num_hidden_layers for name, param in model.named_parameters(): # Canonical names, so wrappers such as activation checkpointing don't change the values. name = resolve_fqn(model, name) @@ -97,6 +116,12 @@ def init_weights(model: nn.Module, config: InitConfig) -> None: else: local_shape, global_offset = tuple(param.shape), (0,) * param.ndim values = normal_shard( - name, tuple(param.shape), tuple(local_shape), tuple(global_offset), config.seed, local.device + name, + tuple(param.shape), + tuple(local_shape), + tuple(global_offset), + config.seed, + local.device, + config.truncate, ) - local.copy_(values.mul_(role_std(role, config, num_layers))) + local.copy_(values.mul_(role_std(role, config, hidden_size, num_layers))) diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index a50bd64233b..054ae16a901 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/src/prime_rl/trainer/optim/__init__.py @@ -121,6 +121,7 @@ def _create_optimizer( lr=lr, weight_decay=config.weight_decay, betas=(config.betas1, config.betas2), + eps=config.eps, fused=fused_adamw, ) case "muon": diff --git a/src/prime_rl/trainer/pretrain/completep.py b/src/prime_rl/trainer/pretrain/completep.py new file mode 100644 index 00000000000..18b73b95bc9 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/completep.py @@ -0,0 +1,32 @@ +from torch import nn + +from prime_rl.configs.trainer import AdamWConfig, InitConfig +from prime_rl.trainer.init import completep_multipliers, param_role +from prime_rl.utils.weights import resolve_fqn + + +def completep_param_groups(model: nn.Module, init: InitConfig, optim: AdamWConfig) -> list[dict]: + """AdamW parameter groups for CompleteP with alpha = 1 (Dey et al. 2025). ``optim.lr``, + ``optim.eps`` and ``optim.weight_decay`` are the values for the base shape; each role scales + them by the width and depth multipliers. The residual output projections and the LM head + carry CompleteP's 1/m_L and 1/m_N forward multipliers folded into their weights (see + ``role_std``), which turns their rules into the ones below.""" + m_n, m_l = completep_multipliers(init, model.config.hidden_size, model.config.num_hidden_layers) + # (lr, eps, weight decay) scale per role. PyTorch AdamW decays by lr * weight_decay. + scales = { + "embedding": (1.0, 1 / m_n, 1.0), + "hidden": (1 / m_n, 1 / (m_n * m_l), m_n), + "residual_out": (1 / (m_n * m_l), 1 / m_n, m_n * m_l), + "output": (1 / m_n, 1.0, m_n), + "norm": (1.0, 1 / m_n, 0.0), + "bias": (1.0, 1 / m_n, 0.0), + } + params = {role: [] for role in scales} + for name, param in model.named_parameters(): + if param.requires_grad: + params[param_role(resolve_fqn(model, name), param)].append(param) + return [ + dict(params=params[role], lr=optim.lr * lr, eps=optim.eps * eps, weight_decay=optim.weight_decay * wd) + for role, (lr, eps, wd) in scales.items() + if params[role] + ] diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index a29bf227c4b..3208a7dee3b 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -11,6 +11,7 @@ import torch import torch.distributed as dist from torch import nn +from torch.optim import AdamW from transformers import AutoConfig from prime_rl import monitors @@ -32,6 +33,7 @@ from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter +from prime_rl.trainer.pretrain.completep import completep_param_groups from prime_rl.trainer.pretrain.data import load_text_dataset, setup_text_dataloader, split_holdout from prime_rl.trainer.pretrain.evals.common import token_byte_lengths from prime_rl.trainer.pretrain.evals.perplexity import build_perplexity_set, evaluate_perplexity, load_perplexity_set @@ -141,16 +143,21 @@ def train(config: PretrainConfig): tokenizer = setup_tokenizer(config.tokenizer) - optimizer, gradient_manager = setup_optimizer( - config.optim, - list(model.named_parameters()), - parallel_dims, - offload=config.model.offload, - model=model, - full_offload_dtype_policy=( - get_full_offload_dtype_policy(model, config.model) if config.model.offload == "full" else None - ), - ) + if config.model.init.scheme == "completep": + param_groups = completep_param_groups(model, config.model.init, config.optim) + optimizer = AdamW(param_groups, betas=(config.optim.betas1, config.optim.betas2), fused=True) + gradient_manager = None + else: + optimizer, gradient_manager = setup_optimizer( + config.optim, + list(model.named_parameters()), + parallel_dims, + offload=config.model.offload, + model=model, + full_offload_dtype_policy=( + get_full_offload_dtype_policy(model, config.model) if config.model.offload == "full" else None + ), + ) scheduler_steps = ( config.max_steps - checkpoint_step diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index fd6ef68b9a3..e8935d21e8a 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -1,8 +1,13 @@ +import copy import itertools import torch +from torch.optim import AdamW +from transformers import LlamaConfig, LlamaForCausalLM -from prime_rl.trainer.init import normal_shard +from prime_rl.configs.trainer import AdamWConfig, InitConfig +from prime_rl.trainer.init import init_weights, normal_shard +from prime_rl.trainer.pretrain.completep import completep_param_groups def test_normal_shard_is_independent_of_sharding(): @@ -19,3 +24,73 @@ def test_normal_shard_is_independent_of_sharding(): column_shard = normal_shard("w", shape, (37, 2, 3), (0, 3, 0), seed=0, device=torch.device("cpu")) assert torch.equal(column_shard, full[:, 3:5]) assert not torch.equal(normal_shard("w", shape, shape, (0, 0, 0), seed=1, device=torch.device("cpu")), full) + + +def test_truncated_normal_shard_stays_within_three_std(): + values = normal_shard("w", (4096, 64), (4096, 64), (0, 0), seed=0, device=torch.device("cpu"), truncate=True) + assert values.abs().max() <= 3 + # Variance of the standard normal truncated at +-3. + assert abs(values.std().item() - 0.9866) < 0.01 + + +def test_completep_folded_multipliers_match_explicit_multipliers(): + """CompleteP folds its 1/m_L residual and 1/m_N logit multipliers into the weights. Training + the folded model must match the multiplier form of CompleteP (Dey et al. 2025, Table 1).""" + torch.manual_seed(0) + config = LlamaConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + tie_word_embeddings=False, + ) + init = InitConfig(scheme="completep", base_hidden_size=32, base_num_layers=2) + optim = AdamWConfig(lr=1e-2, eps=1e-10, weight_decay=0.1) + m_n, m_l = 2.0, 2.0 + + folded = LlamaForCausalLM(config) + init_weights(folded, init) + explicit = copy.deepcopy(folded) + with torch.no_grad(): + for name, param in explicit.named_parameters(): + if name.endswith(("o_proj.weight", "down_proj.weight")): + param.mul_(m_l) + elif "lm_head" in name: + param.mul_(m_n) + for layer in explicit.model.layers: + layer.self_attn.o_proj.register_forward_hook(lambda module, args, out: out / m_l) + layer.mlp.down_proj.register_forward_hook(lambda module, args, out: out / m_l) + explicit.lm_head.register_forward_pre_hook(lambda module, args: (args[0] / m_n,)) + + def explicit_scales(name: str) -> tuple[float, float, float]: + if "embed_tokens" in name or "lm_head" in name: + return 1.0, 1 / m_n, 1.0 + if "norm" in name: + return 1.0, 1 / m_n, 0.0 + return 1 / m_n, 1 / (m_n * m_l), m_n + + explicit_groups = [] + for name, param in explicit.named_parameters(): + lr, eps, wd = explicit_scales(name) + explicit_groups.append( + dict(params=[param], lr=optim.lr * lr, eps=optim.eps * eps, weight_decay=optim.weight_decay * wd) + ) + optimizers = { + "explicit": AdamW(explicit_groups, betas=(0.9, 0.95)), + "folded": AdamW(completep_param_groups(folded, init, optim), betas=(0.9, 0.95)), + } + models = {"explicit": explicit, "folded": folded} + + tokens = torch.randint(0, 128, (4, 16)) + for _ in range(5): + for key, model in models.items(): + model(input_ids=tokens, labels=tokens).loss.backward() + optimizers[key].step() + optimizers[key].zero_grad() + with torch.no_grad(): + torch.testing.assert_close( + folded(input_ids=tokens).logits, explicit(input_ids=tokens).logits, rtol=1e-4, atol=1e-4 + ) From f2187de2358771c392966b244223c92df648dd96 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 13 Oct 2026 18:11:11 +0200 Subject: [PATCH 005/132] feat(pretrain): CompleteP with Muon setup_optimizer accepts custom parameter groups for AdamW and Muon, so the CompleteP groups go through the shared optimizer setup (fused-qkv partitions, meshes, Muon warm-up). With Muon, the hidden matrices and residual output projections are Muon groups whose learning rate is scaled by 1/sqrt(m_N) to offset the sqrt(m_N) growth of the rms_norm update. MuonConfig gets an eps field for its AdamW groups (default unchanged). --- docs/pretraining.md | 11 +++- .../src/prime_rl/configs/pretrain.py | 4 +- .../src/prime_rl/configs/trainer.py | 3 + src/prime_rl/trainer/optim/__init__.py | 39 +++++++++---- src/prime_rl/trainer/pretrain/completep.py | 55 +++++++++++++++---- src/prime_rl/trainer/pretrain/train.py | 31 +++++------ 6 files changed, 102 insertions(+), 41 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 92c0b597abc..5bd89dfd60a 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -50,7 +50,16 @@ Pretraining sets `model.debug.random_init = true`: weights are not downloaded or | LM head | std / m_N | lr / m_N | eps | wd m_N | | Norm weights | 1 | lr | eps / m_N | 0 | -CompleteP scales each residual branch by 1/m_L and the logits by 1/m_N. These multipliers are folded into `o_proj`, `down_proj` and the LM head (the last three rows), so the model, its checkpoints and the exported weights are plain Llama. CompleteP needs AdamW, `model.offload = "none"`, and `scheduler.min_lr = 0` with the cosine scheduler. Keep `head_dim` and the ratios `intermediate_size / hidden_size` and `num_attention_heads / num_key_value_heads` fixed across shapes. +CompleteP scales each residual branch by 1/m_L and the logits by 1/m_N. These multipliers are folded into `o_proj`, `down_proj` and the LM head, so the model, its checkpoints and the exported weights are plain Llama. + +With `optim.type = "muon"`, the hidden matrices and `o_proj`/`down_proj` are trained by Muon (`rms_norm` learning-rate adjustment), the rest by AdamW with the rules above. Muon's `rms_norm` update grows as sqrt(m_N), so its groups use: + +| Parameters | Init std | Learning rate | Weight decay | +|---|---|---|---| +| Hidden matrices | std / sqrt(m_N) | lr / sqrt(m_N) | wd sqrt(m_N) | +| `o_proj`, `down_proj` | std / (sqrt(m_N) m_L) | lr / (sqrt(m_N) m_L) | wd sqrt(m_N) m_L | + +CompleteP works with dense models, needs `model.offload = "none"`, and `scheduler.min_lr = 0` with the cosine scheduler. Keep `head_dim` and the ratios `intermediate_size / hidden_size` and `num_attention_heads / num_key_value_heads` fixed across shapes. For `sft` and `rl`, `random_init_mode` resolves to `skip`: the weights keep arbitrary memory contents, which is the fastest option for throughput benchmarks. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 53233acc1c5..35f46164271 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -236,8 +236,8 @@ def auto_setup_tokenizer(self): def validate_completep(self): if self.model.init.scheme != "completep": return self - if self.optim.type != "adamw" or self.model.offload != "none": - raise ValueError("model.init.scheme = 'completep' needs optim.type = 'adamw' and model.offload = 'none'.") + if self.optim.type not in ("adamw", "muon") or self.model.offload != "none": + raise ValueError("model.init.scheme = 'completep' needs AdamW or Muon and model.offload = 'none'.") if self.scheduler.type == "cosine" and self.scheduler.min_lr > 0: # CosineAnnealingLR decays every parameter group to the same absolute min_lr. raise ValueError("model.init.scheme = 'completep' needs scheduler.min_lr = 0 with the cosine scheduler.") diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 021c1c93a1d..f7333ef9997 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -463,6 +463,9 @@ class MuonConfig(BaseOptimizerConfig): betas2: float = Field(0.95, ge=0) """β2 for the AdamW/Lion sub-optimizer used on non-Muon params.""" + eps: float = Field(1e-8, gt=0) + """Adam epsilon for the AdamW sub-optimizer used on non-Muon params.""" + class SignSGDConfig(BaseOptimizerConfig): type: Literal["sign_sgd"] = "sign_sgd" diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index 054ae16a901..a47913ff7bb 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/src/prime_rl/trainer/optim/__init__.py @@ -54,9 +54,14 @@ def setup_optimizer( offload: OffloadMode = "none", model: nn.Module | None = None, full_offload_dtype_policy: dict[int, tuple[torch.dtype, torch.dtype]] | None = None, + param_groups: list[dict] | None = None, ) -> tuple[OptimizerLike, GradientOffloadManager | None]: + """``param_groups`` replaces the default parameter groups (for Muon, each group needs an + ``algorithm`` key). It is not supported with CPU offloading.""" cpu_offload = offload == "optimizer" full_offload = offload == "full" + if param_groups is not None and offload != "none": + raise ValueError("Custom parameter groups are not supported with CPU offloading") if full_offload and config.type not in ("adamw", "sign_sgd"): raise ValueError("Full optimizer offload only supports AdamW and SignSGD") if full_offload and config.max_norm is not None: @@ -79,6 +84,7 @@ def setup_optimizer( parallel_dims, fused_adamw=config.type == "adamw" and not cpu_offload, model=model, + param_groups=param_groups, ) if full_offload: @@ -105,6 +111,7 @@ def _create_optimizer( lr: float | None = None, fused_adamw: bool = False, model: nn.Module | None = None, + param_groups: list[dict] | None = None, ) -> Optimizer: """Create optimizer. If lr is None, uses config.lr.""" if lr is None: @@ -117,7 +124,7 @@ def _create_optimizer( match config.type: case "adamw": return AdamW( - params=trainable_params, + params=param_groups or trainable_params, lr=lr, weight_decay=config.weight_decay, betas=(config.betas1, config.betas2), @@ -125,7 +132,7 @@ def _create_optimizer( fused=fused_adamw, ) case "muon": - return _create_muon_optimizer(config, named_params, parallel_dims, model, lr) + return _create_muon_optimizer(config, named_params, parallel_dims, model, lr, param_groups) case "sign_sgd": return SignSGD( params=trainable_params, @@ -134,13 +141,12 @@ def _create_optimizer( ) -def _create_muon_optimizer( +def _default_muon_param_groups( config: OptimizerConfig, named_params: list[tuple[str, nn.Parameter]], parallel_dims: ParallelDims, - model: nn.Module | None, - lr: float | None = None, -) -> Optimizer: + lr: float, +) -> list[dict]: def muon_enabled(n, p): if p.ndim < 2: return False @@ -164,8 +170,6 @@ def muon_enabled(n, p): muon_params.append(p) elif p.requires_grad: adamw_params.append(p) - else: - pass param_groups = [] @@ -197,7 +201,22 @@ def muon_enabled(n, p): ) ) - param_groups.append(dict(params=adamw_params, algorithm="adamw", lr=lr, weight_decay=config.weight_decay)) + param_groups.append( + dict(params=adamw_params, algorithm="adamw", lr=lr, weight_decay=config.weight_decay, epsilon=config.eps) + ) + return param_groups + + +def _create_muon_optimizer( + config: OptimizerConfig, + named_params: list[tuple[str, nn.Parameter]], + parallel_dims: ParallelDims, + model: nn.Module | None, + lr: float | None = None, + param_groups: list[dict] | None = None, +) -> Optimizer: + if param_groups is None: + param_groups = _default_muon_param_groups(config, named_params, parallel_dims, lr) if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled: distributed_mesh = parallel_dims.get_mesh("dp_shard_cp") @@ -230,6 +249,6 @@ def muon_enabled(n, p): # Keep both warm-ups after Muon construction and before its first step. The # main and expert groups establish independent NCCL peer connections. _warmup_muon_mesh(distributed_mesh) - if expert_params and parallel_dims.ep_enabled: + if any(group.get("distributed_mesh_name") == "dp_shard_mod_ep" for group in param_groups): _warmup_muon_mesh(parallel_dims.get_mesh("dp_shard_mod_ep")) return optimizer diff --git a/src/prime_rl/trainer/pretrain/completep.py b/src/prime_rl/trainer/pretrain/completep.py index 18b73b95bc9..8068af3b4cc 100644 --- a/src/prime_rl/trainer/pretrain/completep.py +++ b/src/prime_rl/trainer/pretrain/completep.py @@ -1,19 +1,27 @@ from torch import nn -from prime_rl.configs.trainer import AdamWConfig, InitConfig +from prime_rl.configs.trainer import AdamWConfig, InitConfig, MuonConfig from prime_rl.trainer.init import completep_multipliers, param_role from prime_rl.utils.weights import resolve_fqn -def completep_param_groups(model: nn.Module, init: InitConfig, optim: AdamWConfig) -> list[dict]: - """AdamW parameter groups for CompleteP with alpha = 1 (Dey et al. 2025). ``optim.lr``, +def completep_param_groups(model: nn.Module, init: InitConfig, optim: AdamWConfig | MuonConfig) -> list[dict]: + """Optimizer parameter groups for CompleteP with alpha = 1 (Dey et al. 2025). ``optim.lr``, ``optim.eps`` and ``optim.weight_decay`` are the values for the base shape; each role scales them by the width and depth multipliers. The residual output projections and the LM head carry CompleteP's 1/m_L and 1/m_N forward multipliers folded into their weights (see - ``role_std``), which turns their rules into the ones below.""" + ``role_std``), which turns their rules into the ones below. + + With Muon, the hidden matrices and residual output projections are Muon groups with + ``rms_norm`` learning-rate adjustment, whose update size grows as sqrt(m_N); their learning + rate is scaled by 1/sqrt(m_N) to compensate (and weight decay by sqrt(m_N), keeping + lr * weight_decay fixed). Everything else is trained with AdamW, as with ``optim.type = "adamw"``. + """ + if any("mlp.experts" in name for name, _ in model.named_parameters()): + raise ValueError("CompleteP is only implemented for dense models") m_n, m_l = completep_multipliers(init, model.config.hidden_size, model.config.num_hidden_layers) - # (lr, eps, weight decay) scale per role. PyTorch AdamW decays by lr * weight_decay. - scales = { + # (lr, eps, weight decay) scale per role for AdamW. PyTorch AdamW and Dion decay by lr * weight_decay. + adam_scales = { "embedding": (1.0, 1 / m_n, 1.0), "hidden": (1 / m_n, 1 / (m_n * m_l), m_n), "residual_out": (1 / (m_n * m_l), 1 / m_n, m_n * m_l), @@ -21,12 +29,35 @@ def completep_param_groups(model: nn.Module, init: InitConfig, optim: AdamWConfi "norm": (1.0, 1 / m_n, 0.0), "bias": (1.0, 1 / m_n, 0.0), } - params = {role: [] for role in scales} + # (lr, weight decay) scale per role for Muon with rms_norm adjustment. + muon_scales = {"hidden": (m_n**-0.5, m_n**0.5), "residual_out": (m_n**-0.5 / m_l, m_n**0.5 * m_l)} + + params = {role: [] for role in adam_scales} for name, param in model.named_parameters(): if param.requires_grad: params[param_role(resolve_fqn(model, name), param)].append(param) - return [ - dict(params=params[role], lr=optim.lr * lr, eps=optim.eps * eps, weight_decay=optim.weight_decay * wd) - for role, (lr, eps, wd) in scales.items() - if params[role] - ] + + groups = [] + for role, role_params in params.items(): + if not role_params: + continue + if optim.type == "muon" and role in muon_scales: + lr, wd = muon_scales[role] + groups.append( + dict( + params=role_params, + algorithm="muon", + adjust_lr="rms_norm", + lr=optim.lr * lr, + weight_decay=optim.weight_decay * wd, + ) + ) + continue + lr, eps, wd = adam_scales[role] + group = dict(params=role_params, lr=optim.lr * lr, weight_decay=optim.weight_decay * wd) + if optim.type == "muon": + group.update(algorithm="adamw", epsilon=optim.eps * eps) + else: + group["eps"] = optim.eps * eps + groups.append(group) + return groups diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 3208a7dee3b..ca86c26da53 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -11,7 +11,6 @@ import torch import torch.distributed as dist from torch import nn -from torch.optim import AdamW from transformers import AutoConfig from prime_rl import monitors @@ -143,21 +142,21 @@ def train(config: PretrainConfig): tokenizer = setup_tokenizer(config.tokenizer) - if config.model.init.scheme == "completep": - param_groups = completep_param_groups(model, config.model.init, config.optim) - optimizer = AdamW(param_groups, betas=(config.optim.betas1, config.optim.betas2), fused=True) - gradient_manager = None - else: - optimizer, gradient_manager = setup_optimizer( - config.optim, - list(model.named_parameters()), - parallel_dims, - offload=config.model.offload, - model=model, - full_offload_dtype_policy=( - get_full_offload_dtype_policy(model, config.model) if config.model.offload == "full" else None - ), - ) + optimizer, gradient_manager = setup_optimizer( + config.optim, + list(model.named_parameters()), + parallel_dims, + offload=config.model.offload, + model=model, + full_offload_dtype_policy=( + get_full_offload_dtype_policy(model, config.model) if config.model.offload == "full" else None + ), + param_groups=( + completep_param_groups(model, config.model.init, config.optim) + if config.model.init.scheme == "completep" + else None + ), + ) scheduler_steps = ( config.max_steps - checkpoint_step From b560be3acafed50645645f596ee92898d206db4f Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 13 Oct 2026 19:41:46 +0200 Subject: [PATCH 006/132] feat(pretrain): per-hyperparameter multiplier schedules --- docs/pretraining.md | 33 ++++- examples/pretrain/llama3/150m.toml | 6 +- examples/pretrain/llama3/17m.toml | 6 +- examples/pretrain/llama3/1b.toml | 6 +- examples/pretrain/llama3/500m.toml | 6 +- examples/pretrain/llama3/70m.toml | 6 +- .../src/prime_rl/configs/pretrain.py | 112 ++++++++++++++- .../src/prime_rl/configs/trainer.py | 34 ++++- src/prime_rl/trainer/ckpt.py | 4 +- src/prime_rl/trainer/pretrain/schedules.py | 134 ++++++++++++++++++ src/prime_rl/trainer/pretrain/train.py | 29 ++-- tests/unit/train/pretrain/test_schedules.py | 87 ++++++++++++ 12 files changed, 425 insertions(+), 38 deletions(-) create mode 100644 src/prime_rl/trainer/pretrain/schedules.py create mode 100644 tests/unit/train/pretrain/test_schedules.py diff --git a/docs/pretraining.md b/docs/pretraining.md index 8eed93f536b..54b56f5cbc4 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -1,6 +1,6 @@ # Pretraining -`uv run pretrain` trains a model from scratch on raw text. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, scheduler, checkpointing and monitors as the SFT trainer. It runs on a single node. +`uv run pretrain` trains a model from scratch on raw text. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, checkpointing and monitors as the SFT trainer. It runs on a single node. ```bash uv run pretrain @ examples/pretrain/llama3/17m.toml @@ -39,6 +39,37 @@ Pretraining sets `model.debug.random_init = true`: weights are not downloaded or For `sft` and `rl`, `random_init_mode` resolves to `skip`: the weights keep arbitrary memory contents, which is the fastest option for throughput benchmarks. +## Schedules + +A schedule is a multiplier `f(step)`, where `step` counts the optimizer updates completed so far (0 for the first update). Update `step` uses `base * f(step)` for every param group, with `base` the group's own configured value, so ratios between param groups hold at every step. Each optimizer hyperparameter gets its own schedule under `[schedules.optim.]`, keyed by its `optim` field name: `lr`, `weight_decay`, `betas1`, `betas2`, and `mu` for Muon. `[schedules.clip_norm]` scales `optim.max_norm`. Hyperparameters without a schedule stay constant. The default is a WSD schedule on `lr`; setting `schedules.optim` replaces it, so include `lr` when scheduling other hyperparameters. + +```toml +[schedules.optim.lr] +type = "wsd" +warmup = 0.05 # fraction of max_steps, or a number of steps +decay = 0.2 +end_ratio = 0.1 +decay_shape = "1-sqrt" + +[schedules.optim.betas2] +type = "cosine" +warmup = 0 +end_ratio = 1.05 +``` + +| `type` | Shape | +|---|---| +| `wsd` (default) | Linear warmup, stable at 1, then a `linear`, `cosine` or `1-sqrt` decay over the last `decay` updates. Defaults: `warmup = 0.1`, `decay = 0.2`. | +| `cosine` | Linear warmup, then cosine decay over all remaining updates. | +| `constant` | 1 at every step. | +| `custom` | `fn(step, total_steps=max_steps, **kwargs) -> float` imported from `import_path`, with `kwargs` from the config. | + +The configured endpoints are exactly what the first and last updates use: `f(0) = start_ratio`, the warmup reaches 1 at its last update (`warmup - 1`; a warmup of 0 or 1 starts at 1), and the decay reaches `end_ratio` at the last update (`max_steps - 1`). Both ratios default to 0. Fractions resolve against `max_steps`, so fractional phases and decays need `max_steps`. + +AdamW and Muon apply weight decay as `lr * weight_decay`, so a learning rate schedule already scales the decay. + +The values are recomputed from the step, so checkpoints hold no schedule state. Resuming with `ckpt.skip_scheduler = true` restarts every schedule at the resumed step, spanning the remaining steps. Each scheduled value is logged as `optim/` (the first param group's value) and its multiplier as `schedule/`. + ## Data `data.type = "text"` (the default) reads raw documents from a HF dataset column (`data.text_column`, default `text`). Only the selected subsets are downloaded. Each document is tokenized with the tokenizer's BOS and a closing EOS. Documents are packed back to back into rows of `data.seq_len` tokens and continue into the next row, so no token is truncated or padded. A step trains on exactly `data.batch_size * data.seq_len` tokens. diff --git a/examples/pretrain/llama3/150m.toml b/examples/pretrain/llama3/150m.toml index 2bbe3c3a3b6..57dd5fdab82 100644 --- a/examples/pretrain/llama3/150m.toml +++ b/examples/pretrain/llama3/150m.toml @@ -36,10 +36,10 @@ holdout_docs = 10000 lr = 1e-3 betas2 = 0.95 -[scheduler] +[schedules.optim.lr] type = "cosine" -warmup_steps = 300 -min_lr = 0.0001 +warmup = 300 +end_ratio = 0.1 [val] interval = 285 diff --git a/examples/pretrain/llama3/17m.toml b/examples/pretrain/llama3/17m.toml index c4694e6e265..b7aacc818b7 100644 --- a/examples/pretrain/llama3/17m.toml +++ b/examples/pretrain/llama3/17m.toml @@ -36,10 +36,10 @@ holdout_docs = 10000 lr = 3e-3 betas2 = 0.95 -[scheduler] +[schedules.optim.lr] type = "cosine" -warmup_steps = 50 -min_lr = 0.0003 +warmup = 50 +end_ratio = 0.1 [val] interval = 100 diff --git a/examples/pretrain/llama3/1b.toml b/examples/pretrain/llama3/1b.toml index 9aa25f9e8bb..d0c489948bb 100644 --- a/examples/pretrain/llama3/1b.toml +++ b/examples/pretrain/llama3/1b.toml @@ -35,10 +35,10 @@ holdout_docs = 10000 lr = 4e-4 betas2 = 0.95 -[scheduler] +[schedules.optim.lr] type = "cosine" -warmup_steps = 1850 -min_lr = 4e-05 +warmup = 1850 +end_ratio = 0.1 [val] interval = 1860 diff --git a/examples/pretrain/llama3/500m.toml b/examples/pretrain/llama3/500m.toml index c8169205f18..76943ec5b68 100644 --- a/examples/pretrain/llama3/500m.toml +++ b/examples/pretrain/llama3/500m.toml @@ -35,10 +35,10 @@ holdout_docs = 10000 lr = 6e-4 betas2 = 0.95 -[scheduler] +[schedules.optim.lr] type = "cosine" -warmup_steps = 950 -min_lr = 6e-05 +warmup = 950 +end_ratio = 0.1 [val] interval = 950 diff --git a/examples/pretrain/llama3/70m.toml b/examples/pretrain/llama3/70m.toml index 6352b7a95fe..7b6a592c86d 100644 --- a/examples/pretrain/llama3/70m.toml +++ b/examples/pretrain/llama3/70m.toml @@ -36,10 +36,10 @@ holdout_docs = 10000 lr = 1.5e-3 betas2 = 0.95 -[scheduler] +[schedules.optim.lr] type = "cosine" -warmup_steps = 150 -min_lr = 0.00015 +warmup = 150 +end_ratio = 0.1 [val] interval = 140 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 71a896d315e..337b3166893 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -10,13 +10,10 @@ from prime_rl.configs.trainer import ( AdamWConfig, CheckpointConfig, - ConstantSchedulerConfig, GCConfig, ModelConfig, OptimizerConfig, - SchedulerConfig, TokenizerConfig, - validate_scheduler, ) from prime_rl.utils.config import BaseConfig, default_output_dir @@ -118,6 +115,85 @@ class ValConfig(BaseConfig): """Held-out sets. With ``data.holdout_docs > 0`` the training holdout is added as ``holdout``.""" +StepsOrFraction: TypeAlias = Annotated[int, Field(ge=0)] | Annotated[float, Field(gt=0, lt=1)] + + +def resolve_steps(value: int | float, total_steps: int | None) -> int: + """Steps as given, or a fraction of ``total_steps``.""" + return round(value * total_steps) if isinstance(value, float) else value + + +class WSDScheduleConfig(BaseConfig): + """Warmup, stable at 1, decay. ``f(0) = start_ratio``, the warmup reaches 1 at its last update + (``warmup - 1``; a warmup of 0 or 1 starts at 1), and the decay reaches ``end_ratio`` at the last + update of the run.""" + + type: Literal["wsd"] = "wsd" + + warmup: StepsOrFraction = 0.1 + """Warmup updates, or a fraction of ``max_steps``.""" + + decay: StepsOrFraction = 0.2 + """Decay updates at the end of the run, or a fraction of ``max_steps``.""" + + start_ratio: float = Field(0.0, ge=0) + """Multiplier of the first update.""" + + end_ratio: float = Field(0.0, ge=0) + """Multiplier of the last update.""" + + decay_shape: Literal["linear", "cosine", "1-sqrt"] = "linear" + """Shape of the decay from 1 to ``end_ratio``.""" + + +class CosineScheduleConfig(BaseConfig): + """Linear warmup, then cosine decay from 1 to ``end_ratio`` over the remaining updates.""" + + type: Literal["cosine"] = "cosine" + + warmup: StepsOrFraction = 0.1 + """Warmup updates, or a fraction of ``max_steps``.""" + + start_ratio: float = Field(0.0, ge=0) + """Multiplier of the first update.""" + + end_ratio: float = Field(0.0, ge=0) + """Multiplier of the last update.""" + + +class ConstantScheduleConfig(BaseConfig): + """Multiplier 1 at every step.""" + + type: Literal["constant"] = "constant" + + +class CustomScheduleConfig(BaseConfig): + type: Literal["custom"] = "custom" + + import_path: str + """Import path to the schedule function, called as ``fn(step, total_steps=max_steps, **kwargs)`` and returning the multiplier.""" + + kwargs: dict[str, Any] = Field(default_factory=dict) + """Kwargs forwarded to the schedule function.""" + + +ScheduleConfig: TypeAlias = Annotated[ + WSDScheduleConfig | CosineScheduleConfig | ConstantScheduleConfig | CustomScheduleConfig, + Field(discriminator="type"), +] + + +class SchedulesConfig(BaseConfig): + """Each schedule is a multiplier ``f(step)``: update ``step`` (counted from 0) uses ``base * f(step)``, + where ``base`` is the configured value of each param group.""" + + optim: dict[str, ScheduleConfig] = Field(default_factory=lambda: {"lr": WSDScheduleConfig()}) + """Schedules for optimizer hyperparameters, keyed by their ``optim`` field name (``lr``, ``weight_decay``, ``betas1``, ``betas2``, and ``mu`` for Muon). Unscheduled ones stay constant.""" + + clip_norm: ScheduleConfig | None = None + """Schedule for the gradient clipping norm ``optim.max_norm``. If None, it stays constant.""" + + class PretrainDeploymentConfig(BaseConfig): num_train_gpus: int = Field(1, ge=1) """GPUs for the trainer, on this node.""" @@ -138,7 +214,8 @@ class PretrainConfig(BaseConfig): optim: OptimizerConfig = AdamWConfig() - scheduler: SchedulerConfig = ConstantSchedulerConfig() + schedules: SchedulesConfig = SchedulesConfig() + """Hyperparameter schedules. Defaults to a WSD learning rate schedule.""" ckpt: CheckpointConfig | None = CheckpointConfig() """Checkpointing. Defaults to a checkpoint at the end of training.""" @@ -256,6 +333,29 @@ def validate_holdout(self): return self @model_validator(mode="after") - def validate_scheduler_steps(self): - validate_scheduler(self.scheduler, self.max_steps) + def validate_schedules(self): + unsupported = set(self.schedules.optim) - set(self.optim.schedulable) + if unsupported: + raise ValueError( + f"schedules.optim: {sorted(unsupported)} cannot be scheduled with optim.type = {self.optim.type!r} " + f"(schedulable: {sorted(self.optim.schedulable)})" + ) + if self.schedules.clip_norm is not None and self.optim.max_norm is None: + raise ValueError("schedules.clip_norm needs optim.max_norm") + named = {f"schedules.optim.{name}": schedule for name, schedule in self.schedules.optim.items()} + named["schedules.clip_norm"] = self.schedules.clip_norm + for name, schedule in named.items(): + if isinstance(schedule, WSDScheduleConfig): + phases = [schedule.warmup, schedule.decay] + needs_max_steps = schedule.decay != 0 or isinstance(schedule.warmup, float) + elif isinstance(schedule, CosineScheduleConfig): + phases = [schedule.warmup] + needs_max_steps = True + else: + continue + if self.max_steps is None: + if needs_max_steps: + raise ValueError(f"{name} needs max_steps (it decays or has fractional phases)") + elif sum(resolve_steps(phase, self.max_steps) for phase in phases) > self.max_steps: + raise ValueError(f"{name}: its phases exceed max_steps ({self.max_steps})") return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 63b2364e8db..2d6a8d35159 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -1,7 +1,7 @@ import re import warnings from pathlib import Path -from typing import Annotated, Any, Literal, TypeAlias +from typing import Annotated, Any, ClassVar, Literal, NamedTuple, TypeAlias from pydantic import Field, field_validator, model_validator @@ -418,7 +418,26 @@ def validate_scheduler(scheduler: SchedulerConfig, max_steps: int | None) -> Non raise ValueError(f"warmup_steps ({scheduler.warmup_steps}) must be less than max_steps ({max_steps})") +class ParamGroupKey(NamedTuple): + """Where an optimizer config field lives in the optimizer's param groups.""" + + key: str + """Param group key.""" + + index: int | None = None + """Index into the param group value when it is a tuple (e.g. torch AdamW ``betas``).""" + + algorithm: str | None = None + """Only param groups with this ``algorithm`` (Dion Muon). None applies to every group.""" + + class BaseOptimizerConfig(BaseConfig): + schedulable: ClassVar[dict[str, ParamGroupKey]] = { + "lr": ParamGroupKey("lr"), + "weight_decay": ParamGroupKey("weight_decay"), + } + """Hyperparameters a schedule can scale, by config field name.""" + lr: float = Field(1e-6, ge=0) """Peak learning rate.""" @@ -430,6 +449,12 @@ class BaseOptimizerConfig(BaseConfig): class AdamWConfig(BaseOptimizerConfig): + schedulable: ClassVar[dict[str, ParamGroupKey]] = { + **BaseOptimizerConfig.schedulable, + "betas1": ParamGroupKey("betas", index=0), + "betas2": ParamGroupKey("betas", index=1), + } + type: Literal["adamw"] = "adamw" betas1: float = Field(0.9, ge=0) @@ -440,6 +465,13 @@ class AdamWConfig(BaseOptimizerConfig): class MuonConfig(BaseOptimizerConfig): + schedulable: ClassVar[dict[str, ParamGroupKey]] = { + **BaseOptimizerConfig.schedulable, + "mu": ParamGroupKey("mu", algorithm="muon"), + "betas1": ParamGroupKey("beta1", algorithm="adamw"), + "betas2": ParamGroupKey("beta2", algorithm="adamw"), + } + type: Literal["muon"] = "muon" mu: float = Field(0.95, ge=0) diff --git a/src/prime_rl/trainer/ckpt.py b/src/prime_rl/trainer/ckpt.py index 41099e53150..e1a9707a430 100644 --- a/src/prime_rl/trainer/ckpt.py +++ b/src/prime_rl/trainer/ckpt.py @@ -186,7 +186,7 @@ def save_to_path( path: Path, model: nn.Module, optimizers: list[OptimizerLike], - scheduler: LRScheduler, + scheduler: LRScheduler | None, progress: Progress, dataloader: StatefulDataLoader | None = None, ): @@ -273,7 +273,7 @@ def save( step: int, model: nn.Module, optimizers: list[OptimizerLike], - scheduler: LRScheduler, + scheduler: LRScheduler | None, progress: Progress, dataloader: StatefulDataLoader | None = None, ) -> None: diff --git a/src/prime_rl/trainer/pretrain/schedules.py b/src/prime_rl/trainer/pretrain/schedules.py new file mode 100644 index 00000000000..462b7ac2fb3 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/schedules.py @@ -0,0 +1,134 @@ +import math +from functools import partial +from typing import Callable + +from prime_rl.configs.pretrain import ScheduleConfig, SchedulesConfig, resolve_steps +from prime_rl.configs.trainer import OptimizerConfig, ParamGroupKey +from prime_rl.trainer.optim import OptimizerLike +from prime_rl.utils.utils import import_object + +Schedule = Callable[[int], float] +"""Multiplier for an update, given the number of updates completed before it (0 for the first).""" + +DECAY_SHAPES: dict[str, Callable[[float], float]] = { + "linear": lambda progress: 1 - progress, + "cosine": lambda progress: 0.5 * (1 + math.cos(math.pi * progress)), + "1-sqrt": lambda progress: 1 - math.sqrt(progress), +} + + +def wsd( + step: int, + total_steps: int | None, + warmup: int, + decay: int, + start_ratio: float, + end_ratio: float, + decay_shape: str, +) -> float: + """Linear warmup from ``start_ratio`` at step 0 to 1 at step ``warmup - 1``, stable at 1, then + ``decay`` updates that reach ``end_ratio`` at the last step ``total_steps - 1``.""" + peak = max(warmup - 1, 0) + if step < peak: + return start_ratio + (1 - start_ratio) * step / peak + if decay == 0: + return 1.0 + last_stable = max(total_steps - 1 - decay, peak) + if step <= last_stable: + return 1.0 + progress = (step - last_stable) / (total_steps - 1 - last_stable) + return end_ratio + (1 - end_ratio) * DECAY_SHAPES[decay_shape](progress) + + +def cosine(step: int, total_steps: int, warmup: int, start_ratio: float, end_ratio: float) -> float: + """Linear warmup as in ``wsd``, then cosine decay from 1 to ``end_ratio`` at the last step.""" + decay = total_steps - max(warmup, 1) + return wsd(step, total_steps, warmup, decay, start_ratio, end_ratio, "cosine") + + +def constant(step: int, total_steps: int | None) -> float: + return 1.0 + + +def build_schedule(config: ScheduleConfig, total_steps: int | None) -> Schedule: + """Bind a schedule config to the number of updates it spans. Fractional phases resolve against it.""" + match config.type: + case "wsd": + return partial( + wsd, + total_steps=total_steps, + warmup=resolve_steps(config.warmup, total_steps), + decay=resolve_steps(config.decay, total_steps), + start_ratio=config.start_ratio, + end_ratio=config.end_ratio, + decay_shape=config.decay_shape, + ) + case "cosine": + return partial( + cosine, + total_steps=total_steps, + warmup=resolve_steps(config.warmup, total_steps), + start_ratio=config.start_ratio, + end_ratio=config.end_ratio, + ) + case "constant": + return partial(constant, total_steps=total_steps) + case "custom": + return partial(import_object(config.import_path), total_steps=total_steps, **config.kwargs) + + +class Schedules: + """Sets the scheduled hyperparameters of each update. The values are a pure function of the step, + so nothing needs to be checkpointed: on resume they are recomputed from the step. + + The base values are read from the optimizer's param groups at construction, before any checkpoint + load. Each param group keeps its own base, so ratios between groups hold at every step. + """ + + def __init__( + self, + config: SchedulesConfig, + optim_config: OptimizerConfig, + optimizer: OptimizerLike, + total_steps: int | None, + step_offset: int = 0, + ): + self.optimizer = optimizer + self.step_offset = step_offset + # Group indices, not group dicts: loading an optimizer state dict replaces the group dicts. + self.optim: dict[str, tuple[Schedule, ParamGroupKey, list[tuple[int, float]]]] = {} + for name, schedule in config.optim.items(): + target = optim_config.schedulable[name] + bases = [ + (i, group[target.key] if target.index is None else group[target.key][target.index]) + for i, group in enumerate(optimizer.param_groups) + if target.algorithm is None or group["algorithm"] == target.algorithm + ] + self.optim[name] = (build_schedule(schedule, total_steps), target, bases) + self.max_norm = optim_config.max_norm + self.clip_norm = build_schedule(config.clip_norm, total_steps) if config.clip_norm is not None else None + + def apply(self, step: int) -> tuple[float | None, dict[str, float]]: + """Write the values of update ``step`` (updates completed so far) into the optimizer. + Returns the gradient clipping norm and metrics of every scheduled value.""" + step -= self.step_offset + metrics = {} + for name, (schedule, target, bases) in self.optim.items(): + multiplier = schedule(step) + for i, base in bases: + group = self.optimizer.param_groups[i] + if target.index is None: + group[target.key] = base * multiplier + else: + values = list(group[target.key]) + values[target.index] = base * multiplier + group[target.key] = tuple(values) + metrics[f"schedule/{name}"] = multiplier + metrics[f"optim/{name}"] = bases[0][1] * multiplier + max_norm = self.max_norm + if self.clip_norm is not None and max_norm is not None: + multiplier = self.clip_norm(step) + max_norm *= multiplier + metrics["schedule/clip_norm"] = multiplier + metrics["optim/max_norm"] = max_norm + return max_norm, metrics diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index a29bf227c4b..89bc8e2d297 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -35,7 +35,7 @@ from prime_rl.trainer.pretrain.data import load_text_dataset, setup_text_dataloader, split_holdout from prime_rl.trainer.pretrain.evals.common import token_byte_lengths from prime_rl.trainer.pretrain.evals.perplexity import build_perplexity_set, evaluate_perplexity, load_perplexity_set -from prime_rl.trainer.scheduler import setup_scheduler +from prime_rl.trainer.pretrain.schedules import Schedules from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state, setup_dataloader, setup_dataset from prime_rl.trainer.utils import ( GarbageCollection, @@ -152,12 +152,16 @@ def train(config: PretrainConfig): ), ) - scheduler_steps = ( - config.max_steps - checkpoint_step - if config.max_steps is not None and (config.ckpt and config.ckpt.skip_scheduler and checkpoint_step is not None) - else config.max_steps + # With ckpt.skip_scheduler, the schedules restart at the resumed step and span the remaining steps. + restart_schedules = config.ckpt is not None and config.ckpt.skip_scheduler and checkpoint_step is not None + schedule_offset = checkpoint_step if restart_schedules else 0 + schedules = Schedules( + config.schedules, + config.optim, + optimizer, + total_steps=config.max_steps - schedule_offset if config.max_steps is not None else None, + step_offset=schedule_offset, ) - scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) holdout = None if config.data.type == "text": @@ -189,14 +193,12 @@ def train(config: PretrainConfig): checkpoint_step, model, [optimizer], - scheduler if not skip.skip_scheduler else None, + None, progress, dataloader=dataloader if not skip.skip_dataloader else None, path=config.resume.dir / "trainer" if config.resume.dir is not None else None, ) progress.step += 1 - if skip.skip_scheduler: - scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) logger.info(f"Resuming from step {checkpoint_step} (dataset_state={get_dataset_state(dataloader)})") # Create the iterator only after a potential resume: iter() forks the workers with the @@ -254,6 +256,7 @@ def run_evals(step: int) -> None: gc_handler.run(progress.step) is_last_step = config.max_steps is not None and progress.step >= config.max_steps step_start_time = time.perf_counter() + max_norm, schedule_metrics = schedules.apply(progress.step - 1) step_loss_sum = torch.tensor(0.0, device="cuda") nan_loss_count = torch.tensor(0, device="cuda") @@ -319,17 +322,16 @@ def run_evals(step: int) -> None: batch_loss = (step_loss_sum / global_token_count_val).item() nan_loss_count = nan_loss_count.item() - grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm) if config.optim.max_norm else None + grad_norm = clip_grad_norm_(gradient_manager, model, max_norm) if max_norm else None optimizer.step() optimizer.zero_grad() current_lr = optimizer.param_groups[0]["lr"] - scheduler.step() # The last step's checkpoint is written once after the loop. save_ckpt_time = 0.0 if config.ckpt and config.ckpt.interval and progress.step % config.ckpt.interval == 0 and not is_last_step: save_ckpt_start_time = time.perf_counter() - ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) + ckpt_manager.save(progress.step, model, [optimizer], None, progress, dataloader=dataloader) ckpt_manager.maybe_clean() save_ckpt_time = time.perf_counter() - save_ckpt_start_time @@ -379,6 +381,7 @@ def run_evals(step: int) -> None: **get_ckpt_disk_metrics(config.run_dir), **{name: value.item() for name, value in moe_stats.items()}, **expert_load_stats, + **schedule_metrics, "step": progress.step, } if grad_norm is not None: @@ -393,7 +396,7 @@ def run_evals(step: int) -> None: if config.ckpt is not None: logger.info(f"Saving final checkpoint at step {progress.step}") - ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) + ckpt_manager.save(progress.step, model, [optimizer], None, progress, dataloader=dataloader) ckpt_manager.maybe_clean() if gradient_manager is not None: gradient_manager.close() diff --git a/tests/unit/train/pretrain/test_schedules.py b/tests/unit/train/pretrain/test_schedules.py new file mode 100644 index 00000000000..d8483e61e79 --- /dev/null +++ b/tests/unit/train/pretrain/test_schedules.py @@ -0,0 +1,87 @@ +import math + +import pytest +import torch +from pydantic import ValidationError + +from prime_rl.configs.pretrain import ( + CosineScheduleConfig, + CustomScheduleConfig, + PretrainConfig, + SchedulesConfig, + WSDScheduleConfig, +) +from prime_rl.configs.trainer import AdamWConfig +from prime_rl.trainer.pretrain.schedules import Schedules, build_schedule + + +def test_wsd_endpoints_and_shapes(): + config = WSDScheduleConfig(warmup=0.1, decay=0.2, start_ratio=0.1, end_ratio=0.05) + f = build_schedule(config, total_steps=100) + assert f(0) == pytest.approx(0.1) + assert f(9) == 1.0 + assert f(79) == 1.0 + assert f(89) == pytest.approx(0.05 + 0.95 * 0.5) + assert f(99) == pytest.approx(0.05) + + for shape, midpoint in [("cosine", 0.5), ("1-sqrt", 1 - math.sqrt(0.5))]: + f = build_schedule(WSDScheduleConfig(warmup=0, decay=20, decay_shape=shape), total_steps=100) + assert f(0) == 1.0 + assert f(89) == pytest.approx(midpoint) + assert f(99) == 0.0 + + +def test_cosine_endpoints(): + f = build_schedule(CosineScheduleConfig(warmup=10, end_ratio=0.1), total_steps=100) + assert f(0) == 0.0 + assert f(9) == 1.0 + assert f(54) == pytest.approx(0.55) + assert f(99) == pytest.approx(0.1) + assert build_schedule(CosineScheduleConfig(warmup=0), total_steps=10)(0) == 1.0 + + +def inverse_sqrt(step: int, total_steps: int, scale: float) -> float: + return scale / math.sqrt(step + 1) + + +def test_custom_schedule(): + config = CustomScheduleConfig(import_path=f"{__name__}.inverse_sqrt", kwargs={"scale": 2.0}) + assert build_schedule(config, total_steps=10)(3) == 1.0 + + +def test_schedules_config_validation(): + PretrainConfig(max_steps=100, optim={"type": "muon"}, schedules={"optim": {"mu": {"type": "constant"}}}) + for schedules in [{"optim": {"momentum": {"type": "constant"}}}, {"optim": {"mu": {"type": "constant"}}}]: + with pytest.raises(ValidationError, match="cannot be scheduled"): + PretrainConfig(max_steps=100, schedules=schedules) + with pytest.raises(ValidationError, match="needs max_steps"): + PretrainConfig() + with pytest.raises(ValidationError, match="exceed max_steps"): + PretrainConfig(max_steps=100, schedules={"optim": {"lr": {"type": "wsd", "warmup": 50, "decay": 51}}}) + + +def test_apply_keeps_group_ratios(): + params = [torch.nn.Parameter(torch.zeros(2)), torch.nn.Parameter(torch.zeros(2))] + optimizer = torch.optim.AdamW( + [{"params": params[:1], "lr": 1.0}, {"params": params[1:], "lr": 0.1}], betas=(0.9, 0.95) + ) + config = SchedulesConfig( + optim={ + "lr": WSDScheduleConfig(warmup=0, decay=2, end_ratio=0.5), + "betas2": WSDScheduleConfig(warmup=0, decay=1, end_ratio=0.5), + }, + clip_norm=WSDScheduleConfig(warmup=2, decay=0), + ) + schedules = Schedules(config, AdamWConfig(max_norm=2.0), optimizer, total_steps=4) + max_norm, metrics = schedules.apply(0) + assert max_norm == 0.0 + assert [group["lr"] for group in optimizer.param_groups] == [1.0, 0.1] + + # Loading a state dict replaces the param group dicts. + optimizer.load_state_dict(optimizer.state_dict()) + max_norm, metrics = schedules.apply(3) + assert max_norm == 2.0 + assert [group["lr"] for group in optimizer.param_groups] == pytest.approx([0.5, 0.05]) + assert [group["betas"] for group in optimizer.param_groups] == [(0.9, 0.475)] * 2 + assert metrics["schedule/lr"] == 0.5 + assert metrics["optim/lr"] == 0.5 From 00354bfd37cde76860f9e79ef3f4535fe9b569a6 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 10:47:36 +0200 Subject: [PATCH 007/132] feat(pretrain): aux-loss-free MoE load balancing for glm4_moe Update each MoE layer's selection bias once per step with the DeepSeek sign rule (gamma = load_balance_coeff, now a GLM-4.5 config field) from the step's globally reduced tokens per expert; zero the bias on random init; count shared experts in MFU. --- src/prime_rl/trainer/model.py | 9 +++++++++ .../models/glm4_moe/configuration_glm4_moe.py | 4 ++++ .../trainer/models/glm4_moe/modeling_glm4_moe.py | 7 ++++--- src/prime_rl/trainer/perf.py | 5 +++-- src/prime_rl/trainer/pretrain/train.py | 10 +++++++++- tests/unit/train/models/test_moe.py | 16 +++++++++++++++- 6 files changed, 44 insertions(+), 7 deletions(-) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 821f9373adf..136452c558a 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -330,6 +330,15 @@ def get_expert_load_stats(tokens_per_expert: Tensor, group: dist.ProcessGroup) - return {name: value.item() for name, value in compute_expert_load_stats(tokens_per_expert).items()} +@torch.no_grad() +def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor) -> None: + """Aux-loss-free load balancing (https://arxiv.org/abs/2408.15664): move each layer's selection bias by + `load_balance_coeff` towards underloaded experts, given the step's global `[num_moe_layers, num_experts]` counts.""" + for moe, load in zip(iter_moe_blocks(model), tokens_per_expert): + if moe.load_balance_coeff is not None: + moe.router.selection_bias += moe.load_balance_coeff * torch.sign(load.mean() - load) + + def get_model( config: ModelConfig, device: torch.device = torch.device("cpu"), dtype: torch.dtype = torch.bfloat16 ) -> nn.Module: diff --git a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py index 2763679fa93..00d7f286805 100644 --- a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py @@ -102,6 +102,8 @@ class Glm4MoeConfig(PretrainedConfig): Whether to normalize the topk probabilities. use_qk_norm (`bool`, *optional*, defaults to `False`): Whether to use query-key normalization in the attention + load_balance_coeff (`float`, *optional*, defaults to 1e-3): + Update speed of the aux-loss-free load-balancing bias; `None` disables the bias. ```python >>> from transformers import Glm4MoeModel, Glm4MoeConfig >>> # Initializing a Glm4Moe style configuration @@ -162,6 +164,7 @@ def __init__( first_k_dense_replace=1, norm_topk_prob=True, use_qk_norm=False, + load_balance_coeff=1e-3, pad_token_id=None, **kwargs, ): @@ -193,6 +196,7 @@ def __init__( self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob self.use_qk_norm = use_qk_norm + self.load_balance_coeff = load_balance_coeff super().__init__( tie_word_embeddings=tie_word_embeddings, diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py index 19cda1c6259..5122d1a34c7 100644 --- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py @@ -61,7 +61,7 @@ def __init__(self, config: Glm4MoeConfig, layer_idx: int): route_scale=config.routed_scaling_factor, score_before_experts=False, top_k=config.num_experts_per_tok, - load_balance_coeff=1e-3, + load_balance_coeff=config.load_balance_coeff, ) if layer_idx >= config.first_k_dense_replace: shared_expert = None @@ -339,8 +339,9 @@ def init_buffers_post_meta(self): ) rotary_emb.inv_freq.copy_(inv_freq) - # TODO: Init TT MoE buffers - # I think .to_empty() on gpu by default fills 0 so we are ok but this might not be guaranteed behavior + for module in self.modules(): + if isinstance(module, MoE) and module.router.selection_bias is not None: + module.router.selection_bias.zero_() __all__ = ["Glm4MoeConfig", "Glm4MoePreTrainedModel", "Glm4MoeModel", "Glm4MoeForCausalLM"] diff --git a/src/prime_rl/trainer/perf.py b/src/prime_rl/trainer/perf.py index 13ba65d6b99..17b25ff3c0f 100644 --- a/src/prime_rl/trainer/perf.py +++ b/src/prime_rl/trainer/perf.py @@ -156,8 +156,9 @@ def get_active_mm_params(config: PretrainedConfig) -> float: # Some MoE models (e.g. DeepSeek) use moe_intermediate_size, others (e.g. Granite) just use intermediate_size moe_intermediate_size = getattr(config, "moe_intermediate_size", None) or intermediate_size - if hasattr(config, "num_shared_experts") and config.num_shared_experts: # Shared experts - sparse_mlp_params += num_sparse_layers * config.num_shared_experts * 3 * moe_intermediate_size * hidden_size + num_shared_experts = getattr(config, "num_shared_experts", None) or getattr(config, "n_shared_experts", 0) + if num_shared_experts: # Shared experts + sparse_mlp_params += num_sparse_layers * num_shared_experts * 3 * moe_intermediate_size * hidden_size if hasattr(config, "num_experts_per_tok") and config.num_experts_per_tok: # Routed experts sparse_mlp_params += ( num_sparse_layers * config.num_experts_per_tok * 3 * moe_intermediate_size * hidden_size diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index ca86c26da53..9b82094e5ab 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -24,9 +24,11 @@ get_global_moe_stats, get_load_balance_stats, is_tt_moe_model, + iter_moe_blocks, reshard_module, setup_model, setup_tokenizer, + update_selection_biases, ) from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX from prime_rl.trainer.optim import setup_optimizer @@ -311,7 +313,13 @@ def run_evals(step: int) -> None: step_tokens_per_expert += tokens_per_expert forward_backward_time = time.perf_counter() - step_start_time - expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) if is_moe_model else {} + expert_load_stats = {} + if is_moe_model: + expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) + update_selection_biases(model, step_tokens_per_expert) + biases = [moe.router.selection_bias for moe in iter_moe_blocks(model) if moe.load_balance_coeff is not None] + if biases: + expert_load_stats["expert_load/selection_bias_absmax"] = torch.stack(biases).abs().max().item() if gradient_manager is None: global_step_token_count = step_local_token_count.clone() diff --git a/tests/unit/train/models/test_moe.py b/tests/unit/train/models/test_moe.py index 59824b01412..4c1f723a262 100644 --- a/tests/unit/train/models/test_moe.py +++ b/tests/unit/train/models/test_moe.py @@ -4,7 +4,7 @@ from prime_rl.configs.trainer import ModelConfig from prime_rl.trainer.distributed.token_dispatcher import LocalTokenDispatcher -from prime_rl.trainer.model import compute_expert_load_stats, is_tt_moe_model +from prime_rl.trainer.model import compute_expert_load_stats, is_tt_moe_model, update_selection_biases from prime_rl.trainer.models.fusions import fuse_gate_up_projections from prime_rl.trainer.models.layers.activations import ActivationDispatch, ClampedSilu from prime_rl.trainer.models.layers.expert_compute import BF16ExpertCompute, GroupedGemmExpertCompute @@ -251,3 +251,17 @@ def test_expert_load_stats(): assert stats.keys() == expected.keys() for name, value in expected.items(): assert stats[name].item() == pytest.approx(value) + + +def test_update_selection_biases_moves_towards_underloaded_experts(): + moe = MoE.from_args(MoEArgs(num_experts=4, load_balance_coeff=0.1), dim=4, hidden_dim=8, shared_expert=None) + layer = torch.nn.Module() + layer.mlp = moe + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.layers = torch.nn.ModuleList([layer]) + + update_selection_biases(model, torch.tensor([[8.0, 4.0, 4.0, 0.0]])) + update_selection_biases(model, torch.tensor([[8.0, 4.0, 4.0, 0.0]])) + + torch.testing.assert_close(moe.router.selection_bias, torch.tensor([-0.2, 0.0, 0.0, 0.2])) From 7ca7bfd9403c796b389c72200e37957035dafbf0 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 11:27:21 +0200 Subject: [PATCH 008/132] feat(pretrain): DeepSeek-V3 MoE load balancing: centred bias update, final-stage freeze, sequence-wise balance loss --- .../src/prime_rl/configs/pretrain.py | 4 ++ src/prime_rl/trainer/model.py | 7 +++- .../models/glm4_moe/configuration_glm4_moe.py | 4 ++ .../models/glm4_moe/modeling_glm4_moe.py | 1 + src/prime_rl/trainer/models/layers/moe.py | 38 +++++++++++++++++-- src/prime_rl/trainer/pretrain/train.py | 31 ++++++++++++--- tests/unit/train/models/test_moe.py | 16 ++++++-- 7 files changed, 87 insertions(+), 14 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 35f46164271..8b9546a53d6 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -166,6 +166,10 @@ class PretrainConfig(BaseConfig): max_steps: int | None = None """Maximum training steps. If None, runs indefinitely.""" + moe_bias_freeze: float = Field(0.1, ge=0, lt=1) + """Fraction of ``max_steps`` at the end of training during which the MoE selection biases stop updating, as + DeepSeek-V3, GLM-4.5 and Moonlight do in their final stage. Without ``max_steps`` the biases never freeze.""" + gc: GCConfig | None = GCConfig() """Deterministic garbage collection every N steps. Set to null for Python's default GC.""" diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 136452c558a..1e1f193b725 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -333,10 +333,13 @@ def get_expert_load_stats(tokens_per_expert: Tensor, group: dist.ProcessGroup) - @torch.no_grad() def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor) -> None: """Aux-loss-free load balancing (https://arxiv.org/abs/2408.15664): move each layer's selection bias by - `load_balance_coeff` towards underloaded experts, given the step's global `[num_moe_layers, num_experts]` counts.""" + `load_balance_coeff` towards underloaded experts, given the step's global `[num_moe_layers, num_experts]` counts. + The update is mean-centred (as in Moonlight and torchtitan), which leaves the selection unchanged but keeps + the biases from drifting together.""" for moe, load in zip(iter_moe_blocks(model), tokens_per_expert): if moe.load_balance_coeff is not None: - moe.router.selection_bias += moe.load_balance_coeff * torch.sign(load.mean() - load) + delta = torch.sign(load.mean() - load) + moe.router.selection_bias += moe.load_balance_coeff * (delta - delta.mean()) def get_model( diff --git a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py index 00d7f286805..2b0f18d0002 100644 --- a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py @@ -104,6 +104,8 @@ class Glm4MoeConfig(PretrainedConfig): Whether to use query-key normalization in the attention load_balance_coeff (`float`, *optional*, defaults to 1e-3): Update speed of the aux-loss-free load-balancing bias; `None` disables the bias. + seq_aux_loss_coeff (`float`, *optional*): + Weight of the sequence-wise balance loss (DeepSeek-V3); `None` disables it, 0 only logs it. ```python >>> from transformers import Glm4MoeModel, Glm4MoeConfig >>> # Initializing a Glm4Moe style configuration @@ -165,6 +167,7 @@ def __init__( norm_topk_prob=True, use_qk_norm=False, load_balance_coeff=1e-3, + seq_aux_loss_coeff=None, pad_token_id=None, **kwargs, ): @@ -197,6 +200,7 @@ def __init__( self.norm_topk_prob = norm_topk_prob self.use_qk_norm = use_qk_norm self.load_balance_coeff = load_balance_coeff + self.seq_aux_loss_coeff = seq_aux_loss_coeff super().__init__( tie_word_embeddings=tie_word_embeddings, diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py index 5122d1a34c7..f50f329c1ef 100644 --- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py @@ -62,6 +62,7 @@ def __init__(self, config: Glm4MoeConfig, layer_idx: int): score_before_experts=False, top_k=config.num_experts_per_tok, load_balance_coeff=config.load_balance_coeff, + seq_aux_loss_coeff=config.seq_aux_loss_coeff, ) if layer_idx >= config.first_k_dense_replace: shared_expert = None diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index f0494e435e4..b580d118e8d 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -61,6 +61,7 @@ class MoEArgs: # token-choice top_k: int = 1 load_balance_coeff: float | None = 1e-3 + seq_aux_loss_coeff: float | None = None def __post_init__(self) -> None: ActivationDispatch[self.activation] @@ -180,14 +181,14 @@ def forward( self, x: torch.Tensor, routed_experts: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Args: x (torch.Tensor): Input tensor with shape ``(bs*slen, dim)``. routed_experts (torch.Tensor | None, optional): Optional tensor with shape ``(bs * slen, top_k)``. Returns: - tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - top_scores (torch.Tensor): Routing scores for selected experts with shape ``(bs*slen, top_k)``. - selected_experts_indices (torch.Tensor): @@ -196,6 +197,8 @@ def forward( Number of tokens assigned to each expert with shape ``(num_experts,)``. - routing_confidence_sum (torch.Tensor): Sum over tokens of the selected-expert probability mass before route normalization/scaling. + - scores (torch.Tensor): + Scores of all experts with shape ``(bs*slen, num_experts)``. """ # scores shape (bs*slen, num_experts) assert routed_experts is None or routed_experts.shape[-1] == self.top_k, ( @@ -266,7 +269,7 @@ def forward( max=self.num_experts, ).to(torch.int64) - return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum + return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, scores def init_weights(self, init_std: float): nn.init.trunc_normal_(self.gate.weight, mean=0.0, std=init_std) @@ -274,6 +277,25 @@ def init_weights(self, init_std: float): nn.init.zeros_(self.gate.bias) +def sequence_balance(scores: torch.Tensor, selected_experts: torch.Tensor, num_rows: int) -> torch.Tensor: + """DeepSeek-V3's sequence-wise balance loss without its coefficient, per row: ``sum_i f_i P_i`` with + ``f_i`` the row's share of expert ``i``'s selections times ``num_experts / top_k`` (1 when balanced) and + ``P_i`` its mean score normalized over all experts. 1 for a balanced row, up to ``num_experts / top_k``. + + Args: + scores: ``(num_rows * slen, num_experts)`` scores of all experts. + selected_experts: ``(num_rows * slen, top_k)`` selected experts. + """ + num_experts = scores.shape[1] + selections = selected_experts.view(num_rows, -1) + counts = torch.zeros(num_rows, num_experts, device=scores.device).scatter_add_( + 1, selections, torch.ones_like(selections, dtype=torch.float32) + ) + f = counts * (num_experts / selections.shape[1]) + p = (scores / scores.sum(dim=-1, keepdim=True)).view(num_rows, -1, num_experts).mean(dim=1) + return (f * p).sum(dim=-1) + + class MoE(nn.Module): """Token-choice MoE runtime composed from a router, grouped experts, and optional projections.""" @@ -308,6 +330,7 @@ def from_args( shared_expert=shared_expert, score_before_experts=args.score_before_experts, load_balance_coeff=args.load_balance_coeff, + seq_aux_loss_coeff=args.seq_aux_loss_coeff, ) def __init__( @@ -318,6 +341,7 @@ def __init__( shared_expert: FeedForward | None, score_before_experts: bool, load_balance_coeff: float | None, + seq_aux_loss_coeff: float | None = None, ) -> None: super().__init__() self.router = router @@ -344,6 +368,10 @@ def __init__( persistent=False, ) self.register_buffer("routing_confidence_sum", torch.tensor(0.0, dtype=torch.float32), persistent=False) + # Sequence-wise balance loss (DeepSeek-V3 eq. 17-20): with a coefficient set, every forward with grad + # stores the per-row loss in `seq_balance`, for the training loop to add with that coefficient. + self.seq_aux_loss_coeff = seq_aux_loss_coeff + self.seq_balance: torch.Tensor | None = None def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None: self.token_dispatcher = token_dispatcher @@ -383,8 +411,12 @@ def forward( selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, + scores, ) = self.router(x, routed_experts=routed_experts) + if self.seq_aux_loss_coeff is not None and torch.is_grad_enabled(): + self.seq_balance = sequence_balance(scores, selected_experts_indices, bs) + # Accumulate expert usage for selection-bias updates and metrics. with torch.no_grad(): record_moe_routing_statistics( diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 9b82094e5ab..a229c3b866a 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -218,9 +218,13 @@ def train(config: PretrainConfig): dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None is_moe_model = is_tt_moe_model(model) + moes_with_seq_balance = [moe for moe in iter_moe_blocks(model) if moe.seq_aux_loss_coeff is not None] + bias_freeze_step = config.max_steps - round(config.moe_bias_freeze * config.max_steps) if config.max_steps else None - def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: - """Forward pass returning (loss_sum, token_count) over unmasked tokens.""" + def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward pass returning (loss_sum, token_count) over unmasked tokens, and the MoE sequence-wise + balance loss summed over tokens (each token carries its row's loss, so after the division by the + token count it is the mean over rows).""" input_ids = micro_batch["input_ids"].to("cuda", non_blocking=True) position_ids = micro_batch["position_ids"].to("cuda", non_blocking=True) target_ids = micro_batch["target_ids"].to("cuda", non_blocking=True) @@ -237,7 +241,10 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: out = forward( model, input_ids, position_ids, seq_lens=seq_lens, labels=labels, seq_lens_are_pre_shard=cp_enabled ) - return out["loss"], loss_mask.sum(dtype=torch.int64) + aux_loss_sum = torch.zeros((), device="cuda") + for moe in moes_with_seq_balance: + aux_loss_sum += moe.seq_aux_loss_coeff * moe.seq_balance.sum() * input_ids.shape[1] + return out["loss"], loss_mask.sum(dtype=torch.int64), aux_loss_sum def run_evals(step: int) -> None: """Perplexity evals on the weights after ``step`` optimizer updates.""" @@ -287,10 +294,12 @@ def run_evals(step: int) -> None: prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=True) step_tokens_per_expert = 0 + micro_max_mean = 0.0 + seq_balance = 0.0 step_local_num_tokens = 0 for micro_step, micro_batch in enumerate(micro_batches): step_local_num_tokens += micro_batch["input_ids"].shape[1] - local_loss_sum, batch_token_count = compute_loss(micro_batch) + local_loss_sum, batch_token_count, aux_loss_sum = compute_loss(micro_batch) if gradient_manager is None: step_local_token_count += batch_token_count if torch.isnan(local_loss_sum.detach()): @@ -299,7 +308,7 @@ def run_evals(step: int) -> None: scaled_loss = torch.nan_to_num(local_loss_sum, nan=0.0) / grad_accum_steps else: step_loss_sum += local_loss_sum.detach() - scaled_loss = local_loss_sum / grad_accum_steps + scaled_loss = (local_loss_sum + aux_loss_sum) / grad_accum_steps begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) scaled_loss.backward() finish_backward(gradient_manager) @@ -311,12 +320,22 @@ def run_evals(step: int) -> None: if name == "max_vio": moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) step_tokens_per_expert += tokens_per_expert + micro_max_mean += ( + tokens_per_expert.amax(dim=1) / tokens_per_expert.mean(dim=1) + ).mean() / grad_accum_steps + if moes_with_seq_balance: + seq_balance += torch.cat([moe.seq_balance.detach() for moe in moes_with_seq_balance]).mean() forward_backward_time = time.perf_counter() - step_start_time expert_load_stats = {} if is_moe_model: expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) - update_selection_biases(model, step_tokens_per_expert) + # This rank's micro-batches: the imbalance that expert parallelism and the dropless dispatch see. + expert_load_stats["expert_load/micro_max_mean"] = micro_max_mean.item() + if moes_with_seq_balance: + expert_load_stats["expert_load/seq_balance"] = seq_balance.item() / grad_accum_steps + if bias_freeze_step is None or progress.step <= bias_freeze_step: + update_selection_biases(model, step_tokens_per_expert) biases = [moe.router.selection_bias for moe in iter_moe_blocks(model) if moe.load_balance_coeff is not None] if biases: expert_load_stats["expert_load/selection_bias_absmax"] = torch.stack(biases).abs().max().item() diff --git a/tests/unit/train/models/test_moe.py b/tests/unit/train/models/test_moe.py index 4c1f723a262..5f622310a75 100644 --- a/tests/unit/train/models/test_moe.py +++ b/tests/unit/train/models/test_moe.py @@ -13,6 +13,7 @@ GroupedExperts, MoE, MoEArgs, + sequence_balance, ) from prime_rl.trainer.models.qwen3_5 import ( Qwen3_5Config, @@ -261,7 +262,16 @@ def test_update_selection_biases_moves_towards_underloaded_experts(): model.model = torch.nn.Module() model.model.layers = torch.nn.ModuleList([layer]) - update_selection_biases(model, torch.tensor([[8.0, 4.0, 4.0, 0.0]])) - update_selection_biases(model, torch.tensor([[8.0, 4.0, 4.0, 0.0]])) + update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]])) + update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]])) - torch.testing.assert_close(moe.router.selection_bias, torch.tensor([-0.2, 0.0, 0.0, 0.2])) + # sign(mean - load) = [-1, 1, 1, 1], centred to [-1.5, 0.5, 0.5, 0.5] per step + torch.testing.assert_close(moe.router.selection_bias, torch.tensor([-0.3, 0.1, 0.1, 0.1])) + + +def test_sequence_balance(): + scores = torch.tensor([[0.9, 0.1], [0.9, 0.1], [0.9, 0.1], [0.1, 0.9]]) + selected = torch.tensor([[0], [0], [0], [1]]) + + # row 0 sends both tokens to expert 0: f = [2, 0], P = [0.9, 0.1]; row 1 is balanced: f = [1, 1], P = [0.5, 0.5] + torch.testing.assert_close(sequence_balance(scores, selected, num_rows=2), torch.tensor([1.8, 1.0])) From 80093c988972988b14e18288793d9fe743157c09 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 14:22:18 +0200 Subject: [PATCH 009/132] feat(pretrain): wsd defaults to linear decay to 0 after the warmup (decay 0.9) --- docs/pretraining.md | 2 +- packages/prime-rl-configs/src/prime_rl/configs/pretrain.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 54b56f5cbc4..aaf856e8648 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -59,7 +59,7 @@ end_ratio = 1.05 | `type` | Shape | |---|---| -| `wsd` (default) | Linear warmup, stable at 1, then a `linear`, `cosine` or `1-sqrt` decay over the last `decay` updates. Defaults: `warmup = 0.1`, `decay = 0.2`. | +| `wsd` (default) | Linear warmup, stable at 1, then a `linear`, `cosine` or `1-sqrt` decay over the last `decay` updates. Defaults: `warmup = 0.1`, `decay = 0.9`, i.e. a linear decay to 0 over all updates after the warmup. | | `cosine` | Linear warmup, then cosine decay over all remaining updates. | | `constant` | 1 at every step. | | `custom` | `fn(step, total_steps=max_steps, **kwargs) -> float` imported from `import_path`, with `kwargs` from the config. | diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 337b3166893..0f15c151743 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -133,8 +133,9 @@ class WSDScheduleConfig(BaseConfig): warmup: StepsOrFraction = 0.1 """Warmup updates, or a fraction of ``max_steps``.""" - decay: StepsOrFraction = 0.2 - """Decay updates at the end of the run, or a fraction of ``max_steps``.""" + decay: StepsOrFraction = 0.9 + """Decay updates at the end of the run, or a fraction of ``max_steps``. The default decays over all + updates after the default warmup.""" start_ratio: float = Field(0.0, ge=0) """Multiplier of the first update.""" From 42eb2782cf17ee037594cbdfb33ff4ffe9bd1eba Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 15:18:21 +0200 Subject: [PATCH 010/132] feat(pretrain): multi-node SLURM launch through the SFT deployment, slurm config and templates --- docs/pretraining.md | 21 ++- .../src/prime_rl/configs/pretrain.py | 26 ++- .../src/prime_rl/configs/sft.py | 57 +++---- src/prime_rl/entrypoints/pretrain.py | 29 +++- src/prime_rl/entrypoints/sft.py | 152 +++++++++--------- .../templates/multi_node_sft.sbatch.j2 | 4 +- .../templates/single_node_sft.sbatch.j2 | 2 +- src/prime_rl/utils/pathing.py | 2 +- tests/unit/train/pretrain/test_launcher.py | 29 ++++ 9 files changed, 205 insertions(+), 117 deletions(-) create mode 100644 tests/unit/train/pretrain/test_launcher.py diff --git a/docs/pretraining.md b/docs/pretraining.md index 8eed93f536b..360bf931b64 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -1,6 +1,6 @@ # Pretraining -`uv run pretrain` trains a model from scratch on raw text. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, scheduler, checkpointing and monitors as the SFT trainer. It runs on a single node. +`uv run pretrain` trains a model from scratch on raw text. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, scheduler, checkpointing and monitors as the SFT trainer. It runs locally on one node or, with `[slurm]`, as a SLURM job on one or more nodes (see [Multi-node](#multi-node)). ```bash uv run pretrain @ examples/pretrain/llama3/17m.toml @@ -69,3 +69,22 @@ subset = "wikitext-103-raw-v1" split = "test" ``` +## Multi-node + +`uv run pretrain` takes the same `[deployment]` and `[slurm]` blocks as `uv run sft` (see [Scaling § SLURM](scaling.md#slurm)) and submits the same bundled templates. The launcher writes the model dir and downloads the data on the submitting node, then submits one job in which every node runs `torchrun` on the pretrain trainer: + +```toml +output_dir = "/shared/outputs/pretrain" + +[deployment] +type = "multi_node" +num_train_nodes = 4 +gpus_per_node = 8 +nodes_per_fsdp_group = 1 # optional: HSDP, shard within a node and replicate across nodes + +[slurm] +partition = "all" +time = "04:00:00" +``` + +Each document goes to exactly one data stream, and there are `data_parallel_ranks * data.num_workers` streams, each packing its own documents into rows. The loss and gradient are normalized by the token count of the whole global batch, as on one GPU with gradient accumulation. The rows of a step depend on the number of streams: runs with different numbers of GPUs see the same data distribution but not the same batches, unless `data_parallel_ranks * data.num_workers` is the same (e.g. 1 GPU with 32 workers and 32 GPUs with 1 worker train on identical batches). diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 71a896d315e..2c2dee8f2b0 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -5,8 +5,14 @@ from pydantic import Field, model_validator from prime_rl.configs.monitors import TrainMonitorsConfig -from prime_rl.configs.sft import BaseDataConfig, FakeDataConfig -from prime_rl.configs.shared import EnvVars, HeartbeatConfig, ResumeConfig, RunConfig, TrainerLogConfig +from prime_rl.configs.sft import ( + BaseDataConfig, + FakeDataConfig, + SFTDeploymentConfig, + SingleNodeDeploymentConfig, + setup_trainer_deployment, +) +from prime_rl.configs.shared import EnvVars, HeartbeatConfig, ResumeConfig, RunConfig, SlurmConfig, TrainerLogConfig from prime_rl.configs.trainer import ( AdamWConfig, CheckpointConfig, @@ -118,11 +124,6 @@ class ValConfig(BaseConfig): """Held-out sets. With ``data.holdout_docs > 0`` the training holdout is added as ``holdout``.""" -class PretrainDeploymentConfig(BaseConfig): - num_train_gpus: int = Field(1, ge=1) - """GPUs for the trainer, on this node.""" - - class PretrainConfig(BaseConfig): model: PretrainModelConfig = PretrainModelConfig() @@ -175,7 +176,11 @@ class PretrainConfig(BaseConfig): heartbeat: HeartbeatConfig | None = None """BetterStack heartbeat configuration.""" - deployment: PretrainDeploymentConfig = PretrainDeploymentConfig() + deployment: SFTDeploymentConfig = SingleNodeDeploymentConfig() + """``single_node`` (``num_train_gpus`` on this node) or ``multi_node`` (``num_train_nodes`` x ``gpus_per_node``, needs ``[slurm]``). The online-eval fields (``num_infer_*``) are not used.""" + + slurm: SlurmConfig | None = None + """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" dry_run: bool = False """Only validate and dump resolved configs, then exit.""" @@ -232,6 +237,11 @@ def auto_setup_tokenizer(self): self.tokenizer.trust_remote_code = self.model.trust_remote_code return self + @model_validator(mode="after") + def auto_setup_deployment(self): + setup_trainer_deployment(self) + return self + @model_validator(mode="after") def validate_monitors(self): if self.monitors.prime is not None: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index d2171d5580f..b279386ad75 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -1,7 +1,7 @@ import uuid import warnings from pathlib import Path -from typing import Annotated, Literal, TypeAlias +from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias from urllib.parse import urlparse from pydantic import AliasChoices, Field, model_validator @@ -36,6 +36,9 @@ ) from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource +if TYPE_CHECKING: + from prime_rl.configs.pretrain import PretrainConfig + class BaseDataConfig(BaseConfig): batch_size: int = Field(128, ge=1) @@ -216,6 +219,30 @@ class MultiNodeDeploymentConfig(BaseDeploymentConfig): ] +def setup_trainer_deployment(config: "SFTConfig | PretrainConfig") -> None: + """Shared by ``sft`` and ``pretrain``: a multi-node deployment needs SLURM and sets + ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the SLURM job name defaults to the + run name and the template to the bundled single- or multi-node trainer template.""" + deployment = config.deployment + if deployment.type == "multi_node": + if config.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + if deployment.nodes_per_fsdp_group is not None: + if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" + ) + config.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group + if config.slurm is not None: + if config.slurm.job_name is None: + config.slurm.job_name = config.run.name + if config.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + config.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" + + class SFTConfig(BaseConfig): model: ModelConfig = ModelConfig() @@ -410,12 +437,6 @@ def full_optimizer_offload_disables_grad_clipping(self): self.optim.max_norm = None return self - @model_validator(mode="after") - def validate_deployment(self): - if self.deployment.type == "multi_node" and self.slurm is None: - raise ValueError("Must use SLURM for multi-node deployment.") - return self - @model_validator(mode="after") def auto_setup_online_eval(self): """Wire online evals and the trainer-to-inference weight transport.""" @@ -638,25 +659,5 @@ def auto_setup_tokenizer(self): @model_validator(mode="after") def auto_setup_deployment(self): - if self.deployment.type == "multi_node": - if self.deployment.nodes_per_fsdp_group is not None: - if self.deployment.num_train_nodes % self.deployment.nodes_per_fsdp_group != 0: - raise ValueError( - f"deployment.num_train_nodes ({self.deployment.num_train_nodes}) must be divisible by " - f"deployment.nodes_per_fsdp_group ({self.deployment.nodes_per_fsdp_group})" - ) - self.model.dp_replicate = self.deployment.num_train_nodes // self.deployment.nodes_per_fsdp_group - return self - - @model_validator(mode="after") - def auto_setup_slurm_template(self): - if self.slurm is not None and self.slurm.job_name is None: - self.slurm.job_name = self.run.name - if self.slurm is not None and self.slurm.template_path is None: - templates_dir = find_package_resource("templates") - if templates_dir is not None: - if self.deployment.type == "single_node": - self.slurm.template_path = templates_dir / "single_node_sft.sbatch.j2" - else: - self.slurm.template_path = templates_dir / "multi_node_sft.sbatch.j2" + setup_trainer_deployment(self) return self diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py index e94ad716917..fea2fa609b4 100644 --- a/src/prime_rl/entrypoints/pretrain.py +++ b/src/prime_rl/entrypoints/pretrain.py @@ -1,13 +1,16 @@ import json import os import uuid +from pathlib import Path from prime_rl.configs.pretrain import PretrainConfig +from prime_rl.entrypoints.sft import submit_slurm_script, write_trainer_slurm_script from prime_rl.utils.config import cli, dump_resolved_config -from prime_rl.utils.logger import setup_logger +from prime_rl.utils.logger import get_logger, setup_logger from prime_rl.utils.pathing import ( format_config_message, format_log_message, + get_launcher_dir, prepare_attempt_dirs, validate_run_dir, write_launch_artifacts, @@ -21,6 +24,7 @@ ) PRETRAIN_CONFIG = "pretrain.json" +PRETRAIN_SBATCH = "pretrain.sbatch" def write_model_dir(config: PretrainConfig) -> None: @@ -56,6 +60,19 @@ def download_data(config: PretrainConfig) -> None: load_eval_documents(source.dataset, source.subset, source.split, source.revision, source.text_column) +def pretrain_slurm(config: PretrainConfig, config_path: Path, log_dir: Path) -> None: + logger = get_logger() + script_path = get_launcher_dir(config.run_dir) / PRETRAIN_SBATCH + write_trainer_slurm_script(config, "pretrain", config_path, log_dir, script_path) + logger.info(f"Wrote SLURM script to {script_path}") + num_nodes = config.deployment.num_train_nodes if config.deployment.type == "multi_node" else 1 + log_message = format_log_message(log_dir=log_dir, trainer=True, num_train_nodes=num_nodes) + if config.dry_run: + logger.success(f"Dry run complete. To submit manually:\n\n sbatch {script_path}\n\n{log_message}") + return + submit_slurm_script(script_path, log_message, logger) + + def pretrain(config: PretrainConfig): logger = setup_logger(config.log.level or "info", json_logging=config.log.json_logging) os.environ.setdefault("PRL_RUN_ID", uuid.uuid4().hex) @@ -86,8 +103,16 @@ def pretrain(config: PretrainConfig): download_data(config) config_path = config_dir / PRETRAIN_CONFIG - config_path.write_text(json.dumps(dump_resolved_config(config, exclude={"dry_run", "clean"}), indent=2)) + exclude = {"dry_run", "clean"} + if config.slurm is not None: + # A single-node job runs this launcher again on the node; the trainers of a + # multi-node job read the config directly. + exclude |= {"slurm", "deployment"} if config.deployment.type == "multi_node" else {"slurm"} + config_path.write_text(json.dumps(dump_resolved_config(config, exclude=exclude), indent=2)) logger.info(f"Configs:\n{format_config_message(config_dir, 'pretrain', [('Pretrain', config_path)])}") + if config.slurm is not None: + pretrain_slurm(config, config_path, log_dir) + return if config.dry_run: logger.success("Dry run complete. To start the run, remove --dry-run from your command.") return diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 1ec9b367b63..e48963fb06a 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -8,6 +8,7 @@ from prime_rl.configs.eval import SFTOnlineEvalConfig from prime_rl.configs.monitors import EvalMonitorsConfig, PrimeEvalMonitorConfig, TrainMonitorsConfig from prime_rl.configs.orchestrator import OnlineEvalSourceConfig +from prime_rl.configs.pretrain import PretrainConfig from prime_rl.configs.sft import SFTConfig from prime_rl.configs.shared import LogConfig from prime_rl.entrypoints.dashboard import ensure_dashboard, log_dashboard_url @@ -148,10 +149,17 @@ def write_eval_subconfigs(config: SFTConfig, config_dir: Path, strip_router: boo write_env_server_config(config_dir, "eval", source, config.log) -def write_slurm_script( - config: SFTConfig, config_path: Path, log_dir: Path, script_path: Path, prl_run_id: str, wandb_shared: bool +def write_trainer_slurm_script( + config: SFTConfig | PretrainConfig, + entrypoint: str, + config_path: Path, + log_dir: Path, + script_path: Path, + **template_vars, ) -> None: - """Write the SLURM script to disk.""" + """Write the SLURM script of a trainer entrypoint (``sft`` or ``pretrain``). Without + ``template_vars`` a multi-node job runs only trainer nodes; SFT passes its online-eval + variables on top.""" from jinja2 import Environment, FileSystemLoader assert config.slurm is not None @@ -164,80 +172,82 @@ def write_slurm_script( env = Environment(loader=FileSystemLoader(template_dirs), keep_trailing_newline=True) template = env.get_template(config.slurm.template_path.name) - trainer_env_vars = { - **DEFAULT_COMMON_ENV_VARS, - **DEFAULT_TRAINER_ENV_VARS, - **config.env_vars, - } - - if config.deployment.type == "single_node": - script = template.render( - **config.slurm.template_vars, - config_path=config_path, - config_dir=config_path.parent, - log_dir=log_dir, - output_dir=config.run_dir, - launcher_dir=get_launcher_dir(config.run_dir), - launcher_log_dir=get_launcher_log_dir(config.run_dir), - gpus_per_node=config.deployment.gpus_per_node, - ) - else: - online_eval = config.eval is not None - eval_vars = {} - if online_eval: - assert config.inference is not None - inference_env_vars = { - **DEFAULT_COMMON_ENV_VARS, - **DEFAULT_INFERENCE_ENV_VARS, - **config.env_vars, - **config.inference.env_vars, - } - eval_vars = { - "router": config.inference.router, - "router_port": config.inference.server.port, - "backend_port": config.inference.backend_port, - "data_parallel_rpc_port": config.inference.vllm.data_parallel_rpc_port, - "dp_per_node": config.deployment.gpus_per_node // config.inference.vllm.tensor_parallel_size, - "enable_expert_parallel": config.inference.vllm.enable_expert_parallel, - "inference_env_vars": inference_env_vars, - "online_eval_env_vars": { - **DEFAULT_COMMON_ENV_VARS, - "LOGURU_FORCE_COLORS": "1", - **config.env_vars, - }, - "eval_env_names": [source.resolved_name for source in eval_env_servers(config)], - } - script = template.render( - **config.slurm.template_vars, - config_path=config_path, - config_dir=config_path.parent, - log_dir=log_dir, - output_dir=config.run_dir, - launcher_dir=get_launcher_dir(config.run_dir), - launcher_log_dir=get_launcher_log_dir(config.run_dir), - trainer_env_vars=trainer_env_vars, + trainer_vars = {} + if config.deployment.type == "multi_node": + trainer_vars = dict( + trainer_env_vars={**DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, **config.env_vars}, num_nodes=config.deployment.num_train_nodes, num_train_nodes=config.deployment.num_train_nodes, - num_infer_nodes=config.deployment.num_infer_nodes if online_eval else 0, - gpus_per_node=config.deployment.gpus_per_node, + num_infer_nodes=0, ranks_filter=",".join(map(str, config.log.ranks_filter)), - prl_run_id=prl_run_id, - wandb_shared=wandb_shared, + prl_run_id=os.environ["PRL_RUN_ID"], + wandb_shared=False, run_name=config.run.name, - online_eval=online_eval, - use_nccl_broadcast=( - config.eval is not None - and config.weight_broadcast is not None - and config.weight_broadcast.type == "nccl" - ), - **eval_vars, + online_eval=False, + use_nccl_broadcast=False, ) + script = template.render( + **config.slurm.template_vars, + entrypoint=entrypoint, + config_path=config_path, + config_dir=config_path.parent, + log_dir=log_dir, + output_dir=config.run_dir, + launcher_dir=get_launcher_dir(config.run_dir), + launcher_log_dir=get_launcher_log_dir(config.run_dir), + gpus_per_node=config.deployment.gpus_per_node, + **{**trainer_vars, **template_vars}, + ) script_path.parent.mkdir(parents=True, exist_ok=True) get_launcher_log_dir(config.run_dir).mkdir(parents=True, exist_ok=True) script_path.write_text(script) +def submit_slurm_script(script_path: Path, log_message: str, logger) -> None: + logger.info(f"Submitting: sbatch {script_path}") + result = subprocess.run(["sbatch", str(script_path)], capture_output=True, text=True) + if result.returncode != 0: + logger.error(f"sbatch failed: {result.stderr.strip()}") + sys.exit(1) + logger.success(f"{result.stdout.strip()}\n\n{log_message}") + + +def write_slurm_script( + config: SFTConfig, config_path: Path, log_dir: Path, script_path: Path, wandb_shared: bool +) -> None: + """Write the SLURM script to disk, with the online-eval nodes of a multi-node run.""" + eval_vars = {} + if config.deployment.type == "multi_node" and config.eval is not None: + assert config.inference is not None + inference_env_vars = { + **DEFAULT_COMMON_ENV_VARS, + **DEFAULT_INFERENCE_ENV_VARS, + **config.env_vars, + **config.inference.env_vars, + } + eval_vars = { + "num_infer_nodes": config.deployment.num_infer_nodes, + "wandb_shared": wandb_shared, + "online_eval": True, + "use_nccl_broadcast": config.weight_broadcast is not None and config.weight_broadcast.type == "nccl", + "router": config.inference.router, + "router_port": config.inference.server.port, + "backend_port": config.inference.backend_port, + "data_parallel_rpc_port": config.inference.vllm.data_parallel_rpc_port, + "dp_per_node": config.deployment.gpus_per_node // config.inference.vllm.tensor_parallel_size, + "enable_expert_parallel": config.inference.vllm.enable_expert_parallel, + "inference_env_vars": inference_env_vars, + "online_eval_env_vars": { + **DEFAULT_COMMON_ENV_VARS, + "LOGURU_FORCE_COLORS": "1", + **config.env_vars, + }, + "eval_env_names": [source.resolved_name for source in eval_env_servers(config)], + } + write_trainer_slurm_script(config, "sft", config_path, log_dir, script_path, **eval_vars) + + def sft_slurm(config: SFTConfig): """Run SFT training and its online-eval deployment in one SLURM allocation.""" assert config.slurm is not None @@ -269,7 +279,7 @@ def sft_slurm(config: SFTConfig): write_eval_subconfigs(config, config_dir, strip_router=True) logger.info(f"Configs:\n{format_config_message(config_dir, 'sft', sft_config_components(config, config_dir))}") script_path = launcher_dir / SFT_SBATCH - write_slurm_script(config, config_path, log_dir, script_path, os.environ["PRL_RUN_ID"], wandb_shared) + write_slurm_script(config, config_path, log_dir, script_path, wandb_shared) logger.info(f"Wrote SLURM script to {script_path}") num_nodes = config.deployment.num_train_nodes if config.deployment.type == "multi_node" else 1 @@ -289,13 +299,7 @@ def sft_slurm(config: SFTConfig): dashboard_url = ensure_dashboard(config.output_dir, logger) if config.dashboard else None - logger.info(f"Submitting: sbatch {script_path}") - result = subprocess.run(["sbatch", str(script_path)], capture_output=True, text=True) - if result.returncode != 0: - logger.error(f"sbatch failed: {result.stderr.strip()}") - sys.exit(1) - - logger.success(f"{result.stdout.strip()}\n\n{log_message}") + submit_slurm_script(script_path, log_message, logger) log_dashboard_url(logger, dashboard_url) diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 index 32027951d71..6dcb0ca9e44 100644 --- a/src/prime_rl/templates/multi_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -71,7 +71,7 @@ export WANDB_RUN_ID=$PRL_RUN_ID export WANDB_SHARED_MODE=1 export WANDB_SHARED_PRIMARY=trainer export WANDB_SHARED_FINISHER=online-eval -export WANDB_PROGRAM="uv run sft" +export WANDB_PROGRAM="uv run {{ entrypoint }}" {% endif %} # Networking. Inference nodes occupy the first ranks, followed by trainer nodes. @@ -241,7 +241,7 @@ else --tee=3 \ --redirects=3 \ --local-ranks-filter={{ ranks_filter }} \ - -m prime_rl.trainer.sft.train \ + -m prime_rl.trainer.{{ entrypoint }}.train \ @ $TRAIN_CONFIG_PATH \ {%- if use_nccl_broadcast %} --weight-broadcast.host 0.0.0.0 \ diff --git a/src/prime_rl/templates/single_node_sft.sbatch.j2 b/src/prime_rl/templates/single_node_sft.sbatch.j2 index 9723eb491e9..d1183709f45 100644 --- a/src/prime_rl/templates/single_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/single_node_sft.sbatch.j2 @@ -41,4 +41,4 @@ uv sync --all-extras --all-packages # Pre-run command {{ pre_run_command }} {% endif %} -uv run sft @ {{ config_path }} +uv run {{ entrypoint }} @ {{ config_path }} diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index 7dc00a77993..f660684c7a1 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -319,7 +319,7 @@ def has_checkpoints(output_dir: Path) -> bool: # Launcher artifacts may exist before training starts. Everything else is treated as # artifacts of a previous run. -LAUNCHER_ARTIFACTS = ("configs", "launcher") +LAUNCHER_ARTIFACTS = ("configs", "launcher", "model") def has_run_artifacts(run_dir: Path) -> bool: diff --git a/tests/unit/train/pretrain/test_launcher.py b/tests/unit/train/pretrain/test_launcher.py new file mode 100644 index 00000000000..14e71e027cb --- /dev/null +++ b/tests/unit/train/pretrain/test_launcher.py @@ -0,0 +1,29 @@ +import re + +from prime_rl.configs.pretrain import PretrainConfig +from prime_rl.entrypoints.pretrain import pretrain +from prime_rl.utils.config import cli + + +def test_multi_node_dry_run_writes_trainer_sbatch(tmp_path, monkeypatch): + monkeypatch.setenv("PRL_RUN_ID", "test") + monkeypatch.setenv("PRL_RUN_NAME", "test") + config = PretrainConfig.model_validate( + { + "output_dir": tmp_path, + "data": {"type": "fake"}, + "deployment": {"type": "multi_node", "num_train_nodes": 2}, + "slurm": {}, + "dry_run": True, + } + ) + pretrain(config) + + script = (config.run_dir / "launcher" / "pretrain.sbatch").read_text() + assert "#SBATCH --nodes=2" in script + assert "-m prime_rl.trainer.pretrain.train" in script + # The trainers on every node parse the written config, which has no [deployment] or [slurm]. + config_path = re.search(r"^export TRAIN_CONFIG_PATH=(.+)$", script, re.M).group(1) + trainer_config = cli(PretrainConfig, args=["@", config_path]) + assert trainer_config.slurm is None + assert trainer_config.model == config.model From 9c9c8c0d0b8dfc116c2ae0db5a2440d8bc38ccd2 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 15:28:01 +0200 Subject: [PATCH 011/132] fix(pretrain): MoE sequence-wise balance loss per packed row, not per packed sample --- src/prime_rl/trainer/models/layers/moe.py | 6 +++--- src/prime_rl/trainer/pretrain/train.py | 24 +++++++++++++---------- tests/unit/train/models/test_moe.py | 2 ++ 3 files changed, 19 insertions(+), 13 deletions(-) diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index b580d118e8d..c8195cf1320 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -369,9 +369,9 @@ def __init__( ) self.register_buffer("routing_confidence_sum", torch.tensor(0.0, dtype=torch.float32), persistent=False) # Sequence-wise balance loss (DeepSeek-V3 eq. 17-20): with a coefficient set, every forward with grad - # stores the per-row loss in `seq_balance`, for the training loop to add with that coefficient. + # keeps the scores and selections, for the training loop, which knows the packed rows, to compute it. self.seq_aux_loss_coeff = seq_aux_loss_coeff - self.seq_balance: torch.Tensor | None = None + self.routing: tuple[torch.Tensor, torch.Tensor] | None = None def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None: self.token_dispatcher = token_dispatcher @@ -415,7 +415,7 @@ def forward( ) = self.router(x, routed_experts=routed_experts) if self.seq_aux_loss_coeff is not None and torch.is_grad_enabled(): - self.seq_balance = sequence_balance(scores, selected_experts_indices, bs) + self.routing = (scores, selected_experts_indices) # Accumulate expert usage for selection-bias updates and metrics. with torch.no_grad(): diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index a229c3b866a..ff7f2a95b2a 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -31,6 +31,7 @@ update_selection_biases, ) from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX +from prime_rl.trainer.models.layers.moe import sequence_balance from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter @@ -221,10 +222,10 @@ def train(config: PretrainConfig): moes_with_seq_balance = [moe for moe in iter_moe_blocks(model) if moe.seq_aux_loss_coeff is not None] bias_freeze_step = config.max_steps - round(config.moe_bias_freeze * config.max_steps) if config.max_steps else None - def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Forward pass returning (loss_sum, token_count) over unmasked tokens, and the MoE sequence-wise - balance loss summed over tokens (each token carries its row's loss, so after the division by the - token count it is the mean over rows).""" + def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]: + """Forward pass returning (loss_sum, token_count) over unmasked tokens, the MoE sequence-wise balance + loss summed over tokens (each token carries its row's loss, so after the division by the token count it + is the mean over rows), and the per-row balance of each MoE layer that computes it.""" input_ids = micro_batch["input_ids"].to("cuda", non_blocking=True) position_ids = micro_batch["position_ids"].to("cuda", non_blocking=True) target_ids = micro_batch["target_ids"].to("cuda", non_blocking=True) @@ -241,10 +242,13 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.T out = forward( model, input_ids, position_ids, seq_lens=seq_lens, labels=labels, seq_lens_are_pre_shard=cp_enabled ) + # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. + num_rows = config.data.micro_batch_size + seq_balances = [sequence_balance(*moe.routing, num_rows) for moe in moes_with_seq_balance] aux_loss_sum = torch.zeros((), device="cuda") - for moe in moes_with_seq_balance: - aux_loss_sum += moe.seq_aux_loss_coeff * moe.seq_balance.sum() * input_ids.shape[1] - return out["loss"], loss_mask.sum(dtype=torch.int64), aux_loss_sum + for moe, balance in zip(moes_with_seq_balance, seq_balances): + aux_loss_sum += moe.seq_aux_loss_coeff * balance.sum() * (input_ids.shape[1] // num_rows) + return out["loss"], loss_mask.sum(dtype=torch.int64), aux_loss_sum, seq_balances def run_evals(step: int) -> None: """Perplexity evals on the weights after ``step`` optimizer updates.""" @@ -299,7 +303,7 @@ def run_evals(step: int) -> None: step_local_num_tokens = 0 for micro_step, micro_batch in enumerate(micro_batches): step_local_num_tokens += micro_batch["input_ids"].shape[1] - local_loss_sum, batch_token_count, aux_loss_sum = compute_loss(micro_batch) + local_loss_sum, batch_token_count, aux_loss_sum, seq_balances = compute_loss(micro_batch) if gradient_manager is None: step_local_token_count += batch_token_count if torch.isnan(local_loss_sum.detach()): @@ -323,8 +327,8 @@ def run_evals(step: int) -> None: micro_max_mean += ( tokens_per_expert.amax(dim=1) / tokens_per_expert.mean(dim=1) ).mean() / grad_accum_steps - if moes_with_seq_balance: - seq_balance += torch.cat([moe.seq_balance.detach() for moe in moes_with_seq_balance]).mean() + if seq_balances: + seq_balance += torch.cat(seq_balances).detach().mean() forward_backward_time = time.perf_counter() - step_start_time expert_load_stats = {} diff --git a/tests/unit/train/models/test_moe.py b/tests/unit/train/models/test_moe.py index 5f622310a75..3a3341cd0b9 100644 --- a/tests/unit/train/models/test_moe.py +++ b/tests/unit/train/models/test_moe.py @@ -275,3 +275,5 @@ def test_sequence_balance(): # row 0 sends both tokens to expert 0: f = [2, 0], P = [0.9, 0.1]; row 1 is balanced: f = [1, 1], P = [0.5, 0.5] torch.testing.assert_close(sequence_balance(scores, selected, num_rows=2), torch.tensor([1.8, 1.0])) + # the same tokens as one sequence (a packed sample of both rows) hide row 0's imbalance: f = [1.5, 0.5], P = [0.7, 0.3] + torch.testing.assert_close(sequence_balance(scores, selected, num_rows=1), torch.tensor([1.2])) From 351d30e4402514f7af5cf2cdbea7423c5d67682f Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 15:45:52 +0200 Subject: [PATCH 012/132] feat(pretrain): CompleteP for MoE (experts per role, router lr / m_N under Muon) --- docs/pretraining.md | 5 ++- src/prime_rl/trainer/pretrain/completep.py | 47 ++++++++++++++-------- src/prime_rl/trainer/pretrain/train.py | 2 +- tests/unit/train/pretrain/test_init.py | 32 ++++++++++++++- 4 files changed, 65 insertions(+), 21 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 5bd89dfd60a..bed415a9c5e 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -58,8 +58,11 @@ With `optim.type = "muon"`, the hidden matrices and `o_proj`/`down_proj` are tra |---|---|---|---| | Hidden matrices | std / sqrt(m_N) | lr / sqrt(m_N) | wd sqrt(m_N) | | `o_proj`, `down_proj` | std / (sqrt(m_N) m_L) | lr / (sqrt(m_N) m_L) | wd sqrt(m_N) m_L | +| MoE router | std / sqrt(m_N) | lr / m_N | wd m_N | -CompleteP works with dense models, needs `model.offload = "none"`, and `scheduler.min_lr = 0` with the cosine scheduler. Keep `head_dim` and the ratios `intermediate_size / hidden_size` and `num_attention_heads / num_key_value_heads` fixed across shapes. +In MoE models, the expert matrices follow the hidden and `down_proj` rules per expert. The router maps the width to a fixed number of experts, so under Muon its learning rate falls as 1/m_N (with AdamW it follows the hidden rule). + +CompleteP needs `model.offload = "none"` and `scheduler.min_lr = 0` with the cosine scheduler. Keep `head_dim` and the ratios `intermediate_size / hidden_size` (and `moe_intermediate_size / hidden_size`) and `num_attention_heads / num_key_value_heads` fixed across shapes. For `sft` and `rl`, `random_init_mode` resolves to `skip`: the weights keep arbitrary memory contents, which is the fastest option for throughput benchmarks. diff --git a/src/prime_rl/trainer/pretrain/completep.py b/src/prime_rl/trainer/pretrain/completep.py index 8068af3b4cc..4828a1ee738 100644 --- a/src/prime_rl/trainer/pretrain/completep.py +++ b/src/prime_rl/trainer/pretrain/completep.py @@ -1,3 +1,5 @@ +from collections import defaultdict + from torch import nn from prime_rl.configs.trainer import AdamWConfig, InitConfig, MuonConfig @@ -5,7 +7,9 @@ from prime_rl.utils.weights import resolve_fqn -def completep_param_groups(model: nn.Module, init: InitConfig, optim: AdamWConfig | MuonConfig) -> list[dict]: +def completep_param_groups( + model: nn.Module, init: InitConfig, optim: AdamWConfig | MuonConfig, ep_enabled: bool +) -> list[dict]: """Optimizer parameter groups for CompleteP with alpha = 1 (Dey et al. 2025). ``optim.lr``, ``optim.eps`` and ``optim.weight_decay`` are the values for the base shape; each role scales them by the width and depth multipliers. The residual output projections and the LM head @@ -16,42 +20,51 @@ def completep_param_groups(model: nn.Module, init: InitConfig, optim: AdamWConfi ``rms_norm`` learning-rate adjustment, whose update size grows as sqrt(m_N); their learning rate is scaled by 1/sqrt(m_N) to compensate (and weight decay by sqrt(m_N), keeping lr * weight_decay fixed). Everything else is trained with AdamW, as with ``optim.type = "adamw"``. + + MoE experts follow the same roles per expert matrix (``rms_norm`` adjusts each expert by its + own shape). The router maps width to a fixed number of experts, so its spectral-norm target + sqrt(num_experts / d) shrinks as m_N^-1/2: with Muon its learning rate is scaled by 1/m_N; + with AdamW it follows the hidden rule. """ - if any("mlp.experts" in name for name, _ in model.named_parameters()): - raise ValueError("CompleteP is only implemented for dense models") m_n, m_l = completep_multipliers(init, model.config.hidden_size, model.config.num_hidden_layers) # (lr, eps, weight decay) scale per role for AdamW. PyTorch AdamW and Dion decay by lr * weight_decay. adam_scales = { "embedding": (1.0, 1 / m_n, 1.0), "hidden": (1 / m_n, 1 / (m_n * m_l), m_n), + "router": (1 / m_n, 1 / (m_n * m_l), m_n), "residual_out": (1 / (m_n * m_l), 1 / m_n, m_n * m_l), "output": (1 / m_n, 1.0, m_n), "norm": (1.0, 1 / m_n, 0.0), "bias": (1.0, 1 / m_n, 0.0), } # (lr, weight decay) scale per role for Muon with rms_norm adjustment. - muon_scales = {"hidden": (m_n**-0.5, m_n**0.5), "residual_out": (m_n**-0.5 / m_l, m_n**0.5 * m_l)} + muon_scales = { + "hidden": (m_n**-0.5, m_n**0.5), + "router": (1 / m_n, m_n), + "residual_out": (m_n**-0.5 / m_l, m_n**0.5 * m_l), + } - params = {role: [] for role in adam_scales} + params = defaultdict(list) for name, param in model.named_parameters(): if param.requires_grad: - params[param_role(resolve_fqn(model, name), param)].append(param) + name = resolve_fqn(model, name) + role = "router" if ".mlp.router." in name and param.ndim == 2 else param_role(name, param) + params[role, ".mlp.experts." in name].append(param) groups = [] - for role, role_params in params.items(): - if not role_params: - continue + for (role, experts), role_params in params.items(): if optim.type == "muon" and role in muon_scales: lr, wd = muon_scales[role] - groups.append( - dict( - params=role_params, - algorithm="muon", - adjust_lr="rms_norm", - lr=optim.lr * lr, - weight_decay=optim.weight_decay * wd, - ) + group = dict( + params=role_params, + algorithm="muon", + adjust_lr="rms_norm", + lr=optim.lr * lr, + weight_decay=optim.weight_decay * wd, ) + if experts and ep_enabled: + group["distributed_mesh_name"] = "dp_shard_mod_ep" + groups.append(group) continue lr, eps, wd = adam_scales[role] group = dict(params=role_params, lr=optim.lr * lr, weight_decay=optim.weight_decay * wd) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index ff7f2a95b2a..40d0e67acb7 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -155,7 +155,7 @@ def train(config: PretrainConfig): get_full_offload_dtype_policy(model, config.model) if config.model.offload == "full" else None ), param_groups=( - completep_param_groups(model, config.model.init, config.optim) + completep_param_groups(model, config.model.init, config.optim, parallel_dims.ep_enabled) if config.model.init.scheme == "completep" else None ), diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index e8935d21e8a..093c9e9b896 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -5,8 +5,9 @@ from torch.optim import AdamW from transformers import LlamaConfig, LlamaForCausalLM -from prime_rl.configs.trainer import AdamWConfig, InitConfig +from prime_rl.configs.trainer import AdamWConfig, InitConfig, MuonConfig from prime_rl.trainer.init import init_weights, normal_shard +from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM from prime_rl.trainer.pretrain.completep import completep_param_groups @@ -80,7 +81,7 @@ def explicit_scales(name: str) -> tuple[float, float, float]: ) optimizers = { "explicit": AdamW(explicit_groups, betas=(0.9, 0.95)), - "folded": AdamW(completep_param_groups(folded, init, optim), betas=(0.9, 0.95)), + "folded": AdamW(completep_param_groups(folded, init, optim, ep_enabled=False), betas=(0.9, 0.95)), } models = {"explicit": explicit, "folded": folded} @@ -94,3 +95,30 @@ def explicit_scales(name: str) -> tuple[float, float, float]: torch.testing.assert_close( folded(input_ids=tokens).logits, explicit(input_ids=tokens).logits, rtol=1e-4, atol=1e-4 ) + + +def test_completep_moe_groups_cover_experts_and_scale_the_router(): + config = Glm4MoeConfig( + vocab_size=128, + hidden_size=512, + intermediate_size=1024, + moe_intermediate_size=128, + num_attention_heads=8, + num_key_value_heads=2, + n_routed_experts=8, + num_experts_per_tok=2, + n_shared_experts=1, + num_hidden_layers=2, + first_k_dense_replace=1, + ) + config._attn_implementation = "flash_attention_2" + with torch.device("meta"): + model = Glm4MoeForCausalLM._from_config(config) + init = InitConfig(scheme="completep", base_hidden_size=256, base_num_layers=2) + groups = completep_param_groups(model, init, MuonConfig(lr=1.0), ep_enabled=False) + lr = {param: group["lr"] for group in groups for param in group["params"]} + assert sum(len(group["params"]) for group in groups) == len(lr) == len(list(model.parameters())) + moe = model.model.layers[1].mlp + # m_N = 2: Muon matrices get lr / sqrt(m_N) (residual outputs also / m_L = 1), the router lr / m_N. + assert lr[moe.experts.up_proj] == lr[moe.experts.down_proj] == 2**-0.5 + assert lr[moe.router.gate.weight] == 0.5 From 8d9a295243ef304ecf8bb045165806292ab7ea8c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 16:06:15 +0200 Subject: [PATCH 013/132] feat(pretrain): reduce step metrics once per step, MoE balance and grad norms per group, batched perplexity evals Per-step statistics are accumulated on the GPU and reduced with one SUM and one MAX all-reduce after the micro-batches, replacing the per-micro-step MoE all-reduces and host copies. Expert balance comes from the step's global counts (max/mean, min/mean, cv, cold, entropy, dead experts) and per-rank micro-batches (mean and max over ranks); router confidence, bias size and flips; gradient norm per parameter group; loss and grad-norm spike flags; data wait and peak memory as max and mean over ranks. Perplexity evals score micro_batch_size rows per forward. --- src/prime_rl/trainer/pretrain/evals/common.py | 28 +-- .../trainer/pretrain/evals/perplexity.py | 13 +- src/prime_rl/trainer/pretrain/metrics.py | 132 +++++++++++++ src/prime_rl/trainer/pretrain/train.py | 179 ++++++++++-------- tests/unit/train/pretrain/test_metrics.py | 34 ++++ 5 files changed, 292 insertions(+), 94 deletions(-) create mode 100644 src/prime_rl/trainer/pretrain/metrics.py create mode 100644 tests/unit/train/pretrain/test_metrics.py diff --git a/src/prime_rl/trainer/pretrain/evals/common.py b/src/prime_rl/trainer/pretrain/evals/common.py index e6bbeb5a078..db416185dba 100644 --- a/src/prime_rl/trainer/pretrain/evals/common.py +++ b/src/prime_rl/trainer/pretrain/evals/common.py @@ -1,5 +1,3 @@ -from typing import TypeVar - import torch import torch.distributed as dist from transformers.tokenization_utils import PreTrainedTokenizer @@ -7,8 +5,6 @@ from prime_rl.trainer.sft.data import Sample from prime_rl.trainer.world import get_world -T = TypeVar("T") - def bytes_to_unicode() -> dict[int, str]: """The byte -> character map of byte-level BPE tokenizers (GPT-2, Llama 3, Qwen, DeepSeek).""" @@ -41,16 +37,26 @@ def token_byte_lengths(tokenizer: PreTrainedTokenizer) -> torch.Tensor: return lengths -def local_rows(rows: list[T]) -> list[T | None]: - """This rank's share of ``rows``, padded with ``None`` so every rank runs the same number of - forward passes (each forward is an FSDP collective).""" +BATCH_KEYS = ("input_ids", "position_ids", "target_ids", "loss_mask", "seq_lens") + + +def local_batches(rows: list[Sample], batch_size: int) -> list[dict[str, torch.Tensor]]: + """This rank's share of ``rows``, ``batch_size`` rows per forward concatenated into one packed sample, like a + training micro-batch. Every rank runs the same number of forwards (each is an FSDP collective): a rank that + runs out of rows scores ``rows[0]`` with an empty loss mask.""" world = get_world() - per_rank = -(-len(rows) // world.world_size) - mine: list[T | None] = list(rows[world.rank :: world.world_size]) - return mine + [None] * (per_rank - len(mine)) + mine = rows[world.rank :: world.world_size] + batches = [] + for start in range(0, -(-len(rows) // world.world_size), batch_size): + chunk = mine[start : start + batch_size] + batch = to_device({key: [value for row in chunk or rows[:1] for value in row[key]] for key in BATCH_KEYS}) + if not chunk: + batch["loss_mask"].zero_() + batches.append(batch) + return batches -def to_device(row: Sample) -> dict[str, torch.Tensor]: +def to_device(row: dict[str, list[int]]) -> dict[str, torch.Tensor]: return { "input_ids": torch.tensor(row["input_ids"], device="cuda").unsqueeze(0), "position_ids": torch.tensor(row["position_ids"], device="cuda").unsqueeze(0), diff --git a/src/prime_rl/trainer/pretrain/evals/perplexity.py b/src/prime_rl/trainer/pretrain/evals/perplexity.py index cded029475c..f2b9b9ba541 100644 --- a/src/prime_rl/trainer/pretrain/evals/perplexity.py +++ b/src/prime_rl/trainer/pretrain/evals/perplexity.py @@ -10,7 +10,7 @@ from prime_rl.trainer.model import forward from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX from prime_rl.trainer.pretrain.data import StreamPacker, load_eval_documents, tokenize_document -from prime_rl.trainer.pretrain.evals.common import all_reduce_sum, local_rows, to_device +from prime_rl.trainer.pretrain.evals.common import all_reduce_sum, local_batches from prime_rl.trainer.sft.data import Sample @@ -39,15 +39,16 @@ def load_perplexity_set(source: ValSourceConfig, tokenizer: PreTrainedTokenizer, @torch.no_grad() -def evaluate_perplexity(model: nn.Module, sets: list[PerplexitySet], byte_lengths: torch.Tensor) -> dict[str, float]: - """Loss (nats per token), perplexity and bits per byte for every set.""" +def evaluate_perplexity( + model: nn.Module, sets: list[PerplexitySet], byte_lengths: torch.Tensor, batch_size: int +) -> dict[str, float]: + """Loss (nats per token), perplexity and bits per byte for every set, ``batch_size`` rows per forward.""" byte_lengths = byte_lengths.to("cuda") metrics = {} for eval_set in sets: totals = torch.zeros(3, dtype=torch.float64, device="cuda") # loss, tokens, bytes - for row in local_rows(eval_set.rows): - batch = to_device(row if row is not None else eval_set.rows[0]) - loss_mask = batch["loss_mask"] if row is not None else torch.zeros_like(batch["loss_mask"]) + for batch in local_batches(eval_set.rows, batch_size): + loss_mask = batch["loss_mask"] labels = batch["target_ids"].masked_fill(~loss_mask, IGNORE_INDEX) out = forward(model, batch["input_ids"], batch["position_ids"], seq_lens=batch["seq_lens"], labels=labels) totals[0] += out["loss"].double() diff --git a/src/prime_rl/trainer/pretrain/metrics.py b/src/prime_rl/trainer/pretrain/metrics.py new file mode 100644 index 00000000000..227c9292bc5 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/metrics.py @@ -0,0 +1,132 @@ +import math +from collections import defaultdict, deque + +import torch +import torch.distributed as dist +from torch import Tensor, nn +from torch.distributed.tensor import DTensor, Partial, Replicate + +from prime_rl.trainer.model import iter_moe_blocks + + +def reduce_stats(sums: dict[str, Tensor], maxes: dict[str, Tensor], group: dist.ProcessGroup) -> dict[str, Tensor]: + """Sum or max per-rank statistics over ``group`` with one all-reduce per kind, and copy them to the CPU once.""" + packed = {} + for values, op in ((sums, dist.ReduceOp.SUM), (maxes, dist.ReduceOp.MAX)): + if values: + flat = torch.cat([value.reshape(-1).to(torch.float64) for value in values.values()]) + dist.all_reduce(flat, op=op, group=group) + packed[op] = flat + flat = torch.cat(list(packed.values())).cpu() + reduced, offset = {}, 0 + for values in (sums, maxes): + for name, value in values.items(): + reduced[name] = flat[offset : offset + value.numel()].reshape(value.shape) + offset += value.numel() + return reduced + + +def take_routing_counts(model: nn.Module) -> tuple[Tensor, Tensor]: + """This rank's `[num_moe_layers, num_experts]` expert counts and `[num_moe_layers]` routing-confidence sums + since the last call, and reset them.""" + moes = list(iter_moe_blocks(model)) + counts = torch.stack([moe.tokens_per_expert for moe in moes]) + confidence = torch.stack([moe.routing_confidence_sum for moe in moes]) + for moe in moes: + moe.tokens_per_expert.zero_() + moe.routing_confidence_sum.zero_() + return counts, confidence + + +def expert_load_metrics(counts: Tensor) -> dict[str, float]: + """Balance of the `[num_moe_layers, num_experts]` expert counts of a whole step, per layer, summarized as the + mean and the worst layer. `max_mean` / `min_mean`: busiest / least-loaded expert over the mean load; `cv`: + std / mean; `cold_frac`: experts under 0.1x the mean load; `entropy`: entropy of the experts' shares of the + assignments over log(num_experts) (1 = uniform).""" + mean = counts.mean(dim=1) + share = counts / counts.sum(dim=1, keepdim=True) + per_layer = { + "max_mean": (counts.amax(dim=1) / mean, "max"), + "min_mean": (counts.amin(dim=1) / mean, "min"), + "cv": (counts.std(dim=1, correction=0) / mean, "max"), + "cold_frac": ((counts < 0.1 * mean[:, None]).double().mean(dim=1), "max"), + "entropy": (-(share * share.clamp_min(1e-30).log()).sum(dim=1) / math.log(counts.shape[1]), "min"), + } + metrics = {} + for name, (values, worst) in per_layer.items(): + metrics[f"expert_load/{name}/mean"] = values.mean().item() + metrics[f"expert_load/{name}/{worst}"] = (values.amax() if worst == "max" else values.amin()).item() + return metrics + + +class DeadExperts: + """Experts that received no tokens over the last ``window`` steps.""" + + def __init__(self, window: int = 100): + self.counts: deque[Tensor] = deque(maxlen=window) + + def update(self, counts: Tensor) -> int: + self.counts.append(counts) + return int((torch.stack(list(self.counts)).sum(dim=0) == 0).sum()) + + +class SpikeDetector: + """Flags a value more than ``threshold`` standard deviations above the mean of the previous ``window`` values.""" + + def __init__(self, window: int = 128, threshold: float = 6.0): + self.values: deque[float] = deque(maxlen=window) + self.threshold = threshold + + def update(self, value: float) -> float: + spike = 0.0 + if len(self.values) >= 32: + history = torch.tensor(list(self.values), dtype=torch.float64) + spike = float(value > history.mean().item() + self.threshold * history.std().item()) + self.values.append(value) + return spike + + +def param_group(name: str) -> str: + """The parameter group a parameter's gradient norm is logged under.""" + parts = name.split(".") + if "embed_tokens" in parts: + return "embed" + if "lm_head" in parts: + return "lm_head" + if any("norm" in part for part in parts): + return "norm" + if "router" in parts: + return "router" + if "shared_expert" in parts: + return "shared_experts" + if "experts" in parts: + return "experts" + if "self_attn" in parts: + return "attn" + return "mlp" + + +def grad_norms_by_group(model: nn.Module) -> dict[str, Tensor]: + """Gradient L2 norm per parameter group (see `param_group`). The squared norms of the local shards are summed + over the sharded mesh dimensions only (replicated dimensions hold the same gradient), one collective per mesh + and placement.""" + by_layout: dict = defaultdict(lambda: defaultdict(list)) + for name, param in model.named_parameters(): + grad = param.grad + if grad is None: + continue + layout = (grad.device_mesh, tuple(grad.placements)) if isinstance(grad, DTensor) else None + by_layout[layout][param_group(name)].append(grad.to_local() if isinstance(grad, DTensor) else grad) + squares: dict[str, Tensor] = {} + for layout, groups in by_layout.items(): + local = torch.stack( + [torch.stack(torch._foreach_norm(grads)).double().square().sum() for grads in groups.values()] + ) + if layout is not None: + mesh, placements = layout + local = DTensor.from_local( + local, mesh, [Partial() if placement.is_shard() else Replicate() for placement in placements] + ).full_tensor() + for group, square in zip(groups, local): + squares[group] = squares.get(group, 0) + square + return {group: square.sqrt() for group, square in squares.items()} diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index ff7f2a95b2a..083b7f80ab3 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -19,10 +19,7 @@ from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager from prime_rl.trainer.model import ( forward, - get_expert_load_stats, get_full_offload_dtype_policy, - get_global_moe_stats, - get_load_balance_stats, is_tt_moe_model, iter_moe_blocks, reshard_module, @@ -36,6 +33,14 @@ from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter from prime_rl.trainer.pretrain.completep import completep_param_groups +from prime_rl.trainer.pretrain.metrics import ( + DeadExperts, + SpikeDetector, + expert_load_metrics, + grad_norms_by_group, + reduce_stats, + take_routing_counts, +) from prime_rl.trainer.pretrain.data import load_text_dataset, setup_text_dataloader, split_holdout from prime_rl.trainer.pretrain.evals.common import token_byte_lengths from prime_rl.trainer.pretrain.evals.perplexity import build_perplexity_set, evaluate_perplexity, load_perplexity_set @@ -217,7 +222,6 @@ def train(config: PretrainConfig): cp_group = parallel_dims.world_mesh["cp"].get_group() if cp_enabled else None cp_size = parallel_dims.cp dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() - ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None is_moe_model = is_tt_moe_model(model) moes_with_seq_balance = [moe for moe in iter_moe_blocks(model) if moe.seq_aux_loss_coeff is not None] bias_freeze_step = config.max_steps - round(config.moe_bias_freeze * config.max_steps) if config.max_steps else None @@ -252,12 +256,12 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.T def run_evals(step: int) -> None: """Perplexity evals on the weights after ``step`` optimizer updates.""" - metrics = evaluate_perplexity(model, perplexity_sets, byte_lengths) + metrics = evaluate_perplexity(model, perplexity_sets, byte_lengths, config.data.micro_batch_size) # The no-grad eval forwards leave the [lm_head, norm] FSDP group unsharded, and their # routing must not count toward the next step's MoE statistics. reshard_module(model) if is_moe_model: - get_load_balance_stats(model) + take_routing_counts(model) logger.success(f"Evals | Step {step} | " + " | ".join(f"{k} {v:.4f}" for k, v in sorted(metrics.items()))) asyncio.run(monitors.log({**metrics, "step": step}, step=step)) @@ -266,29 +270,18 @@ def run_evals(step: int) -> None: gc_handler = GarbageCollection(config.gc.interval) if config.gc else None max_memory = torch.cuda.mem_get_info()[1] / 1024**3 + num_dp_cp_ranks = dist.get_world_size(dp_cp_group) + dead_experts = DeadExperts() + loss_spikes, grad_norm_spikes = SpikeDetector(), SpikeDetector() + last_bias_moves = None logger.info(f"Starting training loop (max_steps={config.max_steps or 'infinite'})") while True: - torch.cuda.reset_peak_memory_stats() if gc_handler is not None: gc_handler.run(progress.step) is_last_step = config.max_steps is not None and progress.step >= config.max_steps step_start_time = time.perf_counter() - step_loss_sum = torch.tensor(0.0, device="cuda") - nan_loss_count = torch.tensor(0, device="cuda") - moe_stats = ( - { - "max_vio/mean": torch.tensor(0.0), - "max_vio/max": torch.tensor(0.0), - "routing_confidence/mean": torch.tensor(0.0), - } - if is_moe_model - else {} - ) - if gradient_manager is None: - micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) - step_local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") - else: + if gradient_manager is not None: micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] local_token_count = sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches) global_step_token_count = torch.tensor(local_token_count, dtype=torch.int64, device="cuda") @@ -297,65 +290,91 @@ def run_evals(step: int) -> None: grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=True) - step_tokens_per_expert = 0 - micro_max_mean = 0.0 - seq_balance = 0.0 - step_local_num_tokens = 0 - for micro_step, micro_batch in enumerate(micro_batches): - step_local_num_tokens += micro_batch["input_ids"].shape[1] + # Per-rank statistics, kept on the GPU and reduced once after the micro-batches (`reduce_stats`). + zero = torch.zeros((), device="cuda", dtype=torch.float64) + sums = {"loss": zero.clone(), "loss_tokens": zero.clone(), "nan_count": zero.clone(), "tokens": zero.clone()} + maxes = {} + data_time = 0.0 + for micro_step in range(grad_accum_steps): + data_start_time = time.perf_counter() + micro_batch = next(dataiter) if gradient_manager is None else micro_batches[micro_step] + data_time += time.perf_counter() - data_start_time + sums["tokens"] += micro_batch["input_ids"].shape[1] local_loss_sum, batch_token_count, aux_loss_sum, seq_balances = compute_loss(micro_batch) - if gradient_manager is None: - step_local_token_count += batch_token_count - if torch.isnan(local_loss_sum.detach()): - nan_loss_count += 1 - logger.warning("Local loss is nan, excluding this micro step from backward") - scaled_loss = torch.nan_to_num(local_loss_sum, nan=0.0) / grad_accum_steps - else: - step_loss_sum += local_loss_sum.detach() - scaled_loss = (local_loss_sum + aux_loss_sum) / grad_accum_steps + sums["loss_tokens"] += batch_token_count + # A NaN micro-batch is left out of the loss and the backward, without a host sync. + is_nan = torch.isnan(local_loss_sum.detach()) + sums["nan_count"] += is_nan + sums["loss"] += torch.where(is_nan, 0.0, local_loss_sum.detach()) + scaled_loss = torch.where(is_nan, torch.nan_to_num(local_loss_sum, nan=0.0), local_loss_sum + aux_loss_sum) + scaled_loss = scaled_loss / grad_accum_steps begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) scaled_loss.backward() finish_backward(gradient_manager) if is_moe_model: - micro_moe_stats, tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) - for name, value in micro_moe_stats.items(): - moe_stats[f"{name}/mean"] += value / grad_accum_steps - if name == "max_vio": - moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) - step_tokens_per_expert += tokens_per_expert - micro_max_mean += ( - tokens_per_expert.amax(dim=1) / tokens_per_expert.mean(dim=1) - ).mean() / grad_accum_steps + counts, confidence = take_routing_counts(model) + sums["expert_counts"] = sums.get("expert_counts", 0) + counts + sums["routing_confidence"] = sums.get("routing_confidence", 0) + confidence + # One rank's micro-batch: the imbalance that expert parallelism and the dropless dispatch see. + micro_max_mean = counts.amax(dim=1) / counts.mean(dim=1) + sums["micro_max_mean"] = sums.get("micro_max_mean", 0) + micro_max_mean.mean() + maxes["micro_max_mean_max"] = torch.maximum( + maxes.get("micro_max_mean_max", zero), micro_max_mean.amax() + ) if seq_balances: - seq_balance += torch.cat(seq_balances).detach().mean() + rows = torch.cat(seq_balances).detach() + sums["seq_balance"] = sums.get("seq_balance", 0) + rows.sum() + sums["seq_balance_rows"] = sums.get("seq_balance_rows", 0) + rows.numel() + sums["data_time"] = zero + data_time + maxes["data_time_max"] = zero + data_time + # Peak memory since the previous step's reduction: this step's forward/backward and the previous optimizer step. + peak_memory = torch.cuda.max_memory_reserved() / 1024**3 + torch.cuda.reset_peak_memory_stats() + sums["peak_memory"] = zero + peak_memory + maxes["peak_memory_max"] = zero + peak_memory + stats = reduce_stats(sums, maxes, dp_cp_group) forward_backward_time = time.perf_counter() - step_start_time - expert_load_stats = {} - if is_moe_model: - expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) - # This rank's micro-batches: the imbalance that expert parallelism and the dropless dispatch see. - expert_load_stats["expert_load/micro_max_mean"] = micro_max_mean.item() - if moes_with_seq_balance: - expert_load_stats["expert_load/seq_balance"] = seq_balance.item() / grad_accum_steps - if bias_freeze_step is None or progress.step <= bias_freeze_step: - update_selection_biases(model, step_tokens_per_expert) - biases = [moe.router.selection_bias for moe in iter_moe_blocks(model) if moe.load_balance_coeff is not None] - if biases: - expert_load_stats["expert_load/selection_bias_absmax"] = torch.stack(biases).abs().max().item() if gradient_manager is None: - global_step_token_count = step_local_token_count.clone() - dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) - global_token_count_val = global_step_token_count.item() + global_token_count_val = int(stats["loss_tokens"]) grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val scale_gradients_(None, model, grad_scale) + batch_loss = stats["loss"].item() / global_token_count_val + if stats["nan_count"] > 0: + logger.warning(f"{int(stats['nan_count'])} micro-batches had a NaN loss and were left out of this step") - dist.all_reduce(step_loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) - dist.all_reduce(nan_loss_count, op=dist.ReduceOp.SUM) - batch_loss = (step_loss_sum / global_token_count_val).item() - nan_loss_count = nan_loss_count.item() + moe_metrics = {} + if is_moe_model: + counts = stats["expert_counts"] + routed_tokens = counts.sum(dim=1) / next(iter_moe_blocks(model)).router.top_k + moe_metrics = { + **expert_load_metrics(counts), + "expert_load/dead": dead_experts.update(counts), + "expert_load/micro_max_mean/mean": stats["micro_max_mean"].item() + / (grad_accum_steps * num_dp_cp_ranks), + "expert_load/micro_max_mean/max": stats["micro_max_mean_max"].item(), + "router/confidence": (stats["routing_confidence"] / routed_tokens).mean().item(), + } + if "seq_balance" in stats: + moe_metrics["loss/seq_balance"] = (stats["seq_balance"] / stats["seq_balance_rows"]).item() + biases = [moe.router.selection_bias for moe in iter_moe_blocks(model) if moe.load_balance_coeff is not None] + if bias_freeze_step is None or progress.step <= bias_freeze_step: + previous_biases = [bias.clone() for bias in biases] + update_selection_biases(model, counts.to("cuda", torch.float32)) + if biases: + # Share of experts whose bias moved in the opposite direction to the previous step's move. + moves = torch.stack([bias - previous for bias, previous in zip(biases, previous_biases)]).sign() + if last_bias_moves is not None: + moe_metrics["router/bias/flip_frac"] = (moves * last_bias_moves < 0).float().mean().item() + last_bias_moves = moves + if biases: + stacked = torch.stack(biases) + moe_metrics["router/bias/absmax"] = stacked.abs().max().item() + moe_metrics["router/bias/std"] = stacked.std(dim=1).mean().item() + group_grad_norms = grad_norms_by_group(model) grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm) if config.optim.max_norm else None optimizer.step() optimizer.zero_grad() @@ -364,20 +383,21 @@ def run_evals(step: int) -> None: # The last step's checkpoint is written once after the loop. save_ckpt_time = 0.0 + disk_metrics = {} if config.ckpt and config.ckpt.interval and progress.step % config.ckpt.interval == 0 and not is_last_step: save_ckpt_start_time = time.perf_counter() ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) ckpt_manager.maybe_clean() save_ckpt_time = time.perf_counter() - save_ckpt_start_time + disk_metrics = get_ckpt_disk_metrics(config.run_dir) eval_start_time = time.perf_counter() if config.val is not None and (is_last_step or progress.step % config.val.interval == 0): run_evals(progress.step) eval_time = time.perf_counter() - eval_start_time - global_num_tokens = torch.tensor(step_local_num_tokens, dtype=torch.int64, device="cuda") - dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) - num_tokens = global_num_tokens.item() + # Every CP rank counts its rows in full. + num_tokens = int(stats["tokens"]) // cp_size progress.total_tokens += num_tokens dataset_progress = get_dataset_progress(dataloader) progress.total_samples = dataset_progress["step"] @@ -386,40 +406,45 @@ def run_evals(step: int) -> None: perf_counter.count_tokens(num_tokens) throughput = perf_counter.get_step_tokens_per_second(num_tokens, step_time) mfu = perf_counter.get_step_mfu(num_tokens, step_time) - peak_memory = torch.cuda.max_memory_reserved() / 1024**3 + grad_norm = grad_norm.item() if grad_norm is not None else None step_message = f"Step {progress.step} | {step_time:.2f}s | Loss {batch_loss:.4f}" if grad_norm is not None: step_message += f" | Grad. Norm {grad_norm:.4f}" step_message += ( f" | LR {current_lr:.2e} | Throughput {throughput:.0f} tokens/s | MFU {mfu:.1f}%" - f" | Peak Mem. {peak_memory:.1f}/{max_memory:.1f} GiB" + f" | Peak Mem. {stats['peak_memory_max'].item():.1f}/{max_memory:.1f} GiB" ) logger.success(step_message) metrics = { "loss/mean": batch_loss, "loss/perplexity": math.exp(min(batch_loss, 20)), - "loss/nan_count": nan_loss_count, + "loss/nan_count": int(stats["nan_count"]), + "loss/spike": loss_spikes.update(batch_loss), "optim/lr": current_lr, + **{f"optim/grad_norm/{group}": norm.item() for group, norm in group_grad_norms.items()}, "progress/epoch": dataset_progress["epoch"], "progress/num_samples": progress.total_samples, "progress/num_tokens": progress.total_tokens, "perf/throughput": throughput, "perf/throughput_per_gpu": throughput / world.world_size, - "perf/peak_memory": peak_memory, + "perf/peak_memory/max": stats["peak_memory_max"].item(), + "perf/peak_memory/mean": stats["peak_memory"].item() / num_dp_cp_ranks, "perf/mfu": mfu, "time/step": step_time, "time/forward_backward": forward_backward_time, + "time/data/max": stats["data_time_max"].item(), + "time/data/mean": stats["data_time"].item() / num_dp_cp_ranks, "time/save_ckpt": save_ckpt_time, "time/evals": eval_time, - **get_ckpt_disk_metrics(config.run_dir), - **{name: value.item() for name, value in moe_stats.items()}, - **expert_load_stats, + **disk_metrics, + **moe_metrics, "step": progress.step, } if grad_norm is not None: - metrics["optim/grad_norm"] = grad_norm.item() + metrics["optim/grad_norm"] = grad_norm + metrics["optim/grad_norm_spike"] = grad_norm_spikes.update(grad_norm) asyncio.run(monitors.log(metrics, step=progress.step)) if heart is not None: diff --git a/tests/unit/train/pretrain/test_metrics.py b/tests/unit/train/pretrain/test_metrics.py new file mode 100644 index 00000000000..55e8de6b3ba --- /dev/null +++ b/tests/unit/train/pretrain/test_metrics.py @@ -0,0 +1,34 @@ +import torch +from torch import nn + +from prime_rl.trainer.pretrain.metrics import expert_load_metrics, grad_norms_by_group, param_group + + +def test_expert_load_metrics_per_layer_summaries(): + counts = torch.tensor([[1.0, 1.0, 1.0, 1.0], [4.0, 0.0, 0.0, 0.0]], dtype=torch.float64) + metrics = expert_load_metrics(counts) + assert metrics["expert_load/max_mean/mean"] == 2.5 # layers: 1 and 4 + assert metrics["expert_load/max_mean/max"] == 4.0 + assert metrics["expert_load/min_mean/min"] == 0.0 + assert metrics["expert_load/cold_frac/max"] == 0.75 + assert metrics["expert_load/entropy/mean"] == 0.5 # uniform layer 1, one-hot layer 0 + + +def test_grad_norms_by_group_add_up_to_the_total_norm(): + model = nn.ModuleDict( + { + "embed_tokens": nn.Embedding(8, 4), + "self_attn": nn.Linear(4, 4), + "input_layernorm": nn.LayerNorm(4), + "lm_head": nn.Linear(4, 8), + } + ) + for param in model.parameters(): + param.grad = torch.randn_like(param) + norms = grad_norms_by_group(model) + assert set(norms) == {"embed", "attn", "norm", "lm_head"} + total = torch.nn.utils.get_total_norm([param.grad for param in model.parameters()]) + assert torch.allclose(torch.stack(list(norms.values())).square().sum().sqrt().float(), total) + assert param_group("model.layers.3.mlp.router.gate.weight") == "router" + assert param_group("model.layers.3.mlp.shared_expert.w1") == "shared_experts" + assert param_group("model.layers.3.mlp.experts.w1") == "experts" From 68287c410c0801b00eacf89d0e6c8bf3dad8524e Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 16:12:20 +0200 Subject: [PATCH 014/132] test(pretrain): set Muon weight decay in the CompleteP MoE test --- tests/unit/train/pretrain/test_init.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index 093c9e9b896..ba8123c0dc3 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -115,7 +115,7 @@ def test_completep_moe_groups_cover_experts_and_scale_the_router(): with torch.device("meta"): model = Glm4MoeForCausalLM._from_config(config) init = InitConfig(scheme="completep", base_hidden_size=256, base_num_layers=2) - groups = completep_param_groups(model, init, MuonConfig(lr=1.0), ep_enabled=False) + groups = completep_param_groups(model, init, MuonConfig(lr=1.0, weight_decay=0.1), ep_enabled=False) lr = {param: group["lr"] for group in groups for param in group["params"]} assert sum(len(group["params"]) for group in groups) == len(lr) == len(list(model.parameters())) moe = model.model.layers[1].mlp From 2e6744245abf22ffc003a17e0ddaa51fb5cdb3b9 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 16:29:23 +0200 Subject: [PATCH 015/132] feat(pretrain): token-shard data, local architecture spec, batch size in tokens Simplifies the pretraining path: - Data: the only real data type is datatrove token shards (the loader of the closed draft #3869, moved out of SFT): per-source folders read as one stream, mixture phases at fractions of max_steps, a startup check of the planned epochs per source, and a resume state (global row counter and tokens read per source) that does not depend on the world size. The first document of every source must end with the tokenizer's EOS, which catches shards made with another tokenizer. The HF text loader, stream packer and holdout split are gone; data.type = "fake" draws random tokens through the same path. - Batch sizes are in tokens (data.batch_tokens, data.micro_batch_tokens), so changing seq_len keeps the tokens per step and the number of steps. - Model: [model.arch] is the whole architecture (model_type plus the fields of its prime-rl config class; unknown fields are rejected). vocab_size (len(tokenizer) rounded up to 128) and the BOS, EOS and pad ids come from the tokenizer. No baseline HF model is downloaded; model.name is /model. - Init: the trainer initializes the weights itself unless it resumes, so model.debug.random_init_mode, its resolvers in the SFT and RL configs and the forced debug flag are gone. model.init moves into the pretrain config. - [val] reads held-out token shards per source and computes loss, perplexity and bits per byte through the training loss code, so it works with context parallelism and needs no in-memory eval sets. - Defaults: max_steps, optim.lr, data and tokenizer.name are required; the optimizer defaults to Muon. --- configs/debug/fake/pretrain.toml | 11 +- docs/pretraining.md | 72 ++-- docs/training.md | 2 +- examples/pretrain/glm4_moe.toml | 73 ++++ examples/pretrain/llama3/150m.toml | 55 --- examples/pretrain/llama3/17m.toml | 55 --- examples/pretrain/llama3/1b.toml | 54 --- examples/pretrain/llama3/500m.toml | 54 --- examples/pretrain/llama3/70m.toml | 55 --- .../src/prime_rl/configs/pretrain.py | 217 ++++++----- .../src/prime_rl/configs/sft.py | 7 - .../src/prime_rl/configs/trainer.py | 28 +- src/prime_rl/entrypoints/pretrain.py | 49 +-- src/prime_rl/trainer/init.py | 2 +- src/prime_rl/trainer/model.py | 9 +- .../pretrain/{evals/common.py => bytes.py} | 32 -- src/prime_rl/trainer/pretrain/data.py | 343 ++++++++++-------- .../trainer/pretrain/evals/perplexity.py | 61 ---- src/prime_rl/trainer/pretrain/train.py | 102 +++--- .../pretrain/{test_evals.py => test_bytes.py} | 2 +- tests/unit/train/pretrain/test_data.py | 124 ++++--- tests/unit/train/pretrain/test_model_dir.py | 34 ++ 22 files changed, 611 insertions(+), 830 deletions(-) create mode 100644 examples/pretrain/glm4_moe.toml delete mode 100644 examples/pretrain/llama3/150m.toml delete mode 100644 examples/pretrain/llama3/17m.toml delete mode 100644 examples/pretrain/llama3/1b.toml delete mode 100644 examples/pretrain/llama3/500m.toml delete mode 100644 examples/pretrain/llama3/70m.toml rename src/prime_rl/trainer/pretrain/{evals/common.py => bytes.py} (54%) delete mode 100644 src/prime_rl/trainer/pretrain/evals/perplexity.py rename tests/unit/train/pretrain/{test_evals.py => test_bytes.py} (86%) create mode 100644 tests/unit/train/pretrain/test_model_dir.py diff --git a/configs/debug/fake/pretrain.toml b/configs/debug/fake/pretrain.toml index 60ba6332843..db0302142b5 100644 --- a/configs/debug/fake/pretrain.toml +++ b/configs/debug/fake/pretrain.toml @@ -1,17 +1,24 @@ max_steps = 5 [model] -name = "meta-llama/Llama-3.2-1B" compile = "None" [model.arch] +model_type = "llama" hidden_size = 256 num_hidden_layers = 2 num_attention_heads = 4 num_key_value_heads = 2 intermediate_size = 1024 +[tokenizer] +name = "PrimeIntellect/Qwen3-0.6B" + [data] type = "fake" -batch_size = 4 seq_len = 512 +batch_tokens = 2048 +micro_batch_tokens = 512 + +[optim] +lr = 1e-3 diff --git a/docs/pretraining.md b/docs/pretraining.md index 8eed93f536b..884340dbfe1 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -1,35 +1,32 @@ # Pretraining -`uv run pretrain` trains a model from scratch on raw text. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, scheduler, checkpointing and monitors as the SFT trainer. It runs on a single node. +`uv run pretrain` trains a model from scratch on pre-tokenized token shards. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, scheduler, checkpointing and monitors as the SFT trainer. It runs on a single node. ```bash -uv run pretrain @ examples/pretrain/llama3/17m.toml +uv run pretrain @ examples/pretrain/glm4_moe.toml ``` ## Model -`model.name` names a baseline HF model. The baseline provides the architecture class, the tokenizer and the vocabulary. `[model.arch]` overrides fields of its HF config to set the size: +`[model.arch]` is the whole architecture: a `model_type` from `prime_rl.trainer.models` and the fields of its config class. Unknown fields are rejected. The vocabulary size (`len(tokenizer)` rounded up to a multiple of 128) and the BOS, EOS and pad ids come from the tokenizer, so they cannot disagree with it. ```toml -[model] -name = "meta-llama/Llama-3.2-1B" - [model.arch] -hidden_size = 384 +model_type = "glm4_moe" +hidden_size = 512 num_hidden_layers = 8 -num_attention_heads = 6 -num_key_value_heads = 2 -head_dim = 64 -intermediate_size = 1536 -``` +n_routed_experts = 64 +num_experts_per_tok = 6 -The launcher writes the resulting config and the tokenizer to `/model` and trains from there, so the trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all see an ordinary model directory. The trainer logs the non-embedding, embedding (input embedding and LM head) and total parameter counts at startup. +[tokenizer] +name = "tokenizers/deepseek-v4-flash-base" +``` -The [Llama 3 ladder](https://github.com/PrimeIntellect-ai/prime-rl/tree/main/examples/pretrain/llama3) covers 17M, 70M, 150M, 500M and 1B non-embedding parameters. +The launcher writes the config and the tokenizer to `/model` and trains from there, so the trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all see an ordinary model directory. The trainer logs the non-embedding, embedding (input embedding and LM head) and total parameter counts at startup. ## Initialization -Pretraining sets `model.debug.random_init = true`: weights are not downloaded or loaded, and `random_init_mode` resolves to `init`. Set `model.debug.random_init = false` to continue pretraining from the baseline's weights. +Weights start from a random init, unless the run resumes from a checkpoint. | Field | Meaning | |---|---| @@ -37,35 +34,42 @@ Pretraining sets `model.debug.random_init = true`: weights are not downloaded or | `model.init.std` | Base standard deviation (default 0.02). | | `model.init.seed` | Seed. Every rank initializes only its own shard, and the values depend only on the seed and the parameter name, not on the number of GPUs. | -For `sft` and `rl`, `random_init_mode` resolves to `skip`: the weights keep arbitrary memory contents, which is the fastest option for throughput benchmarks. - ## Data -`data.type = "text"` (the default) reads raw documents from a HF dataset column (`data.text_column`, default `text`). Only the selected subsets are downloaded. Each document is tokenized with the tokenizer's BOS and a closing EOS. Documents are packed back to back into rows of `data.seq_len` tokens and continue into the next row, so no token is truncated or padded. A step trains on exactly `data.batch_size * data.seq_len` tokens. +`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of a source are read in name order as one stream; documents continue across rows, and attention restarts at every document and every row. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. -| Field | Meaning | -|---|---| -| `data.document_masking` | Attention stays within each document and positions restart at every document start (default). If false, a row is one sequence. Attention never crosses a row. | -| `data.micro_batch_size` | Rows per forward pass. | -| `data.holdout_docs` | Hold the last N documents (in dataset order) out of training. With `[val]` set, they are evaluated as the `holdout` source. | -| `data.subsets`, `data.splits`, `data.probabilities` | Load and interleave several subsets/splits. | +Each row comes from one source, drawn with the weights of the current mixture phase. Phases start at fractions of `max_steps`. At startup the trainer computes how many passes over each source the run makes and fails above `data.max_epochs`. -`data.type = "fake"` trains on generated tokens, for throughput tests. +```toml +[data] +type = "tokens" +seq_len = 4096 +batch_tokens = 262144 # per optimizer step, over all GPUs +micro_batch_tokens = 65536 # per forward pass on one GPU +max_epochs = 3 + +[data.sources] +web = { path = "data/web" } +code = { path = "data/code" } + +[[data.phases]] +weights = { web = 0.8, code = 0.2 } + +[[data.phases]] +start = 0.9 +weights = { web = 0.5, code = 0.5 } +``` -## Evals +Batch sizes are in tokens, so changing `seq_len` keeps the tokens per step and the number of steps. The resume state is the global row counter and the tokens read per source, the same on every rank. `data.type = "fake"` trains on random tokens, for smoke and throughput tests. -Evals run after the optimizer update, so the metrics at step N describe the weights after N updates, the same weights as the step-N checkpoint. With `eval_on_start` (default), they also run before the first step (logged as step 0), and they always run at the last step. They require `model.cp = 1`. +## Evals -**Perplexity (`[val]`).** Each `[[val.source]]` is a held-out text set. It is tokenized and packed once at startup (with document masking, up to `max_tokens`), so every eval sees the same tokens. Each source logs `val/{name}/loss` (nats per token), `val/{name}/ppl` and `val/{name}/bpb` (bits per UTF-8 byte, comparable across tokenizers). +`[val]` computes the held-out loss (nats per token), perplexity and bits per UTF-8 byte (comparable across tokenizers) of every source in `val.sources`, through the same loss code as training, so it also works with context parallelism. Each eval reads the same leading windows of every source: as many as fill whole micro-batches on every data-parallel rank. Evals run after the optimizer update, so the metrics at step N describe the same weights as the step-N checkpoint. With `eval_on_start` (default), they also run before the first step (logged as step 0), and they always run at the last step. ```toml [val] interval = 500 -[[val.source]] -name = "wikitext" -dataset = "Salesforce/wikitext" -subset = "wikitext-103-raw-v1" -split = "test" +[val.sources] +web = { path = "data/web-val" } ``` - diff --git a/docs/training.md b/docs/training.md index c3867a32e71..5d32a3a0033 100644 --- a/docs/training.md +++ b/docs/training.md @@ -34,7 +34,7 @@ This page covers everything you need to launch, observe, checkpoint, and recover |---|---|---| | `uv run rl` | Wraps the trainer, orchestrator, and inference server in one launch from a merged TOML. | The default for any RL run. Runs locally for single-node experiments; submits to SLURM for single- or multi-node when `[slurm]` is set (see [Scaling § SLURM](scaling.md#slurm)). | | `uv run sft` | Supervised fine-tuning on a HF dataset. | Launches torchrun internally; never call torchrun directly. | -| `uv run pretrain` | Pretraining from scratch on raw text. | Single node; see [Pretraining](pretraining.md). | +| `uv run pretrain` | Pretraining from scratch on token shards. | Single node; see [Pretraining](pretraining.md). | | `uv run inference` | vLLM server. | Always use this entrypoint over `vllm serve` — it adds `/update_weights`, `/load_lora_adapter`, and `/init_broadcaster`. | | `uv run orchestrator` | Standalone orchestrator process. | Pair with a separately-launched inference server and one `env-server` per source. | | `uv run eval` | Multi-env evals against a live inference server. | One epoch per source, pinned (or adaptive) concurrency, cursor checkpoints + `--resume`, dashboard + optional platform upload; see [Eval](eval.md). | diff --git a/examples/pretrain/glm4_moe.toml b/examples/pretrain/glm4_moe.toml new file mode 100644 index 00000000000..09373e7782e --- /dev/null +++ b/examples/pretrain/glm4_moe.toml @@ -0,0 +1,73 @@ +# GLM-4.5-style MoE (64 routed experts, 6 active, 2 shared) with d = 512 and 8 layers, trained from scratch on +# the Dolma 3 mixture tokenized with the DeepSeek-V4 tokenizer. +# 4096 steps x 256K tokens = 1.1B tokens. The learning rate is the optimum of the MoE base-shape sweep. +# +# uv run pretrain @ examples/pretrain/glm4_moe.toml + +max_steps = 4096 + +[model] +offload = "none" +ac = "None" + +[model.arch] +model_type = "glm4_moe" +hidden_size = 512 +num_hidden_layers = 8 +num_attention_heads = 8 +num_key_value_heads = 2 +head_dim = 64 +intermediate_size = 2048 # the dense first layer +moe_intermediate_size = 256 +n_routed_experts = 64 +num_experts_per_tok = 6 +n_shared_experts = 2 +first_k_dense_replace = 1 +norm_topk_prob = true +routed_scaling_factor = 2.5 +attention_bias = false +use_qk_norm = false +partial_rotary_factor = 1.0 +rope_parameters = { rope_type = "default", rope_theta = 500000.0, partial_rotary_factor = 1.0 } +tie_word_embeddings = false + +[tokenizer] +name = "tokenizers/deepseek-v4-flash-base" # a local copy of deepseek-ai/DeepSeek-V4-Flash-Base without config.json + +[data] +type = "tokens" +seq_len = 4096 +batch_tokens = 262144 +micro_batch_tokens = 65536 +max_epochs = 3 + +[data.sources] +web = { path = "data/dolma3-50b-dsv4/web" } +pdfs = { path = "data/dolma3-50b-dsv4/pdfs" } +code = { path = "data/dolma3-50b-dsv4/code" } +math = { path = "data/dolma3-50b-dsv4/math" } +reference = { path = "data/dolma3-50b-dsv4/reference" } + +[[data.phases]] +weights = { web = 0.60, pdfs = 0.20, code = 0.12, math = 0.05, reference = 0.03 } + +[optim] +type = "muon" +lr = 0.00390625 +max_norm = 1.0 + +[scheduler] +type = "linear" +warmup_steps = 410 +decay_steps = 3686 + +[val] +interval = 512 + +[val.sources] +web = { path = "data/dolma3-50b-dsv4/web-val" } +code = { path = "data/dolma3-50b-dsv4/code-val" } +math = { path = "data/dolma3-50b-dsv4/math-val" } + +[ckpt] +interval = 1024 diff --git a/examples/pretrain/llama3/150m.toml b/examples/pretrain/llama3/150m.toml deleted file mode 100644 index 2bbe3c3a3b6..00000000000 --- a/examples/pretrain/llama3/150m.toml +++ /dev/null @@ -1,55 +0,0 @@ -# Llama 3 architecture with ~147.1M non-embedding parameters, trained from scratch on FineWeb-Edu. -# 5,700 steps x 256 rows x 2048 tokens = 3.0B tokens (~20 tokens per non-embedding parameter). -# The learning rate is a starting point, not a tuned value. -# -# uv run pretrain @ examples/pretrain/llama3/150m.toml - -max_steps = 5700 - -[model] -name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) -offload = "none" -ac = "None" # small enough to keep all activations - -[model.arch] -hidden_size = 768 -num_hidden_layers = 17 -num_attention_heads = 12 -num_key_value_heads = 4 -head_dim = 64 -intermediate_size = 3072 -tie_word_embeddings = false -rope_parameters = { rope_type = "default", rope_theta = 500000.0 } - -[model.init] -scheme = "scaled" - -[data] -name = "HuggingFaceFW/fineweb-edu" -subsets = ["sample-10BT"] -seq_len = 2048 -batch_size = 256 -micro_batch_size = 8 -holdout_docs = 10000 - -[optim] -lr = 1e-3 -betas2 = 0.95 - -[scheduler] -type = "cosine" -warmup_steps = 300 -min_lr = 0.0001 - -[val] -interval = 285 - -[[val.source]] -name = "wikitext" -dataset = "Salesforce/wikitext" -subset = "wikitext-103-raw-v1" -split = "test" - - -[ckpt] -interval = 570 diff --git a/examples/pretrain/llama3/17m.toml b/examples/pretrain/llama3/17m.toml deleted file mode 100644 index c4694e6e265..00000000000 --- a/examples/pretrain/llama3/17m.toml +++ /dev/null @@ -1,55 +0,0 @@ -# Llama 3 architecture with ~17.3M non-embedding parameters, trained from scratch on FineWeb-Edu. -# 700 steps x 256 rows x 2048 tokens = 0.4B tokens (~20 tokens per non-embedding parameter). -# The learning rate is a starting point, not a tuned value. -# -# uv run pretrain @ examples/pretrain/llama3/17m.toml - -max_steps = 700 - -[model] -name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) -offload = "none" -ac = "None" # small enough to keep all activations - -[model.arch] -hidden_size = 384 -num_hidden_layers = 8 -num_attention_heads = 6 -num_key_value_heads = 2 -head_dim = 64 -intermediate_size = 1536 -tie_word_embeddings = false -rope_parameters = { rope_type = "default", rope_theta = 500000.0 } - -[model.init] -scheme = "scaled" - -[data] -name = "HuggingFaceFW/fineweb-edu" -subsets = ["sample-10BT"] -seq_len = 2048 -batch_size = 256 -micro_batch_size = 16 -holdout_docs = 10000 - -[optim] -lr = 3e-3 -betas2 = 0.95 - -[scheduler] -type = "cosine" -warmup_steps = 50 -min_lr = 0.0003 - -[val] -interval = 100 - -[[val.source]] -name = "wikitext" -dataset = "Salesforce/wikitext" -subset = "wikitext-103-raw-v1" -split = "test" - - -[ckpt] -interval = 100 diff --git a/examples/pretrain/llama3/1b.toml b/examples/pretrain/llama3/1b.toml deleted file mode 100644 index 9aa25f9e8bb..00000000000 --- a/examples/pretrain/llama3/1b.toml +++ /dev/null @@ -1,54 +0,0 @@ -# Llama 3 architecture with ~973.1M non-embedding parameters, trained from scratch on FineWeb-Edu. -# 37,200 steps x 256 rows x 2048 tokens = 19.5B tokens (~20 tokens per non-embedding parameter). -# The learning rate is a starting point, not a tuned value. -# -# uv run pretrain @ examples/pretrain/llama3/1b.toml - -max_steps = 37200 - -[model] -name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) -offload = "none" - -[model.arch] -hidden_size = 2048 -num_hidden_layers = 16 -num_attention_heads = 32 -num_key_value_heads = 8 -head_dim = 64 -intermediate_size = 8192 -tie_word_embeddings = false -rope_parameters = { rope_type = "default", rope_theta = 500000.0 } - -[model.init] -scheme = "scaled" - -[data] -name = "HuggingFaceFW/fineweb-edu" -subsets = ["sample-100BT"] -seq_len = 2048 -batch_size = 256 -micro_batch_size = 2 -holdout_docs = 10000 - -[optim] -lr = 4e-4 -betas2 = 0.95 - -[scheduler] -type = "cosine" -warmup_steps = 1850 -min_lr = 4e-05 - -[val] -interval = 1860 - -[[val.source]] -name = "wikitext" -dataset = "Salesforce/wikitext" -subset = "wikitext-103-raw-v1" -split = "test" - - -[ckpt] -interval = 3720 diff --git a/examples/pretrain/llama3/500m.toml b/examples/pretrain/llama3/500m.toml deleted file mode 100644 index c8169205f18..00000000000 --- a/examples/pretrain/llama3/500m.toml +++ /dev/null @@ -1,54 +0,0 @@ -# Llama 3 architecture with ~495.5M non-embedding parameters, trained from scratch on FineWeb-Edu. -# 19,000 steps x 256 rows x 2048 tokens = 10.0B tokens (~20 tokens per non-embedding parameter). -# The learning rate is a starting point, not a tuned value. -# -# uv run pretrain @ examples/pretrain/llama3/500m.toml - -max_steps = 19000 - -[model] -name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) -offload = "none" - -[model.arch] -hidden_size = 1280 -num_hidden_layers = 21 -num_attention_heads = 20 -num_key_value_heads = 4 -head_dim = 64 -intermediate_size = 5120 -tie_word_embeddings = false -rope_parameters = { rope_type = "default", rope_theta = 500000.0 } - -[model.init] -scheme = "scaled" - -[data] -name = "HuggingFaceFW/fineweb-edu" -subsets = ["sample-100BT"] -seq_len = 2048 -batch_size = 256 -micro_batch_size = 4 -holdout_docs = 10000 - -[optim] -lr = 6e-4 -betas2 = 0.95 - -[scheduler] -type = "cosine" -warmup_steps = 950 -min_lr = 6e-05 - -[val] -interval = 950 - -[[val.source]] -name = "wikitext" -dataset = "Salesforce/wikitext" -subset = "wikitext-103-raw-v1" -split = "test" - - -[ckpt] -interval = 1900 diff --git a/examples/pretrain/llama3/70m.toml b/examples/pretrain/llama3/70m.toml deleted file mode 100644 index 6352b7a95fe..00000000000 --- a/examples/pretrain/llama3/70m.toml +++ /dev/null @@ -1,55 +0,0 @@ -# Llama 3 architecture with ~70.8M non-embedding parameters, trained from scratch on FineWeb-Edu. -# 2,800 steps x 256 rows x 2048 tokens = 1.5B tokens (~20 tokens per non-embedding parameter). -# The learning rate is a starting point, not a tuned value. -# -# uv run pretrain @ examples/pretrain/llama3/70m.toml - -max_steps = 2800 - -[model] -name = "meta-llama/Llama-3.2-1B" # baseline: architecture class, tokenizer (128,256 tokens) -offload = "none" -ac = "None" # small enough to keep all activations - -[model.arch] -hidden_size = 640 -num_hidden_layers = 12 -num_attention_heads = 10 -num_key_value_heads = 2 -head_dim = 64 -intermediate_size = 2560 -tie_word_embeddings = false -rope_parameters = { rope_type = "default", rope_theta = 500000.0 } - -[model.init] -scheme = "scaled" - -[data] -name = "HuggingFaceFW/fineweb-edu" -subsets = ["sample-10BT"] -seq_len = 2048 -batch_size = 256 -micro_batch_size = 8 -holdout_docs = 10000 - -[optim] -lr = 1.5e-3 -betas2 = 0.95 - -[scheduler] -type = "cosine" -warmup_steps = 150 -min_lr = 0.00015 - -[val] -interval = 140 - -[[val.source]] -name = "wikitext" -dataset = "Salesforce/wikitext" -subset = "wikitext-103-raw-v1" -split = "test" - - -[ckpt] -interval = 280 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 71a896d315e..56e973e1cbc 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -1,3 +1,4 @@ +import math import uuid from pathlib import Path from typing import Annotated, Any, Literal, TypeAlias @@ -5,14 +6,13 @@ from pydantic import Field, model_validator from prime_rl.configs.monitors import TrainMonitorsConfig -from prime_rl.configs.sft import BaseDataConfig, FakeDataConfig from prime_rl.configs.shared import EnvVars, HeartbeatConfig, ResumeConfig, RunConfig, TrainerLogConfig from prime_rl.configs.trainer import ( - AdamWConfig, CheckpointConfig, ConstantSchedulerConfig, GCConfig, ModelConfig, + MuonConfig, OptimizerConfig, SchedulerConfig, TokenizerConfig, @@ -20,93 +20,131 @@ ) from prime_rl.utils.config import BaseConfig, default_output_dir +TOKENIZER_FIELDS = ("vocab_size", "bos_token_id", "eos_token_id", "pad_token_id") + + +class InitConfig(BaseConfig): + """Random weight initialization.""" + + scheme: Literal["normal", "scaled"] = "normal" + """``normal`` draws every weight matrix and the embedding from N(0, std). ``scaled`` additionally divides the std of the residual output projections (``o_proj``, ``down_proj``) by sqrt(2 * num_layers). Norm weights start at 1, biases at 0.""" + + std: float = Field(0.02, gt=0) + """Base standard deviation.""" + + seed: int = 0 + """Seed for the initialization. The values depend only on the seed, not on the number of GPUs.""" + class PretrainModelConfig(ModelConfig): - arch: dict[str, Any] = {} - """HF config fields that override the config of the baseline model ``name``, e.g. ``hidden_size`` or ``num_hidden_layers``. The baseline provides the architecture class, the tokenizer and the vocabulary. The launcher writes the resulting config and tokenizer to ``/model`` and trains from there.""" + arch: dict[str, Any] + """The architecture: ``model_type`` (a model in ``prime_rl.trainer.models``, e.g. ``glm4_moe``) and the fields of its config class. ``vocab_size`` and the special token ids come from the tokenizer. The launcher writes the config and the tokenizer to ``/model``, which ``name`` points at.""" + + init: InitConfig = InitConfig() + """Random weight initialization. Used unless the run resumes from a checkpoint.""" + @model_validator(mode="after") + def validate_arch(self): + if "model_type" not in self.arch: + raise ValueError("model.arch needs a model_type, e.g. 'glm4_moe'") + if derived := [key for key in TOKENIZER_FIELDS if key in self.arch]: + raise ValueError(f"model.arch must not set {derived}: they come from the tokenizer") + return self -class TextDataConfig(BaseDataConfig): - type: Literal["text"] = "text" - name: str = "HuggingFaceFW/fineweb-edu" - """HF dataset name or local path.""" +class TokenSourceConfig(BaseConfig): + path: Path + """Folder of datatrove tokenized shards: ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``. The files are read in name order as one token stream.""" - subsets: list[str] | None = None - """Subsets (HF configs) to load.""" + max_tokens: int | None = Field(None, ge=1) + """Tokens available from this source: the first ``max_tokens`` of its stream. Defaults to all tokens in the shards.""" - splits: list[str] | None = None - """Splits to load. Defaults to ``train``.""" - revision: str | None = None - """HF dataset revision.""" +class MixturePhaseConfig(BaseConfig): + start: float = Field(0.0, ge=0, lt=1) + """Fraction of ``max_steps`` at which this phase starts.""" - probabilities: list[float] | None = None - """Sampling probabilities when interleaving several subsets/splits.""" + weights: dict[str, float] + """Token share of each source in this phase. Shares are non-negative and sum to 1; unlisted sources get 0.""" - stopping_strategy: Literal["first_exhausted", "all_exhausted"] = "all_exhausted" - """Stopping strategy when interleaving several subsets/splits.""" + @model_validator(mode="after") + def validate_weights(self): + if any(weight < 0 for weight in self.weights.values()) or not math.isclose(sum(self.weights.values()), 1.0): + raise ValueError(f"Phase weights must be non-negative and sum to 1, got {self.weights}") + return self - text_column: str = "text" - """Column holding the raw document text.""" - shuffle: bool = True - """Shuffle the documents at the start of each epoch.""" +class BaseBatchConfig(BaseConfig): + seq_len: int = Field(4096, ge=1) + """Tokens per row. Documents continue across rows; attention restarts at every document and every row.""" - seed: int = 0 - """Shuffle seed. Each epoch adds its index to it.""" + batch_tokens: int = Field(ge=1) + """Tokens per optimizer step, over all GPUs. Fixed when ``seq_len`` changes, so the number of steps stays the same.""" - seq_len: int = Field(2048, ge=1) - """Tokens per row. Documents are packed back to back and continue across rows; nothing is truncated or padded.""" + micro_batch_tokens: int = Field(ge=1) + """Tokens per forward pass on one GPU.""" - batch_size: int = Field(256, ge=1) - """Rows per optimizer step, so a step trains on ``batch_size * seq_len`` tokens.""" + num_workers: int = Field(1, ge=1) + """Dataloader worker processes per GPU.""" - document_masking: bool = True - """Restrict attention to the current document and restart positions at every document boundary. If false, a row is one sequence.""" + seed: int = 0 + """Seed for choosing each row's source (and the fake tokens).""" - holdout_docs: int = Field(0, ge=0) - """Hold the last N documents (in dataset order, before shuffling) out of training. With ``[val]`` set, they are evaluated as the ``holdout`` source.""" + @property + def batch_size(self) -> int: + """Rows per optimizer step.""" + return self.batch_tokens // self.seq_len + + @property + def micro_batch_size(self) -> int: + """Rows per forward pass.""" + return self.micro_batch_tokens // self.seq_len @model_validator(mode="after") - def validate_subsets_and_splits(self): - if self.subsets is not None and self.splits is not None and len(self.subsets) != len(self.splits): - raise ValueError("data.subsets and data.splits must have the same length") - sources = len(self.subsets or self.splits or [None]) - if self.probabilities is not None and len(self.probabilities) != sources: - raise ValueError("data.probabilities needs one entry per subset/split") + def validate_batch_tokens(self): + if self.batch_tokens % self.seq_len or self.micro_batch_tokens % self.seq_len: + raise ValueError("data.batch_tokens and data.micro_batch_tokens must be multiples of data.seq_len") + if self.batch_tokens % self.micro_batch_tokens: + raise ValueError("data.batch_tokens must be a multiple of data.micro_batch_tokens") return self -PretrainDataConfig: TypeAlias = Annotated[FakeDataConfig | TextDataConfig, Field(discriminator="type")] +class FakeDataConfig(BaseBatchConfig): + """Uniformly random token ids, for smoke tests.""" + + type: Literal["fake"] = "fake" -class ValSourceConfig(BaseConfig): - name: str - """Label for the metrics (``val/{name}/loss``, ``/ppl``, ``/bpb``).""" +class TokensDataConfig(BaseBatchConfig): + """Rows of pre-tokenized datatrove shards, each row drawn from one source of a weighted mixture.""" - dataset: str - """HF dataset name or local path.""" + type: Literal["tokens"] = "tokens" - subset: str | None = None - """HF config to load.""" + sources: dict[str, TokenSourceConfig] + """Token sources by name. The name labels the source's metrics.""" - split: str = "validation" - """Split to load.""" + phases: list[MixturePhaseConfig] + """Mixture phases in order. The first starts at 0.""" - revision: str | None = None - """HF dataset revision.""" + max_epochs: float = Field(1.0, gt=0) + """Most passes over a source's available tokens that the run (``max_steps``) may make. Checked at startup.""" + + @model_validator(mode="after") + def validate_phases(self): + starts = [phase.start for phase in self.phases] + if not starts or starts[0] != 0 or starts != sorted(set(starts)): + raise ValueError(f"Phase starts must begin at 0 and increase, got {starts}") + for phase in self.phases: + if unknown := phase.weights.keys() - self.sources.keys(): + raise ValueError(f"Phase weights name unknown sources {sorted(unknown)}") + return self - text_column: str = "text" - """Column holding the raw document text.""" - max_tokens: int = Field(10_000_000, ge=1) - """Token budget. Documents are taken in dataset order until it is reached.""" +PretrainDataConfig: TypeAlias = Annotated[FakeDataConfig | TokensDataConfig, Field(discriminator="type")] class ValConfig(BaseConfig): - """Perplexity evals: loss, perplexity and bits per byte on fixed held-out token sets. - Each set is tokenized and packed once at startup, so every eval sees the same tokens.""" + """Held-out loss, perplexity and bits per byte, computed like the training loss on held-out token shards.""" interval: int = Field(500, ge=1) """Evaluate every N steps, after that step's optimizer update.""" @@ -114,8 +152,9 @@ class ValConfig(BaseConfig): eval_on_start: bool = True """Evaluate the initial weights before the first step (logged as step 0).""" - source: list[ValSourceConfig] = [] - """Held-out sets. With ``data.holdout_docs > 0`` the training holdout is added as ``holdout``.""" + sources: dict[str, TokenSourceConfig] + """Held-out sources by name, logged as ``val//loss``, ``/ppl`` and ``/bpb``. Each eval reads the same + windows: as many as fill whole micro-batches on every data-parallel rank.""" class PretrainDeploymentConfig(BaseConfig): @@ -124,19 +163,21 @@ class PretrainDeploymentConfig(BaseConfig): class PretrainConfig(BaseConfig): - model: PretrainModelConfig = PretrainModelConfig() + model: PretrainModelConfig env_vars: EnvVars = {} """Extra environment variables for the trainer processes.""" - tokenizer: TokenizerConfig = TokenizerConfig() + tokenizer: TokenizerConfig + """The tokenizer (``tokenizer.name`` is required). It sets the vocabulary and special tokens of the model.""" - data: PretrainDataConfig = TextDataConfig() + data: PretrainDataConfig val: ValConfig | None = None - """Perplexity evals. If None, none run.""" + """Held-out evals. If None, none run.""" - optim: OptimizerConfig = AdamWConfig() + optim: OptimizerConfig = MuonConfig() + """Optimizer. ``optim.lr`` is required.""" scheduler: SchedulerConfig = ConstantSchedulerConfig() @@ -163,8 +204,8 @@ class PretrainConfig(BaseConfig): matmul_precision: Literal["highest", "high", "medium"] = "high" """Precision for float32 matrix multiplications (``torch.set_float32_matmul_precision``).""" - max_steps: int | None = None - """Maximum training steps. If None, runs indefinitely.""" + max_steps: int = Field(ge=1) + """Training steps.""" gc: GCConfig | None = GCConfig() """Deterministic garbage collection every N steps. Set to null for Python's default GC.""" @@ -185,25 +226,9 @@ def run_dir(self) -> Path: assert self.run.dir is not None # resolved at construction return self.output_dir / self.run.dir - @model_validator(mode="before") - @classmethod - def default_random_init(cls, data): - """Pretraining starts from randomly initialized weights unless ``model.debug.random_init = false`` - (continued pretraining from ``model.name``).""" - if isinstance(data, dict): - model = data.setdefault("model", {}) - if isinstance(model, dict): - debug = model.setdefault("debug", {}) - if isinstance(debug, dict): - debug.setdefault("random_init", True) - return data - @model_validator(mode="after") def resolve_autos(self): - """Resolve the shared ``auto`` fields for pretraining: real random init, a trainable fp32 - MoE router, and weight decay 0.1.""" - if self.model.debug.random_init_mode == "auto": - self.model.debug.random_init_mode = "init" + """Resolve the shared ``auto`` fields for pretraining: a trainable fp32 MoE router and weight decay 0.1.""" if self.model.moe_router_dtype == "auto": self.model.moe_router_dtype = "float32" if self.model.freeze_moe_router == "auto": @@ -212,22 +237,32 @@ def resolve_autos(self): self.optim.weight_decay = 0.1 return self + @model_validator(mode="after") + def validate_lr(self): + if "lr" not in self.optim.model_fields_set: + raise ValueError("optim.lr is required") + return self + @model_validator(mode="after") def auto_setup_run_identity(self): if self.run.name is None: - dataset = str(getattr(self.data, "name", "fake")).split("/")[-1] - model = self.model.name.split("/")[-1] - self.run.name = "--".join([dataset, model, uuid.uuid4().hex[:8]]).lower() + self.run.name = f"{self.model.arch['model_type']}--{uuid.uuid4().hex[:8]}" if self.run.dir is None: self.run.dir = self.run.name if self.monitors.wandb is not None and self.monitors.wandb.name is None: self.monitors.wandb.name = self.run.name return self + @model_validator(mode="after") + def auto_setup_model_dir(self): + """The model is trained from ``/model``, written by the launcher from ``model.arch``.""" + self.model.name = (self.run_dir / "model").as_posix() + return self + @model_validator(mode="after") def auto_setup_tokenizer(self): if self.tokenizer.name is None: - self.tokenizer.name = self.model.name + raise ValueError("tokenizer.name is required") if self.tokenizer.trust_remote_code is None: self.tokenizer.trust_remote_code = self.model.trust_remote_code return self @@ -244,15 +279,7 @@ def validate_cp(self): if self.data.seq_len % self.model.cp != 0: raise ValueError("data.seq_len must be divisible by model.cp") if self.data.micro_batch_size != 1: - raise ValueError("data.micro_batch_size must be 1 with context parallelism") - if self.val is not None: - raise ValueError("[val] does not support context parallelism yet") - return self - - @model_validator(mode="after") - def validate_holdout(self): - if self.val is not None and self.data.type == "text" and self.data.holdout_docs == 0 and not self.val.source: - raise ValueError("[val] needs at least one source, or data.holdout_docs > 0") + raise ValueError("data.micro_batch_tokens must equal data.seq_len with context parallelism") return self @model_validator(mode="after") diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index d2171d5580f..b6cb7accbdd 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -293,13 +293,6 @@ def resolve_freeze_moe_router_auto(self): self.model.freeze_moe_router = False return self - @model_validator(mode="after") - def resolve_random_init_mode_auto(self): - """Resolve ``model.debug.random_init_mode='auto'``: SFT skips weight loading without initializing.""" - if self.model.debug.random_init_mode == "auto": - self.model.debug.random_init_mode = "skip" - return self - @model_validator(mode="after") def auto_setup_run_identity(self): """Auto-generate the run name (``----``) when unset and diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 63b2364e8db..f7035c4e949 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -100,10 +100,7 @@ class DebugModelConfig(BaseConfig): """Override the number of transformer layers (truncates the model).""" random_init: bool = False - """Don't download or load weights from ``model.name``. ``random_init_mode`` sets what the weights hold instead.""" - - random_init_mode: Literal["auto", "skip", "init"] = "auto" - """What ``random_init`` does. ``skip`` leaves the weights uninitialized (arbitrary memory contents) - the fastest option, for throughput benchmarks and debugging. ``init`` initializes the weights as configured in ``model.init``. ``auto`` resolves to ``skip`` for ``sft`` and ``rl`` and to ``init`` for ``pretrain``.""" + """Randomly initialize the model instead of loading weights.""" force_balanced_routing: bool = False """Replace MoE token-choice routing with a round-robin assignment so every expert sees an equal share. Intended for fake-data smoke tests where untrained routing would otherwise OOM under severe imbalance. Gating scores are still gathered from the override indices so the forward pass stays consistent.""" @@ -233,19 +230,6 @@ class MoERuntimeConfig(BaseConfig): dispatch: MoEDispatchConfig = TorchMoEDispatchConfig() -class InitConfig(BaseConfig): - """Random weight initialization, used when ``model.debug.random_init`` resolves to ``init``.""" - - scheme: Literal["normal", "scaled"] = "normal" - """``normal`` draws every weight matrix and the embedding from N(0, std). ``scaled`` additionally divides the std of the residual output projections (``o_proj``, ``down_proj``) by sqrt(2 * num_layers). Norm weights start at 1, biases at 0.""" - - std: float = Field(0.02, gt=0) - """Base standard deviation.""" - - seed: int = 0 - """Seed for the initialization. The values depend only on the seed, not on the number of GPUs.""" - - class ModelConfig(BaseModelConfig): conversion_dir: Path | None = None """Directory for the auto-converted weights (written to a `prime`/`hf` subdirectory). If not set, we write into the model snapshot directory.""" @@ -312,9 +296,6 @@ class ModelConfig(BaseModelConfig): debug: DebugModelConfig = DebugModelConfig() """Debugging knobs for the model and distributed training.""" - init: InitConfig = InitConfig() - """Random weight initialization, used when ``debug.random_init`` resolves to ``init``.""" - fused_lm_head_token_chunk_size: int | Literal["disabled"] = 8192 """Flattened token chunk size for the fused LM head. ``int >= 1`` sets the tokens per LM-head chunk explicitly; ``disabled`` uses the vanilla LM head. In SFT the fused head computes the summed cross-entropy and its gradients chunk by chunk, holding one chunk's full-vocab logits at a time.""" @@ -661,13 +642,6 @@ def resolve_freeze_moe_router_auto(self): self.model.freeze_moe_router = True return self - @model_validator(mode="after") - def resolve_random_init_mode_auto(self): - """Resolve ``model.debug.random_init_mode='auto'``: RL skips weight loading without initializing.""" - if self.model.debug.random_init_mode == "auto": - self.model.debug.random_init_mode = "skip" - return self - @model_validator(mode="after") def deepep_disables_grad_clipping(self): if self.model.ep != 1 and self.model.moe.dispatch.type == "deepep" and self.optim.max_norm is not None: diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py index e94ad716917..7c2c503bf12 100644 --- a/src/prime_rl/entrypoints/pretrain.py +++ b/src/prime_rl/entrypoints/pretrain.py @@ -1,6 +1,7 @@ import json import os import uuid +from pathlib import Path from prime_rl.configs.pretrain import PretrainConfig from prime_rl.utils.config import cli, dump_resolved_config @@ -24,36 +25,31 @@ def write_model_dir(config: PretrainConfig) -> None: - """Write the baseline's config with ``[model.arch]`` applied, plus its tokenizer, to - ``/model``, and point the model and tokenizer at it.""" - from transformers import AutoConfig, AutoTokenizer + """Write the config built from ``[model.arch]`` and the tokenizer to ``/model``. The vocabulary size + (rounded up to a multiple of 128) and the special token ids come from the tokenizer.""" + from transformers import AutoTokenizer + from transformers.models.auto.configuration_auto import CONFIG_MAPPING import prime_rl.trainer.models # noqa: F401 - registers the custom model configs with AutoConfig - model_dir = config.run_dir / "model" - hf_config = AutoConfig.from_pretrained(config.model.name, trust_remote_code=config.model.trust_remote_code) - unknown = [key for key in config.model.arch if not hasattr(hf_config, key)] - if unknown: - raise ValueError(f"[model.arch] fields {unknown} do not exist on the {hf_config.model_type!r} config") - for key, value in config.model.arch.items(): - setattr(hf_config, key, value) - hf_config.save_pretrained(model_dir) + arch = dict(config.model.arch) + config_class = CONFIG_MAPPING[arch.pop("model_type")] + # PretrainedConfig keeps unknown kwargs as attributes, so a typo would silently do nothing. + if unknown := arch.keys() - config_class().to_dict().keys(): + raise ValueError(f"[model.arch] fields {sorted(unknown)} do not exist on {config_class.__name__}") tokenizer = AutoTokenizer.from_pretrained( config.tokenizer.name, trust_remote_code=config.tokenizer.trust_remote_code ) + hf_config = config_class( + **arch, + vocab_size=-(-len(tokenizer) // 128) * 128, + bos_token_id=tokenizer.bos_token_id, + eos_token_id=tokenizer.eos_token_id, + pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id, + ) + model_dir = Path(config.model.name) + hf_config.save_pretrained(model_dir) tokenizer.save_pretrained(model_dir) - config.model.name = config.tokenizer.name = model_dir.as_posix() - - -def download_data(config: PretrainConfig) -> None: - """Download and prepare every dataset once, so the trainer ranks load them from the cache.""" - from prime_rl.trainer.pretrain.data import load_eval_documents, load_text_dataset - - if config.data.type == "text": - load_text_dataset(config.data) - if config.val is not None: - for source in config.val.source: - load_eval_documents(source.dataset, source.subset, source.split, source.revision, source.text_column) def pretrain(config: PretrainConfig): @@ -78,12 +74,7 @@ def pretrain(config: PretrainConfig): write_launch_artifacts(config_dir, "pretrain") if not config.dry_run: - from prime_rl.trainer.model import pre_download_model - - pre_download_model(config.model.name, skip_weights=config.model.debug.random_init) - if config.model.arch: - write_model_dir(config) - download_data(config) + write_model_dir(config) config_path = config_dir / PRETRAIN_CONFIG config_path.write_text(json.dumps(dump_resolved_config(config, exclude={"dry_run", "clean"}), indent=2)) diff --git a/src/prime_rl/trainer/init.py b/src/prime_rl/trainer/init.py index 3d2d49f4c0c..33f1a1c3d1b 100644 --- a/src/prime_rl/trainer/init.py +++ b/src/prime_rl/trainer/init.py @@ -7,7 +7,7 @@ from torch.distributed.tensor import DTensor from torch.distributed.tensor._utils import compute_local_shape_and_global_offset -from prime_rl.configs.trainer import InitConfig +from prime_rl.configs.pretrain import InitConfig from prime_rl.utils.weights import resolve_fqn ParamRole = Literal["embedding", "output", "residual_out", "hidden", "norm", "bias"] diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 821f9373adf..47455fdecf2 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -38,7 +38,6 @@ from prime_rl.multimodal import ForwardPolicy from prime_rl.trainer.activation_checkpointing import get_activation_checkpoint_wrapper from prime_rl.trainer.distributed.embedding_parallel import EmbeddingParallel -from prime_rl.trainer.init import init_weights from prime_rl.trainer.lora import apply_lora_to_model, freeze_all_except_lora_and_specified, strip_lora_from_state_dict from prime_rl.trainer.models import ( AutoModelForCausalLMPrimeRL, @@ -701,13 +700,7 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral logger = get_logger() if config.debug.random_init: - if config.debug.random_init_mode == "init": - logger.info(f"Randomly initializing model weights ({config.init})") - init_weights(model, config.init) - else: - logger.warning( - "Skipping weight loading. Parameters keep arbitrary memory contents (random_init_mode='skip')." - ) + logger.warning("Randomly initializing model. Skipping loading weights from HF.") return if not Path(config.name).exists(): diff --git a/src/prime_rl/trainer/pretrain/evals/common.py b/src/prime_rl/trainer/pretrain/bytes.py similarity index 54% rename from src/prime_rl/trainer/pretrain/evals/common.py rename to src/prime_rl/trainer/pretrain/bytes.py index e6bbeb5a078..98f676695d7 100644 --- a/src/prime_rl/trainer/pretrain/evals/common.py +++ b/src/prime_rl/trainer/pretrain/bytes.py @@ -1,14 +1,6 @@ -from typing import TypeVar - import torch -import torch.distributed as dist from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.trainer.sft.data import Sample -from prime_rl.trainer.world import get_world - -T = TypeVar("T") - def bytes_to_unicode() -> dict[int, str]: """The byte -> character map of byte-level BPE tokenizers (GPT-2, Llama 3, Qwen, DeepSeek).""" @@ -39,27 +31,3 @@ def token_byte_lengths(tokenizer: PreTrainedTokenizer) -> torch.Tensor: else: lengths[token_id] = len(tokenizer.convert_tokens_to_string([token]).encode()) return lengths - - -def local_rows(rows: list[T]) -> list[T | None]: - """This rank's share of ``rows``, padded with ``None`` so every rank runs the same number of - forward passes (each forward is an FSDP collective).""" - world = get_world() - per_rank = -(-len(rows) // world.world_size) - mine: list[T | None] = list(rows[world.rank :: world.world_size]) - return mine + [None] * (per_rank - len(mine)) - - -def to_device(row: Sample) -> dict[str, torch.Tensor]: - return { - "input_ids": torch.tensor(row["input_ids"], device="cuda").unsqueeze(0), - "position_ids": torch.tensor(row["position_ids"], device="cuda").unsqueeze(0), - "target_ids": torch.tensor(row["target_ids"], device="cuda").unsqueeze(0), - "loss_mask": torch.tensor(row["loss_mask"], device="cuda").unsqueeze(0), - "seq_lens": torch.tensor(row["seq_lens"], device="cuda"), - } - - -def all_reduce_sum(tensor: torch.Tensor) -> torch.Tensor: - dist.all_reduce(tensor, op=dist.ReduceOp.SUM) - return tensor diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index ab0d17b3eb1..2ca6e4bb4dc 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -1,183 +1,218 @@ -from collections.abc import Iterable, Iterator -from typing import cast +import bisect +from pathlib import Path -from datasets import Dataset, load_dataset +import numpy as np +import torch from torchdata.stateful_dataloader import StatefulDataLoader -from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.configs.pretrain import TextDataConfig -from prime_rl.trainer.sft.data import Sample, StatefulIterableDataset, cat_collate, setup_and_interleave_datasets +from prime_rl.configs.pretrain import MixturePhaseConfig, PretrainDataConfig, TokenSourceConfig +from prime_rl.trainer.sft.data import StatefulIterableDataset from prime_rl.utils.logger import get_logger -LIST_KEYS = ("input_ids", "target_ids", "loss_mask", "position_ids", "seq_lens") +class TokenShards: + """Datatrove tokenized shards in one folder, read as one token stream in file-name order. + + ``.ds`` holds the token ids (little-endian, 2 or 4 bytes as recorded on the first line of + ``.ds.metadata``) and ``.ds.index`` the uint64 end offset of every document. File + boundaries are document boundaries. The stream wraps around after ``num_tokens``.""" + + def __init__(self, config: TokenSourceConfig): + self.files = sorted(config.path.glob("*.ds")) + if not self.files: + raise FileNotFoundError(f"No datatrove .ds files in {config.path}") + self.dtypes = [token_dtype(file) for file in self.files] + sizes = [file.stat().st_size // dtype.itemsize for file, dtype in zip(self.files, self.dtypes)] + self.starts = np.cumsum([0, *sizes]) + total = int(self.starts[-1]) + if config.max_tokens is not None and config.max_tokens > total: + raise ValueError(f"{config.path} has {total:,} tokens, fewer than max_tokens={config.max_tokens:,}") + self.num_tokens = total if config.max_tokens is None else config.max_tokens + self.memmaps: dict[int, tuple[np.memmap, np.memmap]] = {} + + def open(self, file_index: int) -> tuple[np.memmap, np.memmap]: + # Opened lazily, so every dataloader worker maps the files itself. + if file_index not in self.memmaps: + file = self.files[file_index] + tokens = np.memmap(file, dtype=self.dtypes[file_index], mode="r") + doc_ends = np.memmap(f"{file}.index", dtype=" tuple[np.ndarray, list[int]]: + """Tokens ``[start, start + length)`` of the stream and the offsets in them where a document starts.""" + pieces, doc_starts = [], [] + offset = 0 + while offset < length: + if offset > 0: + doc_starts.append(offset) + position = (start + offset) % self.num_tokens + file_index = int(np.searchsorted(self.starts, position, side="right")) - 1 + local = position - int(self.starts[file_index]) + size = min(length - offset, int(self.starts[file_index + 1]) - position, self.num_tokens - position) + tokens, doc_ends = self.open(file_index) + pieces.append(tokens[local : local + size]) + first, last = np.searchsorted(doc_ends, [local, local + size], side="right") + doc_starts.extend(offset + int(end) - local for end in doc_ends[first:last] if end < local + size) + offset += size + return np.concatenate(pieces).astype(np.int64), doc_starts + + def last_token_of_first_document(self) -> int: + tokens, doc_ends = self.open(0) + return int(tokens[int(doc_ends[0]) - 1]) + + +def token_dtype(file: Path) -> np.dtype: + first_line = Path(f"{file}.metadata").read_text().splitlines()[0] + return np.dtype(f" list[int]: - """Token ids of one document, with the tokenizer's BOS (if it adds one) and a closing EOS.""" - token_ids = tokenizer(text)["input_ids"] - if not token_ids or token_ids[-1] != tokenizer.eos_token_id: - token_ids.append(tokenizer.eos_token_id) - return token_ids + def read(self, start: int, length: int) -> tuple[np.ndarray, list[int]]: + return np.random.default_rng([self.seed, start]).integers(0, self.vocab_size, length), [] -def cut_row(tokens: list[int], doc_starts: list[int], seq_len: int, document_masking: bool) -> Sample: - """One training row from ``seq_len + 1`` tokens: inputs are the first ``seq_len``, targets the - last ``seq_len``. ``doc_starts`` are the offsets in ``tokens`` where a document begins. With - document masking, every document (and the continued document at offset 0) is its own segment - with positions starting at 0.""" - assert len(tokens) == seq_len + 1 - if document_masking: - boundaries = sorted({0, *(start for start in doc_starts if 0 < start < seq_len)}) - seq_lens = [end - start for start, end in zip(boundaries, [*boundaries[1:], seq_len])] - else: - seq_lens = [seq_len] - return { - "input_ids": tokens[:-1], - "target_ids": tokens[1:], - "loss_mask": [True] * seq_len, - "position_ids": [position for length in seq_lens for position in range(length)], - "seq_lens": seq_lens, - "mm_kwargs": None, - "mm_token_type_ids": None, - } +class TokenDataset(StatefulIterableDataset): + """``seq_len`` windows of token sources, each window from one source of a weighted mixture. - -class TextDataset(StatefulIterableDataset): - """Tokenized documents of this data rank, reshuffled every epoch. Documents are addressed by - their global index, so resuming jumps straight to the saved position.""" + Window ``w`` belongs to step ``w // batch_size``; its source is drawn from that step's phase weights + with an RNG keyed by ``(seed, w)``, and it continues the source's stream where that source's last + window ended. Every data rank walks the whole window sequence but reads only its own windows, so the + state (the window counter and the tokens read per source) does not depend on the world size. Each + window reads one token past its end for the last target. With ``num_windows`` the dataset stops + after that many windows.""" def __init__( self, - dataset: Dataset, - tokenizer: PreTrainedTokenizer, - text_column: str = "text", - shuffle: bool = True, + sources: dict[str, TokenShards | RandomTokens], + phases: list[MixturePhaseConfig], + seq_len: int, + batch_size: int, + max_steps: int, seed: int = 0, non_dp_size: int = 1, + num_windows: int | None = None, ): super().__init__(non_dp_size) - self.dataset = dataset - self.num_examples = len(dataset) - self.tokenizer = tokenizer - self.text_column = text_column - self.shuffle = shuffle + self.sources = sources + self.names = list(sources) + self.phase_starts = [round(phase.start * max_steps) for phase in phases] + self.phase_cdfs = [np.cumsum([phase.weights.get(name, 0.0) for name in self.names]) for phase in phases] + self.seq_len = seq_len + self.batch_size = batch_size self.seed = seed + self.num_windows = num_windows + self.tokens = dict.fromkeys(self.names, 0) + + def epochs(self) -> dict[str, float]: + return {name: self.tokens[name] / self.sources[name].num_tokens for name in self.names} - def epoch_dataset(self) -> Dataset: - return self.dataset.shuffle(seed=self.seed + self.epoch) if self.shuffle else self.dataset + def state_dict(self) -> dict: + return {"dataset": {"step": self.step, "tokens": dict(self.tokens), "epochs": self.epochs()}} - def __iter__(self) -> Iterator[list[int]]: + def load_state_dict(self, state_dict: dict): + self.step = state_dict["dataset"]["step"] + self.tokens.update(state_dict["dataset"]["tokens"]) + + def planned_epochs(self, max_steps: int) -> dict[str, float]: + """Expected passes over each source's tokens in ``max_steps`` steps.""" + ends = [*self.phase_starts[1:], max_steps] + step_tokens = self.batch_size * self.seq_len + tokens = sum( + (end - start) * step_tokens * np.diff(cdf, prepend=0.0) + for start, end, cdf in zip(self.phase_starts, ends, self.phase_cdfs) + ) + return {name: float(tokens[i]) / self.sources[name].num_tokens for i, name in enumerate(self.names)} + + def source(self, window: int) -> str: + if len(self.names) == 1: + return self.names[0] + phase = bisect.bisect_right(self.phase_starts, window // self.batch_size) - 1 + cdf = self.phase_cdfs[phase] + draw = np.random.default_rng([self.seed, window]).random() * cdf[-1] + return self.names[min(int(np.searchsorted(cdf, draw, side="right")), len(self.names) - 1)] + + def __iter__(self): self._setup_world_info() - dataset = self.epoch_dataset() - while True: + while self.num_windows is None or self.step < self.num_windows: + window = self.step + name = self.source(window) + start = self.tokens[name] + self.tokens[name] += self.seq_len self.step += 1 - epoch = (self.step - 1) // self.num_examples - if epoch > self.epoch: - self.epoch = epoch - dataset = self.epoch_dataset() - if (self.step - 1) % self.data_world_size != self.data_rank: + if window % self.data_world_size != self.data_rank: continue - example = cast(dict, dataset[(self.step - 1) % self.num_examples]) - token_ids = tokenize_document(self.tokenizer, example[self.text_column]) - source = example.get("__subset") or example.get("__split") - self.num_samples[source] += 1 - self.num_tokens[source] += len(token_ids) - yield token_ids - - -class StreamPacker(StatefulIterableDataset): - """Cuts a document stream into rows of exactly ``seq_len`` tokens (see ``cut_row``) and - yields ``rows_per_sample`` consecutive rows concatenated as one packed sample. Documents - continue across rows, so nothing is truncated or padded, but attention never crosses a - row. The tokens not yet emitted are part of the checkpoint state. Any iterable of - documents works as the source; checkpointing needs a ``TextDataset``.""" - - def __init__( - self, - dataset: TextDataset | Iterable[list[int]], - seq_len: int, - document_masking: bool, - rows_per_sample: int = 1, - ): - self.dataset = dataset - self.seq_len = seq_len - self.document_masking = document_masking - self.rows_per_sample = rows_per_sample - self.buffer: list[int] = [] - self.doc_starts: list[int] = [] + tokens, doc_starts = self.sources[name].read(start, self.seq_len + 1) + yield { + "tokens": tokens, + "seq_lens": np.diff([0, *(s for s in doc_starts if s < self.seq_len), self.seq_len]), + } + + +def collate_rows(rows: list[dict]) -> dict: + """Concatenate windows into one packed sample of shape ``[1, len(rows) * seq_len]``. Positions restart at + every document; attention never crosses a document or a window.""" + tokens = torch.from_numpy(np.stack([row["tokens"] for row in rows])) + seq_lens = torch.from_numpy(np.concatenate([row["seq_lens"] for row in rows])) + return { + "input_ids": tokens[:, :-1].reshape(1, -1), + "target_ids": tokens[:, 1:].reshape(1, -1), + "loss_mask": torch.ones((1, tokens.shape[0] * (tokens.shape[1] - 1)), dtype=torch.bool), + "position_ids": torch.cat([torch.arange(length) for length in seq_lens.tolist()]).unsqueeze(0), + "seq_lens": seq_lens, + "mm_kwargs": None, + "mm_token_type_ids": None, + } - def state_dict(self) -> dict: - return { - "dataset": self.dataset.state_dict(), - "progress": {"num_samples": dict(self.dataset.num_samples), "num_tokens": dict(self.dataset.num_tokens)}, - "buffer": self.buffer, - "doc_starts": self.doc_starts, - } - def load_state_dict(self, state_dict: dict): - self.dataset.load_state_dict(state_dict["dataset"]) - self.dataset.num_samples.update(state_dict["progress"]["num_samples"]) - self.dataset.num_tokens.update(state_dict["progress"]["num_tokens"]) - self.buffer = state_dict["buffer"] - self.doc_starts = state_dict["doc_starts"] - - def next_row(self) -> Sample: - row = cut_row(self.buffer[: self.seq_len + 1], self.doc_starts, self.seq_len, self.document_masking) - self.buffer = self.buffer[self.seq_len :] - self.doc_starts = [start - self.seq_len for start in self.doc_starts if start >= self.seq_len] - return row - - def __iter__(self) -> Iterator[Sample]: - for document in self.dataset: - self.doc_starts.append(len(self.buffer)) - self.buffer.extend(document) - while len(self.buffer) >= self.seq_len * self.rows_per_sample + 1: - rows = [self.next_row() for _ in range(self.rows_per_sample)] - sample = {key: [value for row in rows for value in row[key]] for key in LIST_KEYS} - yield {**sample, "mm_kwargs": None, "mm_token_type_ids": None} - - -def load_text_dataset(config: TextDataConfig) -> Dataset: - """Load (and interleave) the configured subsets/splits. Only their files are downloaded.""" - splits = config.splits or ["train"] * len(config.subsets or [None]) - subsets = config.subsets or [None] * len(splits) - if len(splits) == 1: - dataset = load_dataset(config.name, subsets[0], split=splits[0], revision=config.revision) - return cast(Dataset, dataset).select_columns([config.text_column]) - dataset = setup_and_interleave_datasets( - dataset_name=config.name, - subsets_and_splits=list(zip(subsets, splits)), - probabilities=config.probabilities, - stopping_strategy=config.stopping_strategy, - seed=config.seed, - revision=config.revision, +def setup_dataloader(dataset: TokenDataset, micro_batch_size: int, num_workers: int) -> StatefulDataLoader: + return StatefulDataLoader( + dataset, batch_size=micro_batch_size, collate_fn=collate_rows, num_workers=num_workers, pin_memory=True ) - return dataset.select_columns([config.text_column, "__subset", "__split"]) - - -def split_holdout(dataset: Dataset, holdout_docs: int) -> tuple[Dataset, Dataset | None]: - """Split off the last ``holdout_docs`` documents (in dataset order) as a held-out set.""" - if holdout_docs == 0: - return dataset, None - num_train = len(dataset) - holdout_docs - assert num_train > 0, f"data.holdout_docs ({holdout_docs}) must be smaller than the dataset ({len(dataset)})" - return dataset.select(range(num_train)), dataset.select(range(num_train, len(dataset))) - -def load_eval_documents( - dataset: str, subset: str | None, split: str, revision: str | None, text_column: str -) -> Dataset: - return cast(Dataset, load_dataset(dataset, subset, split=split, revision=revision)).select_columns([text_column]) - -def setup_text_dataloader( - dataset: Dataset, tokenizer: PreTrainedTokenizer, config: TextDataConfig, non_dp_size: int -) -> StatefulDataLoader: - get_logger().info(f"Training on {len(dataset):,} documents from {config.name}") - text_dataset = TextDataset( - dataset, tokenizer, config.text_column, shuffle=config.shuffle, seed=config.seed, non_dp_size=non_dp_size - ) - packer = StreamPacker(text_dataset, config.seq_len, config.document_masking, config.micro_batch_size) - return StatefulDataLoader( - packer, batch_size=1, collate_fn=cat_collate, num_workers=config.num_workers, pin_memory=True +def setup_train_dataset( + config: PretrainDataConfig, max_steps: int, vocab_size: int, eos_token_id: int, non_dp_size: int +) -> TokenDataset: + if config.type == "fake": + sources = {"fake": RandomTokens(vocab_size, config.seed)} + phases = [MixturePhaseConfig(weights={"fake": 1.0})] + else: + sources = {name: TokenShards(source) for name, source in config.sources.items()} + check_eos(sources, eos_token_id) + phases = config.phases + dataset = TokenDataset( + sources, phases, config.seq_len, config.batch_size, max_steps, seed=config.seed, non_dp_size=non_dp_size ) + if config.type == "tokens": + epochs = dataset.planned_epochs(max_steps) + get_logger().info(f"Planned epochs per source: {epochs}") + if over := {name: epoch for name, epoch in epochs.items() if epoch > config.max_epochs}: + raise ValueError(f"Planned epochs {over} exceed data.max_epochs={config.max_epochs}") + return dataset + + +def setup_val_dataset(source: TokenShards, seq_len: int, num_windows: int, non_dp_size: int) -> TokenDataset: + """One pass over the first ``num_windows`` windows of a held-out source.""" + phases = [MixturePhaseConfig(weights={"val": 1.0})] + return TokenDataset({"val": source}, phases, seq_len, 1, 1, non_dp_size=non_dp_size, num_windows=num_windows) + + +def check_eos(sources: dict[str, TokenShards], eos_token_id: int) -> None: + """Every document ends with the tokenizer's EOS. A different token there means the shards were made with + another tokenizer.""" + for name, source in sources.items(): + last = source.last_token_of_first_document() + if last != eos_token_id: + raise ValueError( + f"The first document of source {name!r} ends with token {last}, not the tokenizer's EOS " + f"{eos_token_id}. Were the shards tokenized with tokenizer.name?" + ) diff --git a/src/prime_rl/trainer/pretrain/evals/perplexity.py b/src/prime_rl/trainer/pretrain/evals/perplexity.py deleted file mode 100644 index cded029475c..00000000000 --- a/src/prime_rl/trainer/pretrain/evals/perplexity.py +++ /dev/null @@ -1,61 +0,0 @@ -import math -from dataclasses import dataclass -from itertools import islice - -import torch -from torch import nn -from transformers.tokenization_utils import PreTrainedTokenizer - -from prime_rl.configs.pretrain import ValSourceConfig -from prime_rl.trainer.model import forward -from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX -from prime_rl.trainer.pretrain.data import StreamPacker, load_eval_documents, tokenize_document -from prime_rl.trainer.pretrain.evals.common import all_reduce_sum, local_rows, to_device -from prime_rl.trainer.sft.data import Sample - - -@dataclass -class PerplexitySet: - """A fixed, pre-packed held-out token set.""" - - name: str - rows: list[Sample] - - -def build_perplexity_set( - name: str, documents, tokenizer: PreTrainedTokenizer, seq_len: int, max_tokens: int, text_column: str = "text" -) -> PerplexitySet: - """Pack documents (in order) into full rows until ``max_tokens`` is reached. Rows use - document masking, so every document is scored from its own start.""" - token_stream = (tokenize_document(tokenizer, example[text_column]) for example in documents) - rows = list(islice(StreamPacker(token_stream, seq_len, document_masking=True), max(1, max_tokens // seq_len))) - assert rows, f"Eval source {name!r} has fewer than {seq_len + 1} tokens" - return PerplexitySet(name=name, rows=rows) - - -def load_perplexity_set(source: ValSourceConfig, tokenizer: PreTrainedTokenizer, seq_len: int) -> PerplexitySet: - documents = load_eval_documents(source.dataset, source.subset, source.split, source.revision, source.text_column) - return build_perplexity_set(source.name, documents, tokenizer, seq_len, source.max_tokens, source.text_column) - - -@torch.no_grad() -def evaluate_perplexity(model: nn.Module, sets: list[PerplexitySet], byte_lengths: torch.Tensor) -> dict[str, float]: - """Loss (nats per token), perplexity and bits per byte for every set.""" - byte_lengths = byte_lengths.to("cuda") - metrics = {} - for eval_set in sets: - totals = torch.zeros(3, dtype=torch.float64, device="cuda") # loss, tokens, bytes - for row in local_rows(eval_set.rows): - batch = to_device(row if row is not None else eval_set.rows[0]) - loss_mask = batch["loss_mask"] if row is not None else torch.zeros_like(batch["loss_mask"]) - labels = batch["target_ids"].masked_fill(~loss_mask, IGNORE_INDEX) - out = forward(model, batch["input_ids"], batch["position_ids"], seq_lens=batch["seq_lens"], labels=labels) - totals[0] += out["loss"].double() - totals[1] += loss_mask.sum() - totals[2] += byte_lengths[batch["target_ids"][loss_mask]].sum() - loss_sum, num_tokens, num_bytes = all_reduce_sum(totals).tolist() - loss = loss_sum / num_tokens - metrics[f"val/{eval_set.name}/loss"] = loss - metrics[f"val/{eval_set.name}/ppl"] = math.exp(min(loss, 20)) - metrics[f"val/{eval_set.name}/bpb"] = loss_sum / (math.log(2) * num_bytes) - return metrics diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index a29bf227c4b..2926895237e 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -11,7 +11,6 @@ import torch import torch.distributed as dist from torch import nn -from transformers import AutoConfig from prime_rl import monitors from prime_rl.configs.pretrain import PretrainConfig @@ -32,11 +31,17 @@ from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter -from prime_rl.trainer.pretrain.data import load_text_dataset, setup_text_dataloader, split_holdout -from prime_rl.trainer.pretrain.evals.common import token_byte_lengths -from prime_rl.trainer.pretrain.evals.perplexity import build_perplexity_set, evaluate_perplexity, load_perplexity_set +from prime_rl.trainer.init import init_weights +from prime_rl.trainer.pretrain.bytes import token_byte_lengths +from prime_rl.trainer.pretrain.data import ( + TokenShards, + check_eos, + setup_dataloader, + setup_train_dataset, + setup_val_dataset, +) from prime_rl.trainer.scheduler import setup_scheduler -from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state, setup_dataloader, setup_dataset +from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state from prime_rl.trainer.utils import ( GarbageCollection, begin_backward, @@ -58,18 +63,6 @@ from prime_rl.utils.utils import clean_exit -def check_arch(config: PretrainConfig) -> None: - """``[model.arch]`` is applied by the launcher, which writes the model dir that ``model.name`` - points at. Fail if the trainer was started on a config that doesn't match.""" - hf_config = AutoConfig.from_pretrained(config.model.name, trust_remote_code=config.model.trust_remote_code) - mismatched = {key: value for key, value in config.model.arch.items() if getattr(hf_config, key, None) != value} - if mismatched: - raise ValueError( - f"{config.model.name} does not match [model.arch] {mismatched}. Launch through `uv run pretrain`, " - "which writes the model dir with the overrides applied." - ) - - def count_parameters(model: nn.Module) -> dict[str, int]: """Total, embedding (input embedding and LM head) and non-embedding parameter counts.""" total = embedding = 0 @@ -118,9 +111,6 @@ def train(config: PretrainConfig): ) grad_accum_steps = total_micro_batches // micro_batches_per_step - if config.model.arch: - check_arch(config) - ckpt_manager = setup_ckpt_manager(config.run_dir, config.ckpt, resume=config.resume) checkpoint_step = None if config.resume is not None: @@ -130,7 +120,10 @@ def train(config: PretrainConfig): checkpoint_step = config.resume.step or resolve_latest_ckpt_step(ckpt_manager.ckpt_dir) logger.info(f"Initializing model ({config.model})") - model = setup_model(config.model, parallel_dims, loading_from_checkpoint_later=checkpoint_step is not None) + # Weights come from the checkpoint when resuming and from the random init otherwise. + model = setup_model(config.model, parallel_dims, loading_from_checkpoint_later=True) + if checkpoint_step is None: + init_weights(model, config.model.init) if parallel_dims.cp_enabled: setup_context_parallel(model, config.model, parallel_dims) num_params = count_parameters(model) @@ -154,33 +147,32 @@ def train(config: PretrainConfig): scheduler_steps = ( config.max_steps - checkpoint_step - if config.max_steps is not None and (config.ckpt and config.ckpt.skip_scheduler and checkpoint_step is not None) + if config.ckpt and config.ckpt.skip_scheduler and checkpoint_step is not None else config.max_steps ) scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) - holdout = None - if config.data.type == "text": - documents, holdout = split_holdout(load_text_dataset(config.data), config.data.holdout_docs) - dataloader = setup_text_dataloader(documents, tokenizer, config.data, config.model.cp) - else: - dataloader = setup_dataloader( - setup_dataset(tokenizer, config.data, config.model.cp), config.data, config.model.cp - ) + dataloader = setup_dataloader( + setup_train_dataset( + config.data, config.max_steps, len(tokenizer), tokenizer.eos_token_id, non_dp_size=config.model.cp + ), + config.data.micro_batch_size, + config.data.num_workers, + ) - eval_seq_len = config.data.seq_len - perplexity_sets = [] + val_sources = {} if config.val is not None: - if holdout is not None: - perplexity_sets.append( - build_perplexity_set( - "holdout", holdout, tokenizer, eval_seq_len, max_tokens=10**12, text_column=config.data.text_column - ) - ) - perplexity_sets += [load_perplexity_set(source, tokenizer, eval_seq_len) for source in config.val.source] - for eval_set in perplexity_sets: - logger.info(f"Perplexity eval {eval_set.name!r}: {len(eval_set.rows) * eval_seq_len:,} tokens") - byte_lengths = token_byte_lengths(tokenizer) if perplexity_sets else None + val_sources = {name: TokenShards(source) for name, source in config.val.sources.items()} + check_eos(val_sources, tokenizer.eos_token_id) + byte_lengths = token_byte_lengths(tokenizer).to("cuda") + # Whole micro-batches on every data-parallel rank, as every forward is a collective. + windows_per_round = parallel_dims.get_mesh("dp").size() * config.data.micro_batch_size + val_windows = {} + for name, source in val_sources.items(): + val_windows[name] = (source.num_tokens - 1) // config.data.seq_len // windows_per_round * windows_per_round + if val_windows[name] == 0: + raise ValueError(f"Val source {name!r} has fewer than {windows_per_round} windows of data.seq_len") + logger.info(f"Val source {name!r}: {val_windows[name] * config.data.seq_len:,} tokens") progress = Progress() if checkpoint_step is not None: @@ -231,9 +223,24 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: ) return out["loss"], loss_mask.sum(dtype=torch.int64) + @torch.no_grad() def run_evals(step: int) -> None: - """Perplexity evals on the weights after ``step`` optimizer updates.""" - metrics = evaluate_perplexity(model, perplexity_sets, byte_lengths) + """Held-out loss, perplexity and bits per byte on the weights after ``step`` optimizer updates.""" + metrics = {} + for name, source in val_sources.items(): + val_dataset = setup_val_dataset(source, config.data.seq_len, val_windows[name], config.model.cp) + totals = torch.zeros(3, dtype=torch.float64, device="cuda") # loss sum, tokens, bytes + for micro_batch in setup_dataloader(val_dataset, config.data.micro_batch_size, num_workers=1): + loss_sum, token_count = compute_loss(micro_batch) + # The bytes are counted on the unsharded targets, so every CP rank counts them once. + num_bytes = byte_lengths[micro_batch["target_ids"].to("cuda")].sum() / cp_size + totals += torch.stack([loss_sum.double(), token_count.double(), num_bytes.double()]) + dist.all_reduce(totals, op=dist.ReduceOp.SUM, group=dp_cp_group) + loss_sum, num_tokens, num_bytes = totals.tolist() + loss = loss_sum / num_tokens + metrics[f"val/{name}/loss"] = loss + metrics[f"val/{name}/ppl"] = math.exp(min(loss, 20)) + metrics[f"val/{name}/bpb"] = loss_sum / (math.log(2) * num_bytes) # The no-grad eval forwards leave the [lm_head, norm] FSDP group unsharded, and their # routing must not count toward the next step's MoE statistics. reshard_module(model) @@ -247,12 +254,12 @@ def run_evals(step: int) -> None: gc_handler = GarbageCollection(config.gc.interval) if config.gc else None max_memory = torch.cuda.mem_get_info()[1] / 1024**3 - logger.info(f"Starting training loop (max_steps={config.max_steps or 'infinite'})") + logger.info(f"Starting training loop (max_steps={config.max_steps})") while True: torch.cuda.reset_peak_memory_stats() if gc_handler is not None: gc_handler.run(progress.step) - is_last_step = config.max_steps is not None and progress.step >= config.max_steps + is_last_step = progress.step >= config.max_steps step_start_time = time.perf_counter() step_loss_sum = torch.tensor(0.0, device="cuda") @@ -365,7 +372,8 @@ def run_evals(step: int) -> None: "loss/perplexity": math.exp(min(batch_loss, 20)), "loss/nan_count": nan_loss_count, "optim/lr": current_lr, - "progress/epoch": dataset_progress["epoch"], + **{f"progress/{name}/epochs": epochs for name, epochs in dataset_progress["epochs"].items()}, + **{f"progress/{name}/num_tokens": tokens for name, tokens in dataset_progress["tokens"].items()}, "progress/num_samples": progress.total_samples, "progress/num_tokens": progress.total_tokens, "perf/throughput": throughput, diff --git a/tests/unit/train/pretrain/test_evals.py b/tests/unit/train/pretrain/test_bytes.py similarity index 86% rename from tests/unit/train/pretrain/test_evals.py rename to tests/unit/train/pretrain/test_bytes.py index 35c5dab7ac5..d2403e4d2e7 100644 --- a/tests/unit/train/pretrain/test_evals.py +++ b/tests/unit/train/pretrain/test_bytes.py @@ -1,6 +1,6 @@ from transformers import AutoTokenizer -from prime_rl.trainer.pretrain.evals.common import token_byte_lengths +from prime_rl.trainer.pretrain.bytes import token_byte_lengths def test_token_byte_lengths_count_utf8_bytes(): diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index 8f3f6ab7628..c74006b90eb 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -1,58 +1,66 @@ -from datasets import Dataset -from transformers import AutoTokenizer - -from prime_rl.trainer.pretrain.data import StreamPacker, TextDataset, cut_row, tokenize_document - - -def test_cut_row_segments_documents(): - tokens = list(range(9)) - row = cut_row(tokens, doc_starts=[-2, 3, 5], seq_len=8, document_masking=True) - assert row["input_ids"] == tokens[:8] - assert row["target_ids"] == tokens[1:] - assert row["seq_lens"] == [3, 2, 3] - assert row["position_ids"] == [0, 1, 2, 0, 1, 0, 1, 2] - unmasked = cut_row(tokens, doc_starts=[3, 5], seq_len=8, document_masking=False) - assert unmasked["seq_lens"] == [8] - assert unmasked["position_ids"] == list(range(8)) - - -def test_stream_packer_keeps_every_token(): - documents = [list(range(start, start + length)) for start, length in [(0, 5), (100, 9), (200, 4), (300, 30)]] - rows = list(StreamPacker(documents, seq_len=8, document_masking=True)) - stream = [token for document in documents for token in document] - assert len(rows) == (len(stream) - 1) // 8 - assert [token for row in rows for token in row["input_ids"]] == stream[: len(rows) * 8] - assert all(row["target_ids"][-1] == next_row["input_ids"][0] for row, next_row in zip(rows, rows[1:])) - assert rows[0]["seq_lens"] == [5, 3] - assert rows[1]["seq_lens"] == [6, 2] - - samples = list(StreamPacker(documents, seq_len=8, document_masking=False, rows_per_sample=2)) - assert [sample["seq_lens"] for sample in samples] == [[8, 8]] * len(samples) - assert samples[0]["input_ids"] == rows[0]["input_ids"] + rows[1]["input_ids"] - - -def test_stream_packer_resumes_mid_document(): - tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") - texts = [f"Document {i}: " + "word " * (7 * i + 3) for i in range(40)] - dataset = Dataset.from_list([{"text": text} for text in texts]) - - def packer(): - return StreamPacker(TextDataset(dataset, tokenizer, shuffle=True, seed=1), seq_len=16, document_masking=True) - - reference = packer() - reference_rows = iter(reference) - for _ in range(7): - next(reference_rows) - state = reference.state_dict() - expected = [next(reference_rows) for _ in range(20)] - - resumed = packer() - resumed.load_state_dict(state) - assert [row for row, _ in zip(resumed, range(20))] == expected - - -def test_tokenize_document_appends_eos(): - tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") - token_ids = tokenize_document(tokenizer, "Hello world") - assert token_ids[-1] == tokenizer.eos_token_id - assert token_ids.count(tokenizer.eos_token_id) == 1 +from pathlib import Path + +import numpy as np + +from prime_rl.configs.pretrain import MixturePhaseConfig, TokenSourceConfig +from prime_rl.trainer.pretrain.data import TokenDataset, TokenShards, collate_rows + + +def write_shard(path: Path, documents: list[list[int]]) -> None: + """A datatrove DocumentTokenizer shard: uint32 tokens, uint64 cumulative document ends, metadata.""" + path.parent.mkdir(parents=True, exist_ok=True) + np.array([token for document in documents for token in document], dtype=" PretrainConfig: + return PretrainConfig( + model={"arch": {"model_type": "glm4_moe", "hidden_size": 64, "num_hidden_layers": 2, **arch}}, + tokenizer={"name": TOKENIZER}, + data={"type": "fake", "seq_len": 8, "batch_tokens": 8, "micro_batch_tokens": 8}, + optim={"type": "muon", "lr": 1e-3}, + max_steps=1, + output_dir=tmp_path, + ) + + +def test_model_dir_takes_vocabulary_and_special_tokens_from_tokenizer(tmp_path): + config = make_config(tmp_path) + write_model_dir(config) + hf_config = AutoConfig.from_pretrained(config.model.name) + tokenizer = AutoTokenizer.from_pretrained(TOKENIZER) + assert hf_config.model_type == "glm4_moe" + assert hf_config.hidden_size == 64 + assert hf_config.vocab_size == -(-len(tokenizer) // 128) * 128 + assert hf_config.eos_token_id == tokenizer.eos_token_id + + +def test_model_dir_rejects_unknown_arch_fields(tmp_path): + with pytest.raises(ValueError, match="hiden_size"): + write_model_dir(make_config(tmp_path, hiden_size=64)) From cce76bebed1fa218daa899243cb8684de16dcd86 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 16:32:45 +0200 Subject: [PATCH 016/132] fix(pretrain): accumulate step stats as tensors --- src/prime_rl/trainer/pretrain/train.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 083b7f80ab3..d921a2b4503 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -314,18 +314,18 @@ def run_evals(step: int) -> None: if is_moe_model: counts, confidence = take_routing_counts(model) - sums["expert_counts"] = sums.get("expert_counts", 0) + counts - sums["routing_confidence"] = sums.get("routing_confidence", 0) + confidence + sums["expert_counts"] = sums.get("expert_counts", zero) + counts + sums["routing_confidence"] = sums.get("routing_confidence", zero) + confidence # One rank's micro-batch: the imbalance that expert parallelism and the dropless dispatch see. micro_max_mean = counts.amax(dim=1) / counts.mean(dim=1) - sums["micro_max_mean"] = sums.get("micro_max_mean", 0) + micro_max_mean.mean() + sums["micro_max_mean"] = sums.get("micro_max_mean", zero) + micro_max_mean.mean() maxes["micro_max_mean_max"] = torch.maximum( maxes.get("micro_max_mean_max", zero), micro_max_mean.amax() ) if seq_balances: rows = torch.cat(seq_balances).detach() - sums["seq_balance"] = sums.get("seq_balance", 0) + rows.sum() - sums["seq_balance_rows"] = sums.get("seq_balance_rows", 0) + rows.numel() + sums["seq_balance"] = sums.get("seq_balance", zero) + rows.sum() + sums["seq_balance_rows"] = sums.get("seq_balance_rows", zero) + rows.numel() sums["data_time"] = zero + data_time maxes["data_time_max"] = zero + data_time From 379971f9cbd1c7c980f5a201c2a22f881cf06ee7 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 16:59:20 +0200 Subject: [PATCH 017/132] refactor(pretrain): one WSD learning-rate schedule [schedules] (per-hyperparameter wsd/cosine/constant/custom schedules and a clip-norm schedule) becomes [schedule]: one WSD multiplier on every param group's learning rate (warmup, decay, end_ratio, decay_shape). Cosine decay is a decay_shape; weight decay already follows the learning rate. The schedulable-field map on the shared optimizer configs is gone, so the shared configs are unchanged. --- docs/pretraining.md | 27 +--- .../src/prime_rl/configs/pretrain.py | 99 ++----------- .../src/prime_rl/configs/trainer.py | 34 +---- src/prime_rl/trainer/pretrain/schedules.py | 139 +++++------------- src/prime_rl/trainer/pretrain/train.py | 22 +-- tests/unit/train/pretrain/test_schedules.py | 99 +++++-------- 6 files changed, 94 insertions(+), 326 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 57046a1de50..358f39c4cad 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -34,36 +34,19 @@ Weights start from a random init, unless the run resumes from a checkpoint. | `model.init.std` | Base standard deviation (default 0.02). | | `model.init.seed` | Seed. Every rank initializes only its own shard, and the values depend only on the seed and the parameter name, not on the number of GPUs. | -## Schedules +## Learning-rate schedule -A schedule is a multiplier `f(step)`, where `step` counts the optimizer updates completed so far (0 for the first update). Update `step` uses `base * f(step)` for every param group, with `base` the group's own configured value, so ratios between param groups hold at every step. Each optimizer hyperparameter gets its own schedule under `[schedules.optim.]`, keyed by its `optim` field name: `lr`, `weight_decay`, `betas1`, `betas2`, and `mu` for Muon. `[schedules.clip_norm]` scales `optim.max_norm`. Hyperparameters without a schedule stay constant. The default is a WSD schedule on `lr`; setting `schedules.optim` replaces it, so include `lr` when scheduling other hyperparameters. +`[schedule]` sets a multiplier `f(step)` on the learning rate, where `step` counts the optimizer updates completed so far (0 for the first update). Every param group uses its own learning rate times `f(step)`, so ratios between groups hold at every step. The default warms up over 10% of the steps and decays linearly to 0 over the remaining 90%. ```toml -[schedules.optim.lr] -type = "wsd" +[schedule] warmup = 0.05 # fraction of max_steps, or a number of steps decay = 0.2 end_ratio = 0.1 -decay_shape = "1-sqrt" - -[schedules.optim.betas2] -type = "cosine" -warmup = 0 -end_ratio = 1.05 +decay_shape = "1-sqrt" # or "linear" (default), "cosine" ``` -| `type` | Shape | -|---|---| -| `wsd` (default) | Linear warmup, stable at 1, then a `linear`, `cosine` or `1-sqrt` decay over the last `decay` updates. Defaults: `warmup = 0.1`, `decay = 0.9`, i.e. a linear decay to 0 over all updates after the warmup. | -| `cosine` | Linear warmup, then cosine decay over all remaining updates. | -| `constant` | 1 at every step. | -| `custom` | `fn(step, total_steps=max_steps, **kwargs) -> float` imported from `import_path`, with `kwargs` from the config. | - -The configured endpoints are exactly what the first and last updates use: `f(0) = start_ratio`, the warmup reaches 1 at its last update (`warmup - 1`; a warmup of 0 or 1 starts at 1), and the decay reaches `end_ratio` at the last update (`max_steps - 1`). Both ratios default to 0. Fractions resolve against `max_steps`, so fractional phases and decays need `max_steps`. - -AdamW and Muon apply weight decay as `lr * weight_decay`, so a learning rate schedule already scales the decay. - -The values are recomputed from the step, so checkpoints hold no schedule state. Resuming with `ckpt.skip_scheduler = true` restarts every schedule at the resumed step, spanning the remaining steps. Each scheduled value is logged as `optim/` (the first param group's value) and its multiplier as `schedule/`. +`f(0) = 0`, the warmup reaches 1 at its last update (`warmup - 1`; a warmup of 0 or 1 starts at 1), and the decay reaches `end_ratio` at the last update (`max_steps - 1`). AdamW and Muon apply weight decay as `lr * weight_decay`, so the schedule also scales the decay. The multiplier is recomputed from the step, so checkpoints hold no schedule state. Resuming with `ckpt.skip_scheduler = true` restarts the schedule at the resumed step, spanning the remaining steps. The multiplier is logged as `schedule/lr`. ## Data diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 64624ea4b04..de4a6092981 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -157,17 +157,16 @@ class ValConfig(BaseConfig): StepsOrFraction: TypeAlias = Annotated[int, Field(ge=0)] | Annotated[float, Field(gt=0, lt=1)] -def resolve_steps(value: int | float, total_steps: int | None) -> int: +def resolve_steps(value: int | float, total_steps: int) -> int: """Steps as given, or a fraction of ``total_steps``.""" return round(value * total_steps) if isinstance(value, float) else value -class WSDScheduleConfig(BaseConfig): - """Warmup, stable at 1, decay. ``f(0) = start_ratio``, the warmup reaches 1 at its last update - (``warmup - 1``; a warmup of 0 or 1 starts at 1), and the decay reaches ``end_ratio`` at the last - update of the run.""" - - type: Literal["wsd"] = "wsd" +class ScheduleConfig(BaseConfig): + """Learning-rate multiplier ``f(step)`` (warmup, stable at 1, decay), where ``step`` counts the updates + completed before the current one. Every param group uses its own learning rate times ``f(step)``, so + ratios between groups hold. ``f(0) = 0``, the warmup reaches 1 at its last update (``warmup - 1``; a + warmup of 0 or 1 starts at 1), and the decay reaches ``end_ratio`` at the last update of the run.""" warmup: StepsOrFraction = 0.1 """Warmup updates, or a fraction of ``max_steps``.""" @@ -176,9 +175,6 @@ class WSDScheduleConfig(BaseConfig): """Decay updates at the end of the run, or a fraction of ``max_steps``. The default decays over all updates after the default warmup.""" - start_ratio: float = Field(0.0, ge=0) - """Multiplier of the first update.""" - end_ratio: float = Field(0.0, ge=0) """Multiplier of the last update.""" @@ -186,54 +182,6 @@ class WSDScheduleConfig(BaseConfig): """Shape of the decay from 1 to ``end_ratio``.""" -class CosineScheduleConfig(BaseConfig): - """Linear warmup, then cosine decay from 1 to ``end_ratio`` over the remaining updates.""" - - type: Literal["cosine"] = "cosine" - - warmup: StepsOrFraction = 0.1 - """Warmup updates, or a fraction of ``max_steps``.""" - - start_ratio: float = Field(0.0, ge=0) - """Multiplier of the first update.""" - - end_ratio: float = Field(0.0, ge=0) - """Multiplier of the last update.""" - - -class ConstantScheduleConfig(BaseConfig): - """Multiplier 1 at every step.""" - - type: Literal["constant"] = "constant" - - -class CustomScheduleConfig(BaseConfig): - type: Literal["custom"] = "custom" - - import_path: str - """Import path to the schedule function, called as ``fn(step, total_steps=max_steps, **kwargs)`` and returning the multiplier.""" - - kwargs: dict[str, Any] = Field(default_factory=dict) - """Kwargs forwarded to the schedule function.""" - - -ScheduleConfig: TypeAlias = Annotated[ - WSDScheduleConfig | CosineScheduleConfig | ConstantScheduleConfig | CustomScheduleConfig, - Field(discriminator="type"), -] - - -class SchedulesConfig(BaseConfig): - """Each schedule is a multiplier ``f(step)``: update ``step`` (counted from 0) uses ``base * f(step)``, - where ``base`` is the configured value of each param group.""" - - optim: dict[str, ScheduleConfig] = Field(default_factory=lambda: {"lr": WSDScheduleConfig()}) - """Schedules for optimizer hyperparameters, keyed by their ``optim`` field name (``lr``, ``weight_decay``, ``betas1``, ``betas2``, and ``mu`` for Muon). Unscheduled ones stay constant.""" - - clip_norm: ScheduleConfig | None = None - """Schedule for the gradient clipping norm ``optim.max_norm``. If None, it stays constant.""" - - class PretrainDeploymentConfig(BaseConfig): num_train_gpus: int = Field(1, ge=1) """GPUs for the trainer, on this node.""" @@ -256,8 +204,8 @@ class PretrainConfig(BaseConfig): optim: OptimizerConfig = MuonConfig() """Optimizer. ``optim.lr`` is required.""" - schedules: SchedulesConfig = SchedulesConfig() - """Hyperparameter schedules. Defaults to a WSD learning rate schedule.""" + schedule: ScheduleConfig = ScheduleConfig() + """Learning-rate schedule. Defaults to a warmup over 10% of the steps and a linear decay to 0 over the rest.""" ckpt: CheckpointConfig | None = CheckpointConfig() """Checkpointing. Defaults to a checkpoint at the end of training.""" @@ -361,29 +309,10 @@ def validate_cp(self): return self @model_validator(mode="after") - def validate_schedules(self): - unsupported = set(self.schedules.optim) - set(self.optim.schedulable) - if unsupported: - raise ValueError( - f"schedules.optim: {sorted(unsupported)} cannot be scheduled with optim.type = {self.optim.type!r} " - f"(schedulable: {sorted(self.optim.schedulable)})" - ) - if self.schedules.clip_norm is not None and self.optim.max_norm is None: - raise ValueError("schedules.clip_norm needs optim.max_norm") - named = {f"schedules.optim.{name}": schedule for name, schedule in self.schedules.optim.items()} - named["schedules.clip_norm"] = self.schedules.clip_norm - for name, schedule in named.items(): - if isinstance(schedule, WSDScheduleConfig): - phases = [schedule.warmup, schedule.decay] - needs_max_steps = schedule.decay != 0 or isinstance(schedule.warmup, float) - elif isinstance(schedule, CosineScheduleConfig): - phases = [schedule.warmup] - needs_max_steps = True - else: - continue - if self.max_steps is None: - if needs_max_steps: - raise ValueError(f"{name} needs max_steps (it decays or has fractional phases)") - elif sum(resolve_steps(phase, self.max_steps) for phase in phases) > self.max_steps: - raise ValueError(f"{name}: its phases exceed max_steps ({self.max_steps})") + def validate_schedule(self): + phases = resolve_steps(self.schedule.warmup, self.max_steps) + resolve_steps( + self.schedule.decay, self.max_steps + ) + if phases > self.max_steps: + raise ValueError(f"schedule.warmup + schedule.decay ({phases}) exceed max_steps ({self.max_steps})") return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index f438b5eedab..f7035c4e949 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -1,7 +1,7 @@ import re import warnings from pathlib import Path -from typing import Annotated, Any, ClassVar, Literal, NamedTuple, TypeAlias +from typing import Annotated, Any, Literal, TypeAlias from pydantic import Field, field_validator, model_validator @@ -399,26 +399,7 @@ def validate_scheduler(scheduler: SchedulerConfig, max_steps: int | None) -> Non raise ValueError(f"warmup_steps ({scheduler.warmup_steps}) must be less than max_steps ({max_steps})") -class ParamGroupKey(NamedTuple): - """Where an optimizer config field lives in the optimizer's param groups.""" - - key: str - """Param group key.""" - - index: int | None = None - """Index into the param group value when it is a tuple (e.g. torch AdamW ``betas``).""" - - algorithm: str | None = None - """Only param groups with this ``algorithm`` (Dion Muon). None applies to every group.""" - - class BaseOptimizerConfig(BaseConfig): - schedulable: ClassVar[dict[str, ParamGroupKey]] = { - "lr": ParamGroupKey("lr"), - "weight_decay": ParamGroupKey("weight_decay"), - } - """Hyperparameters a schedule can scale, by config field name.""" - lr: float = Field(1e-6, ge=0) """Peak learning rate.""" @@ -430,12 +411,6 @@ class BaseOptimizerConfig(BaseConfig): class AdamWConfig(BaseOptimizerConfig): - schedulable: ClassVar[dict[str, ParamGroupKey]] = { - **BaseOptimizerConfig.schedulable, - "betas1": ParamGroupKey("betas", index=0), - "betas2": ParamGroupKey("betas", index=1), - } - type: Literal["adamw"] = "adamw" betas1: float = Field(0.9, ge=0) @@ -446,13 +421,6 @@ class AdamWConfig(BaseOptimizerConfig): class MuonConfig(BaseOptimizerConfig): - schedulable: ClassVar[dict[str, ParamGroupKey]] = { - **BaseOptimizerConfig.schedulable, - "mu": ParamGroupKey("mu", algorithm="muon"), - "betas1": ParamGroupKey("beta1", algorithm="adamw"), - "betas2": ParamGroupKey("beta2", algorithm="adamw"), - } - type: Literal["muon"] = "muon" mu: float = Field(0.95, ge=0) diff --git a/src/prime_rl/trainer/pretrain/schedules.py b/src/prime_rl/trainer/pretrain/schedules.py index 462b7ac2fb3..2f8f0cba068 100644 --- a/src/prime_rl/trainer/pretrain/schedules.py +++ b/src/prime_rl/trainer/pretrain/schedules.py @@ -1,14 +1,8 @@ import math -from functools import partial from typing import Callable -from prime_rl.configs.pretrain import ScheduleConfig, SchedulesConfig, resolve_steps -from prime_rl.configs.trainer import OptimizerConfig, ParamGroupKey +from prime_rl.configs.pretrain import ScheduleConfig, resolve_steps from prime_rl.trainer.optim import OptimizerLike -from prime_rl.utils.utils import import_object - -Schedule = Callable[[int], float] -"""Multiplier for an update, given the number of updates completed before it (0 for the first).""" DECAY_SHAPES: dict[str, Callable[[float], float]] = { "linear": lambda progress: 1 - progress, @@ -17,20 +11,12 @@ } -def wsd( - step: int, - total_steps: int | None, - warmup: int, - decay: int, - start_ratio: float, - end_ratio: float, - decay_shape: str, -) -> float: - """Linear warmup from ``start_ratio`` at step 0 to 1 at step ``warmup - 1``, stable at 1, then - ``decay`` updates that reach ``end_ratio`` at the last step ``total_steps - 1``.""" +def wsd(step: int, total_steps: int, warmup: int, decay: int, end_ratio: float, decay_shape: str) -> float: + """Linear warmup from 0 at step 0 to 1 at step ``warmup - 1``, stable at 1, then ``decay`` updates + that reach ``end_ratio`` at the last step ``total_steps - 1``.""" peak = max(warmup - 1, 0) if step < peak: - return start_ratio + (1 - start_ratio) * step / peak + return step / peak if decay == 0: return 1.0 last_stable = max(total_steps - 1 - decay, peak) @@ -40,95 +26,36 @@ def wsd( return end_ratio + (1 - end_ratio) * DECAY_SHAPES[decay_shape](progress) -def cosine(step: int, total_steps: int, warmup: int, start_ratio: float, end_ratio: float) -> float: - """Linear warmup as in ``wsd``, then cosine decay from 1 to ``end_ratio`` at the last step.""" - decay = total_steps - max(warmup, 1) - return wsd(step, total_steps, warmup, decay, start_ratio, end_ratio, "cosine") - - -def constant(step: int, total_steps: int | None) -> float: - return 1.0 - +class LRSchedule: + """Sets the learning rate of each update. It is a pure function of the step, so nothing needs to be + checkpointed: on resume it is recomputed from the step. -def build_schedule(config: ScheduleConfig, total_steps: int | None) -> Schedule: - """Bind a schedule config to the number of updates it spans. Fractional phases resolve against it.""" - match config.type: - case "wsd": - return partial( - wsd, - total_steps=total_steps, - warmup=resolve_steps(config.warmup, total_steps), - decay=resolve_steps(config.decay, total_steps), - start_ratio=config.start_ratio, - end_ratio=config.end_ratio, - decay_shape=config.decay_shape, - ) - case "cosine": - return partial( - cosine, - total_steps=total_steps, - warmup=resolve_steps(config.warmup, total_steps), - start_ratio=config.start_ratio, - end_ratio=config.end_ratio, - ) - case "constant": - return partial(constant, total_steps=total_steps) - case "custom": - return partial(import_object(config.import_path), total_steps=total_steps, **config.kwargs) - - -class Schedules: - """Sets the scheduled hyperparameters of each update. The values are a pure function of the step, - so nothing needs to be checkpointed: on resume they are recomputed from the step. - - The base values are read from the optimizer's param groups at construction, before any checkpoint - load. Each param group keeps its own base, so ratios between groups hold at every step. + The base learning rates are read from the optimizer's param groups at construction, before any + checkpoint load. Each param group keeps its own base, so ratios between groups hold at every step. """ - def __init__( - self, - config: SchedulesConfig, - optim_config: OptimizerConfig, - optimizer: OptimizerLike, - total_steps: int | None, - step_offset: int = 0, - ): + def __init__(self, config: ScheduleConfig, optimizer: OptimizerLike, total_steps: int, step_offset: int = 0): + self.config = config self.optimizer = optimizer + self.total_steps = total_steps self.step_offset = step_offset - # Group indices, not group dicts: loading an optimizer state dict replaces the group dicts. - self.optim: dict[str, tuple[Schedule, ParamGroupKey, list[tuple[int, float]]]] = {} - for name, schedule in config.optim.items(): - target = optim_config.schedulable[name] - bases = [ - (i, group[target.key] if target.index is None else group[target.key][target.index]) - for i, group in enumerate(optimizer.param_groups) - if target.algorithm is None or group["algorithm"] == target.algorithm - ] - self.optim[name] = (build_schedule(schedule, total_steps), target, bases) - self.max_norm = optim_config.max_norm - self.clip_norm = build_schedule(config.clip_norm, total_steps) if config.clip_norm is not None else None - - def apply(self, step: int) -> tuple[float | None, dict[str, float]]: - """Write the values of update ``step`` (updates completed so far) into the optimizer. - Returns the gradient clipping norm and metrics of every scheduled value.""" - step -= self.step_offset - metrics = {} - for name, (schedule, target, bases) in self.optim.items(): - multiplier = schedule(step) - for i, base in bases: - group = self.optimizer.param_groups[i] - if target.index is None: - group[target.key] = base * multiplier - else: - values = list(group[target.key]) - values[target.index] = base * multiplier - group[target.key] = tuple(values) - metrics[f"schedule/{name}"] = multiplier - metrics[f"optim/{name}"] = bases[0][1] * multiplier - max_norm = self.max_norm - if self.clip_norm is not None and max_norm is not None: - multiplier = self.clip_norm(step) - max_norm *= multiplier - metrics["schedule/clip_norm"] = multiplier - metrics["optim/max_norm"] = max_norm - return max_norm, metrics + # Bases by group index: loading an optimizer state dict replaces the group dicts. + self.bases = [group["lr"] for group in optimizer.param_groups] + + def multiplier(self, step: int) -> float: + config = self.config + return wsd( + step - self.step_offset, + self.total_steps, + resolve_steps(config.warmup, self.total_steps), + resolve_steps(config.decay, self.total_steps), + config.end_ratio, + config.decay_shape, + ) + + def apply(self, step: int) -> float: + """Write the learning rates of update ``step`` (updates completed so far) and return the multiplier.""" + multiplier = self.multiplier(step) + for group, base in zip(self.optimizer.param_groups, self.bases): + group["lr"] = base * multiplier + return multiplier diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index d4aaaa1eb63..b6feedb89a9 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -40,7 +40,7 @@ setup_train_dataset, setup_val_dataset, ) -from prime_rl.trainer.pretrain.schedules import Schedules +from prime_rl.trainer.pretrain.schedules import LRSchedule from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state from prime_rl.trainer.utils import ( GarbageCollection, @@ -145,16 +145,10 @@ def train(config: PretrainConfig): ), ) - # With ckpt.skip_scheduler, the schedules restart at the resumed step and span the remaining steps. - restart_schedules = config.ckpt is not None and config.ckpt.skip_scheduler and checkpoint_step is not None - schedule_offset = checkpoint_step if restart_schedules else 0 - schedules = Schedules( - config.schedules, - config.optim, - optimizer, - total_steps=config.max_steps - schedule_offset if config.max_steps is not None else None, - step_offset=schedule_offset, - ) + # With ckpt.skip_scheduler, the schedule restarts at the resumed step and spans the remaining steps. + restart_schedule = config.ckpt is not None and config.ckpt.skip_scheduler and checkpoint_step is not None + schedule_offset = checkpoint_step if restart_schedule else 0 + lr_schedule = LRSchedule(config.schedule, optimizer, config.max_steps - schedule_offset, schedule_offset) dataloader = setup_dataloader( setup_train_dataset( @@ -263,7 +257,7 @@ def run_evals(step: int) -> None: gc_handler.run(progress.step) is_last_step = progress.step >= config.max_steps step_start_time = time.perf_counter() - max_norm, schedule_metrics = schedules.apply(progress.step - 1) + lr_multiplier = lr_schedule.apply(progress.step - 1) step_loss_sum = torch.tensor(0.0, device="cuda") nan_loss_count = torch.tensor(0, device="cuda") @@ -329,7 +323,7 @@ def run_evals(step: int) -> None: batch_loss = (step_loss_sum / global_token_count_val).item() nan_loss_count = nan_loss_count.item() - grad_norm = clip_grad_norm_(gradient_manager, model, max_norm) if max_norm else None + grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm) if config.optim.max_norm else None optimizer.step() optimizer.zero_grad() current_lr = optimizer.param_groups[0]["lr"] @@ -389,7 +383,7 @@ def run_evals(step: int) -> None: **get_ckpt_disk_metrics(config.run_dir), **{name: value.item() for name, value in moe_stats.items()}, **expert_load_stats, - **schedule_metrics, + "schedule/lr": lr_multiplier, "step": progress.step, } if grad_norm is not None: diff --git a/tests/unit/train/pretrain/test_schedules.py b/tests/unit/train/pretrain/test_schedules.py index d8483e61e79..fa6dc486e4a 100644 --- a/tests/unit/train/pretrain/test_schedules.py +++ b/tests/unit/train/pretrain/test_schedules.py @@ -4,84 +4,51 @@ import torch from pydantic import ValidationError -from prime_rl.configs.pretrain import ( - CosineScheduleConfig, - CustomScheduleConfig, - PretrainConfig, - SchedulesConfig, - WSDScheduleConfig, -) -from prime_rl.configs.trainer import AdamWConfig -from prime_rl.trainer.pretrain.schedules import Schedules, build_schedule +from prime_rl.configs.pretrain import PretrainConfig, ScheduleConfig +from prime_rl.trainer.pretrain.schedules import LRSchedule -def test_wsd_endpoints_and_shapes(): - config = WSDScheduleConfig(warmup=0.1, decay=0.2, start_ratio=0.1, end_ratio=0.05) - f = build_schedule(config, total_steps=100) - assert f(0) == pytest.approx(0.1) - assert f(9) == 1.0 - assert f(79) == 1.0 - assert f(89) == pytest.approx(0.05 + 0.95 * 0.5) - assert f(99) == pytest.approx(0.05) - - for shape, midpoint in [("cosine", 0.5), ("1-sqrt", 1 - math.sqrt(0.5))]: - f = build_schedule(WSDScheduleConfig(warmup=0, decay=20, decay_shape=shape), total_steps=100) - assert f(0) == 1.0 - assert f(89) == pytest.approx(midpoint) - assert f(99) == 0.0 - - -def test_cosine_endpoints(): - f = build_schedule(CosineScheduleConfig(warmup=10, end_ratio=0.1), total_steps=100) - assert f(0) == 0.0 - assert f(9) == 1.0 - assert f(54) == pytest.approx(0.55) - assert f(99) == pytest.approx(0.1) - assert build_schedule(CosineScheduleConfig(warmup=0), total_steps=10)(0) == 1.0 - - -def inverse_sqrt(step: int, total_steps: int, scale: float) -> float: - return scale / math.sqrt(step + 1) +def multipliers(config: ScheduleConfig, total_steps: int) -> list[float]: + params = [torch.nn.Parameter(torch.zeros(1))] + schedule = LRSchedule(config, torch.optim.AdamW(params, lr=1.0), total_steps) + return [schedule.multiplier(step) for step in range(total_steps)] -def test_custom_schedule(): - config = CustomScheduleConfig(import_path=f"{__name__}.inverse_sqrt", kwargs={"scale": 2.0}) - assert build_schedule(config, total_steps=10)(3) == 1.0 - +def test_wsd_endpoints_and_shapes(): + f = multipliers(ScheduleConfig(warmup=0.1, decay=0.2, end_ratio=0.05), total_steps=100) + assert f[0] == 0.0 + assert f[9] == 1.0 + assert f[79] == 1.0 + assert f[89] == pytest.approx(0.05 + 0.95 * 0.5) + assert f[99] == pytest.approx(0.05) -def test_schedules_config_validation(): - PretrainConfig(max_steps=100, optim={"type": "muon"}, schedules={"optim": {"mu": {"type": "constant"}}}) - for schedules in [{"optim": {"momentum": {"type": "constant"}}}, {"optim": {"mu": {"type": "constant"}}}]: - with pytest.raises(ValidationError, match="cannot be scheduled"): - PretrainConfig(max_steps=100, schedules=schedules) - with pytest.raises(ValidationError, match="needs max_steps"): - PretrainConfig() + for shape, midpoint in [("cosine", 0.5), ("1-sqrt", 1 - math.sqrt(0.5))]: + f = multipliers(ScheduleConfig(warmup=0, decay=20, decay_shape=shape), total_steps=100) + assert f[0] == 1.0 + assert f[89] == pytest.approx(midpoint) + assert f[99] == 0.0 + + +def test_schedule_phases_must_fit_max_steps(): + config = PretrainConfig( + model={"arch": {"model_type": "llama"}}, + tokenizer={"name": "t"}, + data={"type": "fake", "seq_len": 8, "batch_tokens": 8, "micro_batch_tokens": 8}, + optim={"lr": 1e-3}, + max_steps=100, + ) with pytest.raises(ValidationError, match="exceed max_steps"): - PretrainConfig(max_steps=100, schedules={"optim": {"lr": {"type": "wsd", "warmup": 50, "decay": 51}}}) + PretrainConfig(**{**config.model_dump(), "schedule": {"warmup": 50, "decay": 51}}) def test_apply_keeps_group_ratios(): params = [torch.nn.Parameter(torch.zeros(2)), torch.nn.Parameter(torch.zeros(2))] - optimizer = torch.optim.AdamW( - [{"params": params[:1], "lr": 1.0}, {"params": params[1:], "lr": 0.1}], betas=(0.9, 0.95) - ) - config = SchedulesConfig( - optim={ - "lr": WSDScheduleConfig(warmup=0, decay=2, end_ratio=0.5), - "betas2": WSDScheduleConfig(warmup=0, decay=1, end_ratio=0.5), - }, - clip_norm=WSDScheduleConfig(warmup=2, decay=0), - ) - schedules = Schedules(config, AdamWConfig(max_norm=2.0), optimizer, total_steps=4) - max_norm, metrics = schedules.apply(0) - assert max_norm == 0.0 + optimizer = torch.optim.AdamW([{"params": params[:1], "lr": 1.0}, {"params": params[1:], "lr": 0.1}]) + schedule = LRSchedule(ScheduleConfig(warmup=0, decay=2, end_ratio=0.5), optimizer, total_steps=4) + assert schedule.apply(0) == 1.0 assert [group["lr"] for group in optimizer.param_groups] == [1.0, 0.1] # Loading a state dict replaces the param group dicts. optimizer.load_state_dict(optimizer.state_dict()) - max_norm, metrics = schedules.apply(3) - assert max_norm == 2.0 + assert schedule.apply(3) == 0.5 assert [group["lr"] for group in optimizer.param_groups] == pytest.approx([0.5, 0.05]) - assert [group["betas"] for group in optimizer.param_groups] == [(0.9, 0.475)] * 2 - assert metrics["schedule/lr"] == 0.5 - assert metrics["optim/lr"] == 0.5 From f54899f7f5fdb3c230ed2bba70c96088c68ad37b Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 17:58:12 +0200 Subject: [PATCH 018/132] feat(pretrain): asynchronous eval worker on model-only snapshots With [eval_snapshots], the trainer writes a model-only DCP snapshot every interval steps and at the last step, in the background. `python -m prime_rl.trainer.pretrain.eval_worker` (torchrun on dedicated GPUs) builds the run's model once, scores the newest finished snapshot on the eval suite (ppl / ctxcont / pairs sets, ported from the pretrain-evals offline suite: fp32 scoring head, packed document-masked rows, rows dealt over the worker's GPUs), logs eval/{set}/{metric} at the snapshot's step through the run's monitors, deletes scored snapshots and stops after the last step. --- docs/pretraining.md | 13 + .../src/prime_rl/configs/pretrain.py | 41 ++- src/prime_rl/trainer/pretrain/eval_worker.py | 131 +++++++++ src/prime_rl/trainer/pretrain/evals/suite.py | 270 ++++++++++++++++++ src/prime_rl/trainer/pretrain/snapshots.py | 61 ++++ src/prime_rl/trainer/pretrain/train.py | 7 + tests/unit/train/pretrain/test_suite.py | 26 ++ 7 files changed, 548 insertions(+), 1 deletion(-) create mode 100644 src/prime_rl/trainer/pretrain/eval_worker.py create mode 100644 src/prime_rl/trainer/pretrain/evals/suite.py create mode 100644 src/prime_rl/trainer/pretrain/snapshots.py create mode 100644 tests/unit/train/pretrain/test_suite.py diff --git a/docs/pretraining.md b/docs/pretraining.md index 5bd89dfd60a..ada97c666af 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -93,3 +93,16 @@ subset = "wikitext-103-raw-v1" split = "test" ``` +**Asynchronous evals on dedicated GPUs.** With `[eval_snapshots]`, the trainer writes a model-only snapshot every `interval` steps and at the last step to `/snapshots/step_N`, in the background (it never waits for the eval side). The eval worker runs on its own GPUs (e.g. one node next to the training nodes), takes the newest finished snapshot, scores it on the eval suite (`ppl`, `ctxcont` and `pairs` sets in `//{test.parquet,meta.json}`), logs `eval/{set}/{metric}` at the snapshot's step, deletes that snapshot and every older one, and repeats until it has scored the last step. The eval frequency therefore follows from the cost of a round and the number of eval GPUs. + +```toml +[eval_snapshots] +interval = 50 +``` + +```bash +torchrun --nproc-per-node 8 -m prime_rl.trainer.pretrain.eval_worker \ + --run-dir / --data-dir ~/evals-data --sets fwedu-holdout lambada math +``` + +To log into the same W&B run as the trainer, start both in W&B shared mode with the same run id: `WANDB_SHARED_MODE=1 WANDB_RUN_ID= WANDB_SHARED_PRIMARY=trainer WANDB_SHARED_FINISHER=eval` for both, plus `WANDB_SHARED_LABEL=trainer` for the trainer and `WANDB_SHARED_LABEL=eval` for the worker. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 8b9546a53d6..9c6a48d81b8 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -6,7 +6,7 @@ from prime_rl.configs.monitors import TrainMonitorsConfig from prime_rl.configs.sft import BaseDataConfig, FakeDataConfig -from prime_rl.configs.shared import EnvVars, HeartbeatConfig, ResumeConfig, RunConfig, TrainerLogConfig +from prime_rl.configs.shared import EnvVars, HeartbeatConfig, LogConfig, ResumeConfig, RunConfig, TrainerLogConfig from prime_rl.configs.trainer import ( AdamWConfig, CheckpointConfig, @@ -118,6 +118,42 @@ class ValConfig(BaseConfig): """Held-out sets. With ``data.holdout_docs > 0`` the training holdout is added as ``holdout``.""" +class EvalSnapshotsConfig(BaseConfig): + """Model-only snapshots for the asynchronous eval worker (`prime_rl.trainer.pretrain.eval_worker`), written in + the background to ``/snapshots/step_N``. The worker deletes each one after scoring it.""" + + interval: int = Field(ge=1) + """Write a snapshot every N steps, and at the final step.""" + + +class EvalWorkerConfig(BaseConfig): + """The asynchronous eval worker: scores the newest snapshot of a training run on the eval suite, logs the + results into the run's W&B run and ``metrics.jsonl`` at the snapshot's step, and deletes scored snapshots.""" + + run_dir: Path + """The training run's directory (with ``configs/latest/resolved/pretrain.json`` and ``snapshots/``).""" + + data_dir: Path + """Eval sets, one directory per set with ``test.parquet`` and ``meta.json``.""" + + sets: list[str] + """Names of the sets to score.""" + + seq_len: int = Field(2048, ge=1) + """Window length for ``ppl`` sets and the context limit for ``ctxcont`` sets.""" + + token_budget: int = Field(32768, ge=1) + """Tokens per packed row (one forward).""" + + poll_seconds: int = Field(30, ge=1) + """Seconds to wait before looking for a new snapshot again.""" + + save_arrays: bool = True + """Write per-document arrays (npz) and summaries to ``/evals/step_N/``.""" + + log: LogConfig = LogConfig() + + class PretrainDeploymentConfig(BaseConfig): num_train_gpus: int = Field(1, ge=1) """GPUs for the trainer, on this node.""" @@ -136,6 +172,9 @@ class PretrainConfig(BaseConfig): val: ValConfig | None = None """Perplexity evals. If None, none run.""" + eval_snapshots: EvalSnapshotsConfig | None = None + """Model-only snapshots for the asynchronous eval worker. If None, none are written.""" + optim: OptimizerConfig = AdamWConfig() scheduler: SchedulerConfig = ConstantSchedulerConfig() diff --git a/src/prime_rl/trainer/pretrain/eval_worker.py b/src/prime_rl/trainer/pretrain/eval_worker.py new file mode 100644 index 00000000000..9530345f8d8 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/eval_worker.py @@ -0,0 +1,131 @@ +"""Asynchronous eval worker for pretraining runs: on its own GPUs, it scores the newest model snapshot of a +running training run (`[eval_snapshots]` in the run's config) on the eval suite, logs the results into the run's +W&B run and `metrics.jsonl` at the snapshot's step, deletes the snapshot and every older one, and repeats until +it has scored the final step. Started with torchrun on the eval GPUs: + + torchrun --nproc-per-node 8 -m prime_rl.trainer.pretrain.eval_worker --run-dir --data-dir --sets + +The model is built once with the trainer's code (`setup_model` from the run's resolved config) and kept resident; +each round loads only the weights. Rows of every set are scored data-parallel over the worker's GPUs.""" + +import asyncio +import json +import time +from datetime import timedelta + +import numpy as np +import torch +import torch.distributed as dist +from torch.distributed.checkpoint.state_dict_loader import load as dcp_load +from transformers import AutoTokenizer + +from prime_rl import monitors +from prime_rl.configs.pretrain import EvalWorkerConfig, PretrainConfig +from prime_rl.trainer.ckpt import AppState +from prime_rl.trainer.model import setup_model +from prime_rl.trainer.parallel_dims import get_parallel_dims +from prime_rl.trainer.pretrain.evals.suite import prepare_set, score_requests, summarize, use_scoring_head +from prime_rl.trainer.pretrain.snapshots import delete_snapshots_up_to, finished_snapshot_steps, snapshot_dir +from prime_rl.trainer.utils import get_dcp_group, setup_torch_distributed +from prime_rl.trainer.world import get_world +from prime_rl.utils.config import cli +from prime_rl.utils.logger import setup_logger +from prime_rl.utils.process import set_proc_title +from prime_rl.utils.utils import clean_exit + + +def run_config(config: EvalWorkerConfig) -> PretrainConfig: + """The training run's resolved config, with the parallelism of the worker: plain FSDP over its GPUs.""" + run = PretrainConfig.model_validate_json((config.run_dir / "configs/latest/resolved/pretrain.json").read_text()) + run.model.compile = None # packed rows of varying length would recompile at every forward + run.model.ep = 1 + run.model.cp = 1 + run.model.dp_replicate = 1 + return run + + +def broadcast(value: int) -> int: + values = [value] + dist.broadcast_object_list(values, src=0, group=get_dcp_group()) + return values[0] + + +@clean_exit +def main_loop(config: EvalWorkerConfig) -> None: + world = get_world() + logger = setup_logger(config.log.level, json_logging=config.log.json_logging) + run = run_config(config) + setup_torch_distributed(timeout=timedelta(seconds=run.dist_timeout_seconds)) + torch.set_float32_matmul_precision(run.matmul_precision) + torch.backends.cuda.matmul.allow_tf32 = False + + asyncio.run( + monitors.setup( + producer="eval", + wandb=run.monitors.wandb, + file=run.monitors.file, + output_dir=config.run_dir, + run_config=run, + overview_flavor="sft", + ) + ) + + model = setup_model(run.model, get_parallel_dims(run.model), loading_from_checkpoint_later=True) + model.eval() + use_scoring_head(model) + tokenizer = AutoTokenizer.from_pretrained(run.tokenizer.name, trust_remote_code=run.tokenizer.trust_remote_code) + start = time.perf_counter() + eval_sets = [prepare_set(name, config.data_dir, tokenizer, config.seq_len) for name in config.sets] + num_tokens = sum(len(ids) - begin for eval_set in eval_sets for ids, begin in eval_set.requests) + logger.info(f"Prepared {len(eval_sets)} sets ({num_tokens:,} scored tokens) in {time.perf_counter() - start:.0f}s") + + last_step = -1 + while True: + newest = broadcast(max(finished_snapshot_steps(config.run_dir), default=-1) if world.is_master else 0) + if newest <= last_step: + if run.max_steps is not None and last_step >= run.max_steps: + break + time.sleep(config.poll_seconds) + continue + + round_start = time.perf_counter() + dcp_load( + {"app": AppState(model, [], None, None)}, + checkpoint_id=snapshot_dir(config.run_dir) / f"step_{newest}", + process_group=get_dcp_group(), + ) + load_time = time.perf_counter() - round_start + metrics = {} + for eval_set in eval_sets: + scored = score_requests(model, eval_set.requests, config.token_budget, get_dcp_group()) + if not world.is_master: + continue + summary, arrays = summarize(eval_set, *scored) + metrics |= { + f"eval/{eval_set.name}/{key}": value for key, value in summary.items() if isinstance(value, float) + } + if config.save_arrays: + out = config.run_dir / "evals" / f"step_{newest}" + out.mkdir(parents=True, exist_ok=True) + np.savez(out / f"{eval_set.name}.npz", **arrays) + (out / f"{eval_set.name}.json").write_text( + json.dumps(summary | {"data_sha256": eval_set.sha256}, default=float) + ) + if world.is_master: + round_time = time.perf_counter() - round_start + metrics |= {"time/eval_round": round_time, "time/eval_load": load_time, "step": newest} + asyncio.run(monitors.log(metrics, step=newest)) + delete_snapshots_up_to(config.run_dir, newest) + logger.success(f"Evals | Step {newest} | {round_time:.0f}s | {len(eval_sets)} sets") + last_step = newest + + asyncio.run(monitors.finalize()) + + +def main(): + set_proc_title("PretrainEvalWorker") + main_loop(cli(EvalWorkerConfig)) + + +if __name__ == "__main__": + main() diff --git a/src/prime_rl/trainer/pretrain/evals/suite.py b/src/prime_rl/trainer/pretrain/evals/suite.py new file mode 100644 index 00000000000..4f03696ad01 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/evals/suite.py @@ -0,0 +1,270 @@ +"""The pretraining eval suite: log-likelihood scoring of text eval sets, as in the offline studies +(`pretrain-evals` offline suite). A set lives in ``//`` as ``test.parquet`` plus ``meta.json`` +(``kind``, ``parquet_sha256``). Kinds: + +- ``ppl``: column ``text``. Each document is BOS + text + EOS, cut into non-overlapping windows of ``seq_len`` + scored tokens (later windows start cold). Every token after BOS is scored. +- ``ctxcont``: ``context``/``continuation`` (+ ``choices``/``gold``, optional ``choice_contexts``). The + continuation is scored given BOS + context; context and continuation are tokenized jointly and the + continuation span is found by character offsets. +- ``pairs``: ``good``/``bad`` sentences, each scored as BOS + sentence. + +A set is tokenized once (`prepare_set`); every round scores its requests (`score_requests`, the model's LM head +replaced by `scoring_head`) and summarizes them (`summarize`).""" + +import json +import math +import types +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow.parquet as pq +import torch +import torch.distributed as dist +from torch import nn +from transformers.tokenization_utils import PreTrainedTokenizer + +from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX +from prime_rl.trainer.world import get_world + +POSITION_BINS = (0, 128, 512, 2048, 4096, 16384) +"""ppl loss is also reported per bin of the target's position in its window.""" + +Request = tuple[list[int], int] +"""Token ids and the index of the first scored token: ``ids[start:]`` is scored, each token given everything before it.""" + + +@dataclass +class EvalSet: + name: str + kind: str + table: Any + requests: list[Request] + context: dict + sha256: str + + +def head_logprobs(hidden, weight, targets, chunk: int = 16384): + """Log-prob and greedy hit of ``targets`` under ``hidden @ weight.T``, in fp32 chunks over tokens.""" + logprobs, hits = [], [] + for start in range(0, len(targets), chunk): + logits = hidden[start : start + chunk].float() @ weight.float().T + target = targets[start : start + chunk] + logprobs.append(logits.gather(-1, target[:, None])[:, 0] - torch.logsumexp(logits, -1)) + hits.append(logits.argmax(-1) == target) + return torch.cat(logprobs), torch.cat(hits) + + +def scoring_head(self, hidden_states, labels, temperature=None, sampling_mask=None): + """Replaces the forward of prime-rl's LM head: log-prob and greedy hit of every label that is not + IGNORE_INDEX, in fp32 without TF32, returned in the ``logprobs`` / ``entropy`` slots of its output dict.""" + hidden = hidden_states.reshape(-1, hidden_states.shape[-1]) + labels = labels.reshape(-1) + keep = labels != IGNORE_INDEX + precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") + logprobs, hits = head_logprobs(hidden[keep], self.weight, labels[keep]) + torch.set_float32_matmul_precision(precision) + return {"logprobs": logprobs, "entropy": hits} + + +def use_scoring_head(model: nn.Module) -> None: + model.lm_head.forward = types.MethodType(scoring_head, model.lm_head) + + +def pack_requests(requests: list[Request], token_budget: int) -> list[list[int]]: + """Request indices grouped into rows of at most ``token_budget`` tokens, longest first.""" + order = sorted(range(len(requests)), key=lambda i: -len(requests[i][0])) + batches, batch, total = [], [], 0 + for index in order: + length = len(requests[index][0]) - 1 + if batch and total + length > token_budget: + batches.append(batch) + batch, total = [], 0 + batch.append(index) + total += length + if batch: + batches.append(batch) + return batches + + +@torch.no_grad() +def score_requests( + model: nn.Module, requests: list[Request], token_budget: int, group: dist.ProcessGroup +) -> tuple[list[np.ndarray], list[np.ndarray]] | None: + """Per-token log-probs and greedy hits of every request, on rank 0 (None on the others). Requests are + packed into document-masked rows (positions restart per request, ``seq_lens`` marks them); the rows are + dealt round-robin over the ranks, which run the same number of forwards (each is an FSDP collective), and + the results are gathered on rank 0 over ``group``.""" + world = get_world() + batches = pack_requests(requests, token_budget) + mine = batches[world.rank :: world.world_size] + results = [] + for i in range(-(-len(batches) // world.world_size)): + batch = mine[i] if i < len(mine) else batches[0] + inputs, labels, positions, lengths = [], [], [], [] + for index in batch: + ids, start = requests[index] + inputs += ids[:-1] + labels += [IGNORE_INDEX] * (start - 1) + ids[start:] + positions += range(len(ids) - 1) + lengths.append(len(ids) - 1) + out = model( + input_ids=torch.tensor([inputs], device="cuda"), + position_ids=torch.tensor([positions], device="cuda"), + labels=torch.tensor([labels], device="cuda"), + seq_lens=torch.tensor(lengths, device="cuda"), + ) + if i < len(mine): + results.append((batch, out["logprobs"].cpu(), out["entropy"].cpu())) + gathered = [None] * world.world_size if world.is_master else None + dist.gather_object(results, gathered, dst=0, group=group) + if not world.is_master: + return None + logprobs: list[np.ndarray] = [None] * len(requests) + hits: list[np.ndarray] = [None] * len(requests) + for batch, flat_logprobs, flat_hits in (result for rank_results in gathered for result in rank_results): + offset = 0 + for index in batch: + n = len(requests[index][0]) - requests[index][1] + logprobs[index] = flat_logprobs[offset : offset + n].double().numpy() + hits[index] = flat_hits[offset : offset + n].numpy() + offset += n + return logprobs, hits + + +def continuation_ids(tokenizer: PreTrainedTokenizer, context: str, continuation: str) -> Request: + """BOS + jointly tokenized context + continuation, and the index of the first continuation token (the + first token that ends after the context; a token straddling the boundary counts as continuation).""" + encoding = tokenizer(context + continuation, add_special_tokens=False, return_offsets_mapping=True) + first = next(i for i, (_, end) in enumerate(encoding["offset_mapping"]) if end > len(context)) + return [tokenizer.bos_token_id, *encoding["input_ids"]], first + 1 + + +def truncate_left(ids: list[int], start: int, seq_len: int) -> Request: + cut = max(0, len(ids) - seq_len - 1) + assert start - cut >= 1, "continuation longer than the window" + return ids[cut:], start - cut + + +def requests_of(row: dict) -> list[tuple[str, str]]: + """(context, continuation) per option: alternative continuations (``choices``), alternative contexts with a + shared continuation (``choice_contexts``, WinoGrande), or the single gold one.""" + if row.get("choice_contexts") is not None: + return [(context, row["continuation"]) for context in row["choice_contexts"]] + if row.get("choices") is not None: + return [(row["context"], choice) for choice in row["choices"]] + return [(row["context"], row["continuation"])] + + +def prepare_set(name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_len: int) -> EvalSet: + """Load and tokenize one eval set.""" + meta = json.loads((data_dir / name / "meta.json").read_text()) + table = pq.read_table(data_dir / name / "test.parquet") + kind, requests, context = meta["kind"], [], {} + if kind == "ppl": + texts = table.column("text").to_pylist() + owner = [] + for doc, text_ids in enumerate(tokenizer(texts, add_special_tokens=False)["input_ids"]): + ids = [tokenizer.bos_token_id, *text_ids, tokenizer.eos_token_id] + for begin in range(0, len(ids) - 1, seq_len): + requests.append((ids[begin : begin + seq_len + 1], 1)) + owner.append(doc) + context = {"texts": texts, "owner": owner} + elif kind == "ctxcont": + rows = table.to_pylist() + options = [requests_of(row) for row in rows] + requests = [ + truncate_left(*continuation_ids(tokenizer, ctx, cont), seq_len) for item in options for ctx, cont in item + ] + context = {"rows": rows, "options": options} + elif kind == "pairs": + good, bad = table.column("good").to_pylist(), table.column("bad").to_pylist() + encoded = tokenizer(good + bad, add_special_tokens=False)["input_ids"] + requests = [([tokenizer.bos_token_id, *ids], 1) for ids in encoded] + context = {"num_good": len(good)} + else: + raise ValueError(f"Unknown eval set kind {kind!r} ({name})") + return EvalSet(name, kind, table, requests, context, meta["parquet_sha256"]) + + +def summarize(eval_set: EvalSet, logprobs: list[np.ndarray], hits: list[np.ndarray]) -> tuple[dict, dict]: + """Summary metrics and per-document / per-item arrays of one scored set.""" + if eval_set.kind == "ppl": + return summarize_ppl(eval_set.context["texts"], eval_set.context["owner"], logprobs) + if eval_set.kind == "ctxcont": + return summarize_ctxcont(eval_set.context["rows"], eval_set.context["options"], logprobs, hits) + return summarize_pairs(eval_set.context["num_good"], logprobs) + + +def summarize_ppl(texts: list[str], owner: list[int], logprobs: list[np.ndarray]) -> tuple[dict, dict]: + bins = list(zip(POSITION_BINS[:-1], POSITION_BINS[1:])) + nll = np.zeros(len(texts)) + tokens = np.zeros(len(texts), dtype=np.int64) + nll_bins = np.zeros((len(texts), len(bins))) + for doc, window in zip(owner, logprobs): + nll[doc] -= window.sum() + tokens[doc] += len(window) + for b, (lo, hi) in enumerate(bins): + nll_bins[doc, b] -= window[lo:hi].sum() + num_bytes = np.array([len(text.encode()) for text in texts]) + summary = { + "loss": nll.sum() / tokens.sum(), + "bpb": nll.sum() / (num_bytes.sum() * math.log(2)), + "docs": len(texts), + "tokens": int(tokens.sum()), + } + for b, (lo, hi) in enumerate(bins): + bin_tokens = sum(max(0, min(hi, len(w)) - lo) for w in logprobs) + if bin_tokens: + summary[f"loss_pos{lo}-{hi}"] = nll_bins[:, b].sum() / bin_tokens + return summary, {"nll": nll, "bytes": num_bytes, "tokens": tokens, "nll_bins": nll_bins} + + +def summarize_ctxcont(rows, options, logprobs, hits) -> tuple[dict, dict]: + choice_logprob, choice_hits, choice_tokens = [], [], [] + offset = 0 + for item in options: + item_logprobs = logprobs[offset : offset + len(item)] + choice_logprob.append(np.array([lp.sum() for lp in item_logprobs])) + choice_tokens.append([len(lp) for lp in item_logprobs]) + choice_hits.append(hits[offset : offset + len(item)]) + offset += len(item) + gold = np.array([row["gold"] if len(item) > 1 else 0 for row, item in zip(rows, options)]) + nll = np.array([-choice_logprob[i][g] for i, g in enumerate(gold)]) + tokens = np.array([choice_tokens[i][g] for i, g in enumerate(gold)]) + greedy = np.array([choice_hits[i][g].all() for i, g in enumerate(gold)]) + greedy_tokens = np.array([choice_hits[i][g].sum() for i, g in enumerate(gold)]) + num_bytes = np.array([len(row["continuation"].encode()) for row in rows]) + summary = { + "bpb": nll.sum() / (num_bytes.sum() * math.log(2)), + "loss": nll.sum() / tokens.sum(), + "greedy_em": greedy.mean(), + "greedy_token_acc": greedy_tokens.sum() / tokens.sum(), + "items": len(rows), + } + arrays = {"nll": nll, "bytes": num_bytes, "tokens": tokens, "greedy": greedy, "greedy_tokens": greedy_tokens} + mc = np.array([len(item) > 1 for item in options]) + if mc.any(): + acc = np.array([lp.argmax() == g for lp, g in zip(choice_logprob, gold)]) & mc + # length normalization only where the options differ in their continuation + has_norm = mc & np.array([row.get("choice_contexts") is None for row in rows]) + acc_norm = np.zeros(len(rows), dtype=bool) + for i in np.flatnonzero(has_norm): + per_byte = choice_logprob[i] / [len(c.encode()) for _, c in options[i]] + acc_norm[i] = per_byte.argmax() == gold[i] + summary["acc"] = acc[mc].mean() + if has_norm.any(): + summary["acc_norm"] = acc_norm[has_norm].mean() + arrays |= {"acc": acc, "acc_norm": acc_norm, "mc": mc} + return summary, arrays + + +def summarize_pairs(num_good: int, logprobs) -> tuple[dict, dict]: + sums = np.array([lp.sum() for lp in logprobs]) + good_logprob, bad_logprob = sums[:num_good], sums[num_good:] + correct = good_logprob > bad_logprob + summary = {"acc": correct.mean(), "items": num_good} + return summary, {"good_logprob": good_logprob, "bad_logprob": bad_logprob, "acc": correct} diff --git a/src/prime_rl/trainer/pretrain/snapshots.py b/src/prime_rl/trainer/pretrain/snapshots.py new file mode 100644 index 00000000000..d89487fd959 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/snapshots.py @@ -0,0 +1,61 @@ +import shutil +from pathlib import Path + +import torch.distributed.checkpoint as dcp +from torch import nn + +from prime_rl.trainer.ckpt import AppState +from prime_rl.trainer.utils import get_dcp_group +from prime_rl.trainer.world import get_world + +FINISHED_MARKER = ".finished" + + +def snapshot_dir(run_dir: Path) -> Path: + return run_dir / "snapshots" + + +def finished_snapshot_steps(run_dir: Path) -> list[int]: + """Steps of the complete snapshots, ascending.""" + paths = snapshot_dir(run_dir).glob(f"step_*/{FINISHED_MARKER}") + return sorted(int(path.parent.name.removeprefix("step_")) for path in paths) + + +def delete_snapshots_up_to(run_dir: Path, step: int) -> None: + for path in snapshot_dir(run_dir).glob("step_*"): + if int(path.name.removeprefix("step_")) <= step: + shutil.rmtree(path, ignore_errors=True) + + +class SnapshotWriter: + """Model-only DCP snapshots for the eval worker (`prime_rl.trainer.pretrain.eval_worker`), written in the + background: the trainer only waits for the copy of its shards to CPU memory. A snapshot is complete once its + `.finished` marker exists; the eval worker deletes it after scoring it.""" + + def __init__(self, run_dir: Path): + self.run_dir = run_dir + self.pending = None + + def save(self, model: nn.Module, step: int) -> None: + self.wait() + path = snapshot_dir(self.run_dir) / f"step_{step}" + future = dcp.async_save( + {"app": AppState(model, [], None, None)}, checkpoint_id=path, process_group=get_dcp_group() + ) + if get_world().is_master: + # The coordinator's save finishes after every rank's shards and the metadata are written. + future.add_done_callback(lambda done: mark_finished(path) if done.exception() is None else None) + self.pending = (future, path) + + def wait(self) -> None: + """Block until the last snapshot is on disk (at the end of training).""" + if self.pending is not None: + future, path = self.pending + future.result() + if get_world().is_master: + mark_finished(path) + self.pending = None + + +def mark_finished(path: Path) -> None: + (path / FINISHED_MARKER).touch() diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index d921a2b4503..c1d9f7e2521 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -33,6 +33,7 @@ from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter from prime_rl.trainer.pretrain.completep import completep_param_groups +from prime_rl.trainer.pretrain.snapshots import SnapshotWriter from prime_rl.trainer.pretrain.metrics import ( DeadExperts, SpikeDetector, @@ -274,6 +275,7 @@ def run_evals(step: int) -> None: dead_experts = DeadExperts() loss_spikes, grad_norm_spikes = SpikeDetector(), SpikeDetector() last_bias_moves = None + snapshots = SnapshotWriter(config.run_dir) if config.eval_snapshots is not None else None logger.info(f"Starting training loop (max_steps={config.max_steps or 'infinite'})") while True: if gc_handler is not None: @@ -391,6 +393,9 @@ def run_evals(step: int) -> None: save_ckpt_time = time.perf_counter() - save_ckpt_start_time disk_metrics = get_ckpt_disk_metrics(config.run_dir) + if snapshots is not None and (is_last_step or progress.step % config.eval_snapshots.interval == 0): + snapshots.save(model, progress.step) + eval_start_time = time.perf_counter() if config.val is not None and (is_last_step or progress.step % config.val.interval == 0): run_evals(progress.step) @@ -453,6 +458,8 @@ def run_evals(step: int) -> None: break progress.step += 1 + if snapshots is not None: + snapshots.wait() if config.ckpt is not None: logger.info(f"Saving final checkpoint at step {progress.step}") ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) diff --git a/tests/unit/train/pretrain/test_suite.py b/tests/unit/train/pretrain/test_suite.py new file mode 100644 index 00000000000..d74fba40940 --- /dev/null +++ b/tests/unit/train/pretrain/test_suite.py @@ -0,0 +1,26 @@ +import math + +import numpy as np + +from prime_rl.trainer.pretrain.evals.suite import pack_requests, summarize_ctxcont + + +def test_pack_requests_covers_every_request_within_the_budget(): + requests = [([0] * n, 1) for n in (5, 9, 3, 9, 2)] + batches = pack_requests(requests, token_budget=10) + assert sorted(i for batch in batches for i in batch) == list(range(5)) + assert all(sum(len(requests[i][0]) - 1 for i in batch) <= 10 for batch in batches) + + +def test_summarize_ctxcont_scores_the_gold_choice(): + rows = [ + {"context": "Q", "continuation": "ab", "choices": ["ab", "cd"], "gold": 0}, + {"context": "Q", "continuation": "e", "choices": ["x", "e"], "gold": 1}, + ] + options = [[("Q", c) for c in row["choices"]] for row in rows] + logprobs = [np.array([-1.0, -1.0]), np.array([-3.0]), np.array([-0.5]), np.array([-2.0])] + hits = [np.array([True, False]), np.array([False]), np.array([True]), np.array([False])] + summary, _ = summarize_ctxcont(rows, options, logprobs, hits) + assert summary["acc"] == 0.5 # item 0: gold wins; item 1: the other choice wins + assert math.isclose(summary["bpb"], (2.0 + 2.0) / (3 * math.log(2))) # gold nll / gold continuation bytes + assert summary["greedy_em"] == 0.0 From bf315f2ad4850705c8b5e91991c52ce73447ae2a Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 17:59:55 +0200 Subject: [PATCH 019/132] fix(pretrain): eval worker waits for the run's config --- src/prime_rl/trainer/pretrain/eval_worker.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/eval_worker.py b/src/prime_rl/trainer/pretrain/eval_worker.py index 9530345f8d8..c5bcb8f33d9 100644 --- a/src/prime_rl/trainer/pretrain/eval_worker.py +++ b/src/prime_rl/trainer/pretrain/eval_worker.py @@ -35,8 +35,12 @@ def run_config(config: EvalWorkerConfig) -> PretrainConfig: - """The training run's resolved config, with the parallelism of the worker: plain FSDP over its GPUs.""" - run = PretrainConfig.model_validate_json((config.run_dir / "configs/latest/resolved/pretrain.json").read_text()) + """The training run's resolved config (waiting for the launcher to write it), with the parallelism of the + worker: plain FSDP over its GPUs.""" + path = config.run_dir / "configs/latest/resolved/pretrain.json" + while not path.exists(): + time.sleep(config.poll_seconds) + run = PretrainConfig.model_validate_json(path.read_text()) run.model.compile = None # packed rows of varying length would recompile at every forward run.model.ep = 1 run.model.cp = 1 From 90f739568184f64bfb5b19a0e320f8fb4a20b5e5 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:01:12 +0200 Subject: [PATCH 020/132] fix(glm4_moe): head_dim is a config field The model already reads config.head_dim (falling back to hidden_size // num_attention_heads). As an explicit field it is part of the config class, so [model.arch] accepts it. --- .../trainer/models/glm4_moe/configuration_glm4_moe.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py index 2763679fa93..30dce92cd34 100644 --- a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py @@ -29,6 +29,8 @@ class Glm4MoeConfig(PretrainedConfig): converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group. For more details, check out [this paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`. + head_dim (`int`, *optional*): + Dimension of each attention head. Defaults to `hidden_size // num_attention_heads`. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. max_position_embeddings (`int`, *optional*, defaults to 131072): @@ -144,6 +146,7 @@ def __init__( num_attention_heads=96, partial_rotary_factor=0.5, num_key_value_heads=8, + head_dim=None, hidden_act="silu", max_position_embeddings=131072, initializer_range=0.02, @@ -174,6 +177,7 @@ def __init__( self.partial_rotary_factor = partial_rotary_factor self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps From 32b00bde6404e9deade343fb60fac5ef1114484c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:04:45 +0200 Subject: [PATCH 021/132] fix(pretrain): eval worker keeps the trainer's matmul precision API --- src/prime_rl/trainer/pretrain/eval_worker.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/prime_rl/trainer/pretrain/eval_worker.py b/src/prime_rl/trainer/pretrain/eval_worker.py index c5bcb8f33d9..42daa3ca691 100644 --- a/src/prime_rl/trainer/pretrain/eval_worker.py +++ b/src/prime_rl/trainer/pretrain/eval_worker.py @@ -61,7 +61,6 @@ def main_loop(config: EvalWorkerConfig) -> None: run = run_config(config) setup_torch_distributed(timeout=timedelta(seconds=run.dist_timeout_seconds)) torch.set_float32_matmul_precision(run.matmul_precision) - torch.backends.cuda.matmul.allow_tf32 = False asyncio.run( monitors.setup( From 8eb569612209589d0b29af8ee5c49c381d2dc8ab Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:12:43 +0200 Subject: [PATCH 022/132] test(pretrain): read a bounded number of windows from the endless token dataset --- tests/unit/train/pretrain/test_data.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index c74006b90eb..4ae2bbb36e3 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -1,3 +1,4 @@ +from itertools import islice from pathlib import Path import numpy as np @@ -36,7 +37,7 @@ def test_token_dataset_resumes_and_follows_phases(tmp_path): def dataset(): return TokenDataset(sources, phases, seq_len=4, batch_size=2, max_steps=4) - full = [window["tokens"].tolist() for window in dataset()][:8] + full = [window["tokens"].tolist() for window in islice(dataset(), 8)] assert full[:4] == [ [100, 101, 102, 103, 104], [104, 105, 106, 107, 108], @@ -51,7 +52,7 @@ def dataset(): next(iterator) resumed = dataset() resumed.load_state_dict(first.state_dict()) - assert [window["tokens"].tolist() for window in resumed][:5] == full[3:8] + assert [window["tokens"].tolist() for window in islice(resumed, 5)] == full[3:8] def test_collate_rows_restarts_positions_at_documents(): From 917b3e2d1316972dd50c4e0798315bfdc0dedafc Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:13:36 +0200 Subject: [PATCH 023/132] refactor(trainer): run_micro_batches shared by the SFT and pretrain loops The forward/backward loop over the micro-batches of a step (token-normalized gradient scaling, the offload path, NaN micro-steps, MoE routing statistics, the global loss) was copied between trainer/sft/train.py and trainer/pretrain/train.py. It is now trainer/step.py:run_micro_batches. --- src/prime_rl/trainer/pretrain/train.py | 85 ++-------------- src/prime_rl/trainer/sft/train.py | 131 ++++++------------------- src/prime_rl/trainer/step.py | 122 +++++++++++++++++++++++ 3 files changed, 159 insertions(+), 179 deletions(-) create mode 100644 src/prime_rl/trainer/step.py diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index b6feedb89a9..4d680c49f83 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -20,7 +20,6 @@ forward, get_expert_load_stats, get_full_offload_dtype_policy, - get_global_moe_stats, get_load_balance_stats, is_tt_moe_model, reshard_module, @@ -42,14 +41,11 @@ ) from prime_rl.trainer.pretrain.schedules import LRSchedule from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state +from prime_rl.trainer.step import run_micro_batches from prime_rl.trainer.utils import ( GarbageCollection, - begin_backward, clip_grad_norm_, - finish_backward, get_ckpt_disk_metrics, - prepare_gradient_offload, - scale_gradients_, setup_full_cpu_optimizer_offload, setup_torch_distributed, ) @@ -196,7 +192,6 @@ def train(config: PretrainConfig): cp_group = parallel_dims.world_mesh["cp"].get_group() if cp_enabled else None cp_size = parallel_dims.cp dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() - ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None is_moe_model = is_tt_moe_model(model) def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: @@ -259,69 +254,9 @@ def run_evals(step: int) -> None: step_start_time = time.perf_counter() lr_multiplier = lr_schedule.apply(progress.step - 1) - step_loss_sum = torch.tensor(0.0, device="cuda") - nan_loss_count = torch.tensor(0, device="cuda") - moe_stats = ( - { - "max_vio/mean": torch.tensor(0.0), - "max_vio/max": torch.tensor(0.0), - "routing_confidence/mean": torch.tensor(0.0), - } - if is_moe_model - else {} - ) - if gradient_manager is None: - micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) - step_local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") - else: - micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] - local_token_count = sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches) - global_step_token_count = torch.tensor(local_token_count, dtype=torch.int64, device="cuda") - dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) - global_token_count_val = global_step_token_count.item() // cp_size - grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val - prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=True) - - step_tokens_per_expert = 0 - step_local_num_tokens = 0 - for micro_step, micro_batch in enumerate(micro_batches): - step_local_num_tokens += micro_batch["input_ids"].shape[1] - local_loss_sum, batch_token_count = compute_loss(micro_batch) - if gradient_manager is None: - step_local_token_count += batch_token_count - if torch.isnan(local_loss_sum.detach()): - nan_loss_count += 1 - logger.warning("Local loss is nan, excluding this micro step from backward") - scaled_loss = torch.nan_to_num(local_loss_sum, nan=0.0) / grad_accum_steps - else: - step_loss_sum += local_loss_sum.detach() - scaled_loss = local_loss_sum / grad_accum_steps - begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) - scaled_loss.backward() - finish_backward(gradient_manager) - - if is_moe_model: - micro_moe_stats, tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) - for name, value in micro_moe_stats.items(): - moe_stats[f"{name}/mean"] += value / grad_accum_steps - if name == "max_vio": - moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) - step_tokens_per_expert += tokens_per_expert - + result = run_micro_batches(model, dataiter, compute_loss, grad_accum_steps, parallel_dims, gradient_manager) forward_backward_time = time.perf_counter() - step_start_time - expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) if is_moe_model else {} - - if gradient_manager is None: - global_step_token_count = step_local_token_count.clone() - dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) - global_token_count_val = global_step_token_count.item() - grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val - scale_gradients_(None, model, grad_scale) - - dist.all_reduce(step_loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) - dist.all_reduce(nan_loss_count, op=dist.ReduceOp.SUM) - batch_loss = (step_loss_sum / global_token_count_val).item() - nan_loss_count = nan_loss_count.item() + expert_load_stats = get_expert_load_stats(result.tokens_per_expert, dp_cp_group) if is_moe_model else {} grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm) if config.optim.max_norm else None optimizer.step() @@ -341,9 +276,7 @@ def run_evals(step: int) -> None: run_evals(progress.step) eval_time = time.perf_counter() - eval_start_time - global_num_tokens = torch.tensor(step_local_num_tokens, dtype=torch.int64, device="cuda") - dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) - num_tokens = global_num_tokens.item() + num_tokens = result.num_tokens progress.total_tokens += num_tokens dataset_progress = get_dataset_progress(dataloader) progress.total_samples = dataset_progress["step"] @@ -354,7 +287,7 @@ def run_evals(step: int) -> None: mfu = perf_counter.get_step_mfu(num_tokens, step_time) peak_memory = torch.cuda.max_memory_reserved() / 1024**3 - step_message = f"Step {progress.step} | {step_time:.2f}s | Loss {batch_loss:.4f}" + step_message = f"Step {progress.step} | {step_time:.2f}s | Loss {result.loss:.4f}" if grad_norm is not None: step_message += f" | Grad. Norm {grad_norm:.4f}" step_message += ( @@ -364,9 +297,9 @@ def run_evals(step: int) -> None: logger.success(step_message) metrics = { - "loss/mean": batch_loss, - "loss/perplexity": math.exp(min(batch_loss, 20)), - "loss/nan_count": nan_loss_count, + "loss/mean": result.loss, + "loss/perplexity": math.exp(min(result.loss, 20)), + "loss/nan_count": result.nan_count, "optim/lr": current_lr, **{f"progress/{name}/epochs": epochs for name, epochs in dataset_progress["epochs"].items()}, **{f"progress/{name}/num_tokens": tokens for name, tokens in dataset_progress["tokens"].items()}, @@ -381,7 +314,7 @@ def run_evals(step: int) -> None: "time/save_ckpt": save_ckpt_time, "time/evals": eval_time, **get_ckpt_disk_metrics(config.run_dir), - **{name: value.item() for name, value in moe_stats.items()}, + **result.moe_stats, **expert_load_stats, "schedule/lr": lr_multiplier, "step": progress.step, diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 98727b861b9..8aa104721a5 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -28,7 +28,6 @@ forward, get_full_offload_dtype_policy, get_expert_load_stats, - get_global_moe_stats, get_load_balance_stats, is_tt_moe_model, reshard_module, @@ -45,16 +44,13 @@ setup_dataloader, setup_dataset, ) +from prime_rl.trainer.step import run_micro_batches from prime_rl.trainer.utils import ( GarbageCollection, MemoryProfiler, - begin_backward, clip_grad_norm_, - finish_backward, get_ckpt_disk_metrics, - prepare_gradient_offload, print_sample, - scale_gradients_, setup_full_cpu_optimizer_offload, setup_torch_distributed, ) @@ -235,8 +231,8 @@ def train(config: SFTConfig): cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() if cp_enabled else 0 cp_group = parallel_dims.world_mesh["cp"].get_group() if cp_enabled else None dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() - ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None cp_size = parallel_dims.cp + is_moe_model = is_tt_moe_model(model) def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: """Forward pass returning (loss_sum, token_count) over unmasked tokens.""" @@ -298,6 +294,13 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: del out return loss_sum, token_count + def compute_train_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: + if config.log.log_data: + print_sample( + micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer + ) + return compute_loss(micro_batch) + maybe_record_function = nullcontext def run_eval_loop(data_iter): @@ -397,7 +400,7 @@ def is_online_eval_step(step: int) -> bool: if config.trace_path: logger.info(f"Tracing to {config.trace_path}") prof = profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True).__enter__() - maybe_record_function = record_function # noqa: F841 – captured by run_forward_loop closure + maybe_record_function = record_function max_peak_memory = 0.0 while True: # Reset peak memory stats @@ -413,88 +416,22 @@ def is_online_eval_step(step: int) -> bool: step_start_time = time.perf_counter() forward_backward_start_time = time.perf_counter() - step_loss_sum = torch.tensor(0.0, device="cuda") - nan_loss_count = torch.tensor(0, device="cuda") - is_moe_model = is_tt_moe_model(model) - moe_stats = ( - { - "max_vio/mean": torch.tensor(0.0), - "max_vio/max": torch.tensor(0.0), - "routing_confidence/mean": torch.tensor(0.0), - } - if is_moe_model - else {} - ) run_validation_this_step = config.val is not None and ( (is_first_step and config.val.eval_on_start) or (not is_first_step and progress.step % config.val.interval == 0) ) - if gradient_manager is None: - micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) - step_local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") - else: - micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] - local_token_count = sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches) - global_step_token_count = torch.tensor(local_token_count, dtype=torch.int64, device="cuda") - dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) - global_token_count_val = global_step_token_count.item() // cp_size - grad_scale = ( - parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val - if global_token_count_val > 0 - else 1.0 - ) - prepare_gradient_offload( - gradient_manager, - grad_scale, - overlap_optimizer=not run_validation_this_step, - ) - - step_tokens_per_expert = 0 - step_local_num_tokens = 0 - for micro_step, micro_batch in enumerate(micro_batches): - step_local_num_tokens += micro_batch["input_ids"].shape[1] - if config.log.log_data: - print_sample( - micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer - ) - - with maybe_record_function("forward"): - local_loss_sum, batch_token_count = compute_loss(micro_batch) - - if gradient_manager is None: - step_local_token_count += batch_token_count - - if torch.isnan(local_loss_sum.detach()): - nan_loss_count += 1 - logger.warning("Local loss is nan, excluding this micro step from backward") - scaled_loss = torch.nan_to_num(local_loss_sum, nan=0.0) / grad_accum_steps - else: - step_loss_sum += local_loss_sum.detach() - scaled_loss = local_loss_sum / grad_accum_steps - - with maybe_record_function("backward"): - begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) - scaled_loss.backward() - finish_backward(gradient_manager) - - if is_moe_model: - micro_moe_stats, tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) - for name, value in micro_moe_stats.items(): - moe_stats[f"{name}/mean"] += value / grad_accum_steps - if name == "max_vio": - moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) - step_tokens_per_expert += tokens_per_expert - + result = run_micro_batches( + model, + dataiter, + compute_train_loss, + grad_accum_steps, + parallel_dims, + gradient_manager, + overlap_optimizer=not run_validation_this_step, + record_function=maybe_record_function, + ) forward_backward_time = time.perf_counter() - forward_backward_start_time - expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) if is_moe_model else {} - - if gradient_manager is None: - global_step_token_count = step_local_token_count.clone() - dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) - global_token_count_val = global_step_token_count.item() - if global_token_count_val > 0: - grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val - scale_gradients_(None, model, grad_scale) + expert_load_stats = get_expert_load_stats(result.tokens_per_expert, dp_cp_group) if is_moe_model else {} # Run validation after forward-backward (so torch.compile sees training graph first) but before # optimizer step (so eval_on_start evaluates untrained weights) @@ -505,15 +442,6 @@ def is_online_eval_step(step: int) -> bool: # unsharded parameters. Reshard so clipping sees every gradient. reshard_module(model) - # Compute the global mean loss for logging. - dist.all_reduce(step_loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) - dist.all_reduce(nan_loss_count, op=dist.ReduceOp.SUM) - if global_token_count_val > 0: - batch_loss = (step_loss_sum / global_token_count_val).item() - else: - batch_loss = 0.0 - nan_loss_count = nan_loss_count.item() - grad_norm: torch.Tensor | None = None if config.optim.max_norm is not None: logger.debug(f"Clipping gradients with max norm {config.optim.max_norm}") @@ -550,10 +478,7 @@ def is_online_eval_step(step: int) -> bool: if memory_profiler is not None: memory_profiler.step() - # The dp mesh excludes cp, whose ranks hold the same rows. - global_num_tokens = torch.tensor(step_local_num_tokens, dtype=torch.int64, device="cuda") - dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) - num_tokens = global_num_tokens.item() + num_tokens = result.num_tokens progress.total_tokens += num_tokens dataset_progress = get_dataset_progress(dataloader) progress.total_samples = dataset_progress["step"] @@ -568,13 +493,13 @@ def is_online_eval_step(step: int) -> bool: max_peak_memory = max(max_peak_memory, peak_memory) # Log step metrics - step_message = f"Step {progress.step} | {format_time(step_time):>7} | Loss {batch_loss:.4f}" + step_message = f"Step {progress.step} | {format_time(step_time):>7} | Loss {result.loss:.4f}" if grad_norm is not None: step_message += f" | Grad. Norm {grad_norm:.4f}" step_message += f" | LR {current_lr:.2e} | Throughput {throughput:.0f} tokens/s | MFU {mfu:.1f}% | Peak Mem. {peak_memory:.1f}/{max_memory:.1f} GiB ({peak_memory / max_memory * 100:.1f}%)" if is_moe_model: for name, label in (("max_vio/mean", "Max Vio"), ("routing_confidence/mean", "Routing Conf.")): - value = moe_stats[name].item() + value = result.moe_stats[name] if value > 0: step_message += f" | {label} {value:.4f}" logger.success(step_message) @@ -624,9 +549,9 @@ def is_online_eval_step(step: int) -> bool: asyncio.run(monitors.log(optim_metrics, step=progress.step)) loss_log_metrics = { - "loss/mean": batch_loss, - "loss/perplexity": math.exp(min(batch_loss, 20)), - "loss/nan_count": nan_loss_count, + "loss/mean": result.loss, + "loss/perplexity": math.exp(min(result.loss, 20)), + "loss/nan_count": result.nan_count, "step": progress.step, } # Log tensor stats @@ -647,7 +572,7 @@ def is_online_eval_step(step: int) -> bool: disk_metrics["step"] = progress.step asyncio.run(monitors.log(disk_metrics, step=progress.step)) - moe_log_metrics = {name: value.item() for name, value in moe_stats.items()} | expert_load_stats + moe_log_metrics = result.moe_stats | expert_load_stats if moe_log_metrics: asyncio.run(monitors.log({**moe_log_metrics, "step": progress.step}, step=progress.step)) diff --git a/src/prime_rl/trainer/step.py b/src/prime_rl/trainer/step.py new file mode 100644 index 00000000000..c94b656536f --- /dev/null +++ b/src/prime_rl/trainer/step.py @@ -0,0 +1,122 @@ +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch import Tensor, nn + +from prime_rl.trainer.model import get_global_moe_stats, is_tt_moe_model +from prime_rl.trainer.parallel_dims import ParallelDims +from prime_rl.trainer.utils import begin_backward, finish_backward, prepare_gradient_offload, scale_gradients_ +from prime_rl.utils.logger import get_logger + + +@dataclass +class StepResult: + loss: float + """Mean loss per unmasked token over the global batch.""" + + nan_count: int + """Micro-batches with a NaN loss over all ranks. They are left out of the backward pass.""" + + num_tokens: int + """Input tokens of the global batch.""" + + moe_stats: dict[str, float] + """Routing statistics averaged over the micro-batches (``max_vio/mean``, ``max_vio/max``, + ``routing_confidence/mean``). Empty for dense models.""" + + tokens_per_expert: Tensor | None + """Global ``[num_moe_layers, num_experts]`` token counts of the step. None for dense models.""" + + +def run_micro_batches( + model: nn.Module, + dataiter: Iterator[dict], + compute_loss: Callable[[dict], tuple[Tensor, Tensor]], + grad_accum_steps: int, + parallel_dims: ParallelDims, + gradient_manager=None, + *, + overlap_optimizer: bool = True, + record_function: Callable[[str], AbstractContextManager] = nullcontext, +) -> StepResult: + """Forward and backward passes over the ``grad_accum_steps`` micro-batches of one step, leaving the + gradients of the summed token loss divided by the global token count (the mean loss per token, as on one + GPU without accumulation). ``compute_loss`` returns the summed loss and the number of unmasked tokens of a + micro-batch.""" + logger = get_logger() + dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() + ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None + is_moe_model = is_tt_moe_model(model) + + if gradient_manager is None: + micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) + local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") + else: + # The offloaded gradients are scaled as they stream out, so the token count is needed up front. + micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] + global_token_count = torch.tensor( + sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches), device="cuda" + ) + dist.all_reduce(global_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count = global_token_count.item() // parallel_dims.cp + grad_scale = ( + parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count + if global_token_count > 0 + else 1.0 + ) + prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=overlap_optimizer) + + loss_sum = torch.tensor(0.0, device="cuda") + nan_count = torch.tensor(0, device="cuda") + moe_stats = {name: torch.tensor(0.0) for name in ("max_vio/mean", "max_vio/max", "routing_confidence/mean")} + tokens_per_expert = None + local_num_tokens = 0 + for micro_step, micro_batch in enumerate(micro_batches): + local_num_tokens += micro_batch["input_ids"].shape[1] + with record_function("forward"): + micro_loss_sum, micro_token_count = compute_loss(micro_batch) + if gradient_manager is None: + local_token_count += micro_token_count + if torch.isnan(micro_loss_sum.detach()): + nan_count += 1 + logger.warning("Local loss is nan, excluding this micro step from backward") + scaled_loss = torch.nan_to_num(micro_loss_sum, nan=0.0) / grad_accum_steps + else: + loss_sum += micro_loss_sum.detach() + scaled_loss = micro_loss_sum / grad_accum_steps + with record_function("backward"): + begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) + scaled_loss.backward() + finish_backward(gradient_manager) + + if is_moe_model: + micro_moe_stats, micro_tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) + for name, value in micro_moe_stats.items(): + moe_stats[f"{name}/mean"] += value / grad_accum_steps + if name == "max_vio": + moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) + tokens_per_expert = micro_tokens_per_expert + (0 if tokens_per_expert is None else tokens_per_expert) + + if gradient_manager is None: + dist.all_reduce(local_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count = local_token_count.item() + if global_token_count > 0: + scale_gradients_( + None, model, parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count + ) + + dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) + dist.all_reduce(nan_count, op=dist.ReduceOp.SUM) + # The dp mesh excludes cp, whose ranks hold the same rows. + num_tokens = torch.tensor(local_num_tokens, dtype=torch.int64, device="cuda") + dist.all_reduce(num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) + return StepResult( + loss=(loss_sum / global_token_count).item() if global_token_count > 0 else 0.0, + nan_count=int(nan_count.item()), + num_tokens=int(num_tokens.item()), + moe_stats={name: value.item() for name, value in moe_stats.items()} if is_moe_model else {}, + tokens_per_expert=tokens_per_expert, + ) From 20578b007ce1bbe4b18c2151fb8a3338cbedb1c3 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:13:36 +0200 Subject: [PATCH 024/132] refactor(configs): TrainRunConfig shared by SFTConfig and PretrainConfig The run, model, optimizer, checkpoint, monitor, deployment and SLURM fields and their validators (DeepEP and full offload disable clipping, full offload needs AdamW/SignSGD, multi-node needs SLURM and sets dp_replicate, run identity, tokenizer trust_remote_code) live in one base class. Pretrain gains the DeepEP and offload validators it was missing; each config only names its run (default_run_name). --- .../src/prime_rl/configs/pretrain.py | 76 +---- .../src/prime_rl/configs/sft.py | 264 +++++++++--------- 2 files changed, 142 insertions(+), 198 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index f42c8de5c49..c21dea1ec9f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -5,18 +5,15 @@ from pydantic import Field, model_validator -from prime_rl.configs.monitors import TrainMonitorsConfig -from prime_rl.configs.sft import SFTDeploymentConfig, SingleNodeDeploymentConfig, setup_trainer_deployment -from prime_rl.configs.shared import EnvVars, HeartbeatConfig, ResumeConfig, RunConfig, SlurmConfig, TrainerLogConfig +from prime_rl.configs.sft import TrainRunConfig from prime_rl.configs.trainer import ( CheckpointConfig, - GCConfig, ModelConfig, MuonConfig, OptimizerConfig, TokenizerConfig, ) -from prime_rl.utils.config import BaseConfig, default_output_dir +from prime_rl.utils.config import BaseConfig TOKENIZER_FIELDS = ("vocab_size", "bos_token_id", "eos_token_id", "pad_token_id") @@ -183,12 +180,9 @@ class ScheduleConfig(BaseConfig): """Shape of the decay from 1 to ``end_ratio``.""" -class PretrainConfig(BaseConfig): +class PretrainConfig(TrainRunConfig): model: PretrainModelConfig - env_vars: EnvVars = {} - """Extra environment variables for the trainer processes.""" - tokenizer: TokenizerConfig """The tokenizer (``tokenizer.name`` is required). It sets the vocabulary and special tokens of the model.""" @@ -206,51 +200,12 @@ class PretrainConfig(BaseConfig): ckpt: CheckpointConfig | None = CheckpointConfig() """Checkpointing. Defaults to a checkpoint at the end of training.""" - resume: ResumeConfig | None = None - """Resume the run from a checkpoint. If None, does not resume.""" - - log: TrainerLogConfig = TrainerLogConfig() - - monitors: TrainMonitorsConfig = TrainMonitorsConfig() - """Metric monitors (``monitors.wandb``, ``monitors.file``).""" - - run: RunConfig = Field(default_factory=RunConfig) - """Run metadata. ``run.name`` names the run directory under ``output_dir``.""" - - output_dir: Path = Field(default_factory=default_output_dir) - """Directory that groups related runs. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``.""" - - clean: bool = False - """Delete the run directory before starting. Required to overwrite a previous run's artifacts when not resuming.""" - - matmul_precision: Literal["highest", "high", "medium"] = "high" - """Precision for float32 matrix multiplications (``torch.set_float32_matmul_precision``).""" - max_steps: int = Field(ge=1) """Training steps.""" - gc: GCConfig | None = GCConfig() - """Deterministic garbage collection every N steps. Set to null for Python's default GC.""" - - dist_timeout_seconds: int = 900 - """Timeout in seconds for torch distributed ops.""" - - heartbeat: HeartbeatConfig | None = None - """BetterStack heartbeat configuration.""" - - deployment: SFTDeploymentConfig = SingleNodeDeploymentConfig() - """``single_node`` (``num_train_gpus`` on this node) or ``multi_node`` (``num_train_nodes`` x ``gpus_per_node``, needs ``[slurm]``). The online-eval fields (``num_infer_*``) are not used.""" - - slurm: SlurmConfig | None = None - """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" - - dry_run: bool = False - """Only validate and dump resolved configs, then exit.""" - - @property - def run_dir(self) -> Path: - assert self.run.dir is not None # resolved at construction - return self.output_dir / self.run.dir + def default_run_name(self) -> str: + """``--``.""" + return f"{self.model.arch['model_type']}--{uuid.uuid4().hex[:8]}" @model_validator(mode="after") def resolve_autos(self): @@ -269,16 +224,6 @@ def validate_lr(self): raise ValueError("optim.lr is required") return self - @model_validator(mode="after") - def auto_setup_run_identity(self): - if self.run.name is None: - self.run.name = f"{self.model.arch['model_type']}--{uuid.uuid4().hex[:8]}" - if self.run.dir is None: - self.run.dir = self.run.name - if self.monitors.wandb is not None and self.monitors.wandb.name is None: - self.monitors.wandb.name = self.run.name - return self - @model_validator(mode="after") def auto_setup_model_dir(self): """The model is trained from ``/model``, written by the launcher from ``model.arch``.""" @@ -286,16 +231,9 @@ def auto_setup_model_dir(self): return self @model_validator(mode="after") - def auto_setup_tokenizer(self): + def validate_tokenizer(self): if self.tokenizer.name is None: raise ValueError("tokenizer.name is required") - if self.tokenizer.trust_remote_code is None: - self.tokenizer.trust_remote_code = self.model.trust_remote_code - return self - - @model_validator(mode="after") - def auto_setup_deployment(self): - setup_trainer_deployment(self) return self @model_validator(mode="after") diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 682974a3770..05c709d1415 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -1,7 +1,7 @@ import uuid import warnings from pathlib import Path -from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias +from typing import Annotated, Literal, TypeAlias from urllib.parse import urlparse from pydantic import AliasChoices, Field, model_validator @@ -36,9 +36,6 @@ ) from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource -if TYPE_CHECKING: - from prime_rl.configs.pretrain import PretrainConfig - class BaseDataConfig(BaseConfig): batch_size: int = Field(128, ge=1) @@ -219,62 +216,18 @@ class MultiNodeDeploymentConfig(BaseDeploymentConfig): ] -def setup_trainer_deployment(config: "SFTConfig | PretrainConfig") -> None: - """Shared by ``sft`` and ``pretrain``: a multi-node deployment needs SLURM and sets - ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the SLURM job name defaults to the - run name and the template to the bundled single- or multi-node trainer template.""" - deployment = config.deployment - if deployment.type == "multi_node": - if config.slurm is None: - raise ValueError("Must use SLURM for multi-node deployment.") - if deployment.nodes_per_fsdp_group is not None: - if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: - raise ValueError( - f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " - f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" - ) - config.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group - if config.slurm is not None: - if config.slurm.job_name is None: - config.slurm.job_name = config.run.name - if config.slurm.template_path is None: - templates_dir = find_package_resource("templates") - if templates_dir is not None: - config.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" - +class TrainRunConfig(BaseConfig): + """Fields and validators shared by the single-process trainers (``sft`` and ``pretrain``).""" -class SFTConfig(BaseConfig): model: ModelConfig = ModelConfig() env_vars: EnvVars = {} - """Extra environment variables for the SFT trainer process(es). Merged on top of the launcher defaults.""" + """Extra environment variables for the trainer process(es). Merged on top of the launcher defaults.""" tokenizer: TokenizerConfig = TokenizerConfig() - renderer: RendererConfig = AutoRendererConfig() - """Renderer config. Defaults to auto-selecting from the tokenizer model name.""" - - data: DataConfig = SFTDataConfig() - - val: SFTValConfig | None = None - """Validation configuration. If None, no validation runs.""" - - eval: SFTOnlineEvalConfig | None = None - """Online evaluation configuration: rollout-based evals against a live inference - server that receives the trainer's weight broadcasts. If None, no online evals run.""" - - inference: InferenceConfig | None = None - """Inference server for online evals. If None (with ``eval`` set), the launcher - does not start a server and evals connect to ``eval.client.base_url``.""" - - weight_broadcast: WeightBroadcastConfig | None = None - """Trainer-to-inference weight transport for online evals. Defaults to NCCL. - LoRA and external inference use filesystem broadcast.""" - optim: OptimizerConfig = AdamWConfig() - scheduler: SchedulerConfig = ConstantSchedulerConfig() - ckpt: CheckpointConfig | None = None resume: ResumeConfig | None = None @@ -294,65 +247,15 @@ class SFTConfig(BaseConfig): clean: bool = False """Delete the run directory (``output_dir / run.name``) before starting training. Required to overwrite a run directory that contains artifacts from a previous run when not resuming.""" - @property - def run_dir(self) -> Path: - assert self.run.dir is not None # resolved at construction - return self.output_dir / self.run.dir - - @model_validator(mode="after") - def resolve_moe_router_dtype_auto(self): - """Resolve ``model.moe_router_dtype='auto'``: SFT defaults to bf16 routing, skipping the fp32 gate GEMM and its fp32 FSDP unit.""" - if self.model.moe_router_dtype == "auto": - self.model.moe_router_dtype = "bfloat16" - return self - - @model_validator(mode="after") - def resolve_weight_decay_auto(self): - """Resolve ``optim.weight_decay='auto'``: SFT keeps the historical 0.01 default — standard L2 regularization for supervised training.""" - if self.optim.weight_decay == "auto": - self.optim.weight_decay = 0.01 - return self - - @model_validator(mode="after") - def resolve_freeze_moe_router_auto(self): - """Resolve ``model.freeze_moe_router='auto'``: SFT trains the router.""" - if self.model.freeze_moe_router == "auto": - self.model.freeze_moe_router = False - return self - - @model_validator(mode="after") - def auto_setup_run_identity(self): - """Auto-generate the run name (``----``) when unset and - default the run directory, W&B run name and platform run name to it when not - set explicitly.""" - if self.run.name is None: - dataset = str(getattr(self.data, "name", "")).split("/")[-1] - model = self.model.name.split("/")[-1] - parts = [part for part in (dataset, model) if part] - self.run.name = "--".join([*parts, uuid.uuid4().hex[:8]]).lower() - if self.run.dir is None: - self.run.dir = self.run.name - if self.monitors.wandb is not None and self.monitors.wandb.name is None: - self.monitors.wandb.name = self.run.name - if self.monitors.prime is not None and self.monitors.prime.name is None: - self.monitors.prime.name = self.run.name - return self - matmul_precision: Literal["highest", "high", "medium"] = "high" """Precision for float32 matrix multiplications. ``highest`` is full FP32 (required on ROCm/AMD GPUs to avoid catastrophic precision loss in softmax over large vocabularies). ``high`` enables TF32 on NVIDIA GPUs for a speedup with minor precision tradeoff. See ``torch.set_float32_matmul_precision``.""" max_steps: int | None = None """Maximum training steps. If None, runs indefinitely.""" - memory_profiler_path: Path | None = None - """Path to write the memory profile to.""" - gc: GCConfig | None = GCConfig() """Garbage collection config. Disables automatic GC and runs deterministic collections every N steps to avoid stragglers. Set to null to use Python's default GC behavior.""" - trace_path: Path | None = None - """Path to write the PyTorch profiler trace to.""" - dist_timeout_seconds: int = 900 """Timeout in seconds for torch distributed ops: a hung collective (e.g. a dead rank) fails the run after this long.""" @@ -364,14 +267,13 @@ def auto_setup_run_identity(self): slurm: SlurmConfig | None = None """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" - dashboard: bool = True - """Make sure a local dashboard daemon serves this run's output dir (started on - demand in interactive sessions; an already-running daemon's URL is logged).""" - dry_run: bool = False """Only validate and dump resolved configs, then exit early.""" - ### Pre-validation normalization + @property + def run_dir(self) -> Path: + assert self.run.dir is not None # resolved at construction + return self.output_dir / self.run.dir @model_validator(mode="before") @classmethod @@ -384,23 +286,22 @@ def normalize_deployment(cls, data): deployment.pop(key, None) return data - @model_validator(mode="before") - @classmethod - def propagate_model_to_inference(cls, data): - """Fill ``inference.vllm.model`` from ``model.name`` before InferenceConfig - construction, so its parser auto-resolution sees the right model name.""" - if not isinstance(data, dict): - return data - inference = data.get("inference") - model = data.get("model") - name = model.get("name") if isinstance(model, dict) else None - if isinstance(inference, dict) and name: - vllm = inference.setdefault("vllm", {}) - if isinstance(vllm, dict): - vllm.setdefault("model", name) - return data + def default_run_name(self) -> str: + raise NotImplementedError - ### Validate configs (e.g. raise for unsupported (combinations of) configs) + @model_validator(mode="after") + def auto_setup_run_identity(self): + """Default the run name (see ``default_run_name``) and default the run directory, W&B run name and + platform run name to it when not set explicitly.""" + if self.run.name is None: + self.run.name = self.default_run_name() + if self.run.dir is None: + self.run.dir = self.run.name + if self.monitors.wandb is not None and self.monitors.wandb.name is None: + self.monitors.wandb.name = self.run.name + if self.monitors.prime is not None and self.monitors.prime.name is None: + self.monitors.prime.name = self.run.name + return self @model_validator(mode="after") def deepep_disables_grad_clipping(self): @@ -430,6 +331,108 @@ def full_optimizer_offload_disables_grad_clipping(self): self.optim.max_norm = None return self + @model_validator(mode="after") + def auto_setup_deployment(self): + """A multi-node deployment needs SLURM and sets ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the + SLURM job name defaults to the run name and the template to the bundled single- or multi-node template.""" + deployment = self.deployment + if deployment.type == "multi_node": + if self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + if deployment.nodes_per_fsdp_group is not None: + if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" + ) + self.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group + if self.slurm is not None: + if self.slurm.job_name is None: + self.slurm.job_name = self.run.name + if self.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + self.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" + return self + + @model_validator(mode="after") + def auto_setup_tokenizer_trust_remote_code(self): + if self.tokenizer.trust_remote_code is None: + self.tokenizer.trust_remote_code = self.model.trust_remote_code + return self + + +class SFTConfig(TrainRunConfig): + def default_run_name(self) -> str: + """``----``.""" + dataset = str(getattr(self.data, "name", "")).split("/")[-1] + model = self.model.name.split("/")[-1] + parts = [part for part in (dataset, model) if part] + return "--".join([*parts, uuid.uuid4().hex[:8]]).lower() + + renderer: RendererConfig = AutoRendererConfig() + """Renderer config. Defaults to auto-selecting from the tokenizer model name.""" + + data: DataConfig = SFTDataConfig() + + val: SFTValConfig | None = None + """Validation configuration. If None, no validation runs.""" + + eval: SFTOnlineEvalConfig | None = None + """Online evaluation configuration: rollout-based evals against a live inference + server that receives the trainer's weight broadcasts. If None, no online evals run.""" + + inference: InferenceConfig | None = None + """Inference server for online evals. If None (with ``eval`` set), the launcher + does not start a server and evals connect to ``eval.client.base_url``.""" + + weight_broadcast: WeightBroadcastConfig | None = None + """Trainer-to-inference weight transport for online evals. Defaults to NCCL. + LoRA and external inference use filesystem broadcast.""" + + scheduler: SchedulerConfig = ConstantSchedulerConfig() + + @model_validator(mode="after") + def resolve_moe_router_dtype_auto(self): + """Resolve ``model.moe_router_dtype='auto'``: SFT defaults to bf16 routing, skipping the fp32 gate GEMM and its fp32 FSDP unit.""" + if self.model.moe_router_dtype == "auto": + self.model.moe_router_dtype = "bfloat16" + return self + + @model_validator(mode="after") + def resolve_weight_decay_auto(self): + """Resolve ``optim.weight_decay='auto'``: SFT keeps the historical 0.01 default — standard L2 regularization for supervised training.""" + if self.optim.weight_decay == "auto": + self.optim.weight_decay = 0.01 + return self + + @model_validator(mode="after") + def resolve_freeze_moe_router_auto(self): + """Resolve ``model.freeze_moe_router='auto'``: SFT trains the router.""" + if self.model.freeze_moe_router == "auto": + self.model.freeze_moe_router = False + return self + + ### Pre-validation normalization + + @model_validator(mode="before") + @classmethod + def propagate_model_to_inference(cls, data): + """Fill ``inference.vllm.model`` from ``model.name`` before InferenceConfig + construction, so its parser auto-resolution sees the right model name.""" + if not isinstance(data, dict): + return data + inference = data.get("inference") + model = data.get("model") + name = model.get("name") if isinstance(model, dict) else None + if isinstance(inference, dict) and name: + vllm = inference.setdefault("vllm", {}) + if isinstance(vllm, dict): + vllm.setdefault("model", name) + return data + + ### Validate configs (e.g. raise for unsupported (combinations of) configs) + @model_validator(mode="after") def auto_setup_online_eval(self): """Wire online evals and the trainer-to-inference weight transport.""" @@ -642,15 +645,18 @@ def validate_scheduler_steps(self): ### Auto-setup and validate shared configs + memory_profiler_path: Path | None = None + """Path to write the memory profile to.""" + + trace_path: Path | None = None + """Path to write the PyTorch profiler trace to.""" + + dashboard: bool = True + """Make sure a local dashboard daemon serves this run's output dir (started on + demand in interactive sessions; an already-running daemon's URL is logged).""" + @model_validator(mode="after") def auto_setup_tokenizer(self): if self.tokenizer.name is None: self.tokenizer.name = self.model.name - if self.tokenizer.trust_remote_code is None: - self.tokenizer.trust_remote_code = self.model.trust_remote_code - return self - - @model_validator(mode="after") - def auto_setup_deployment(self): - setup_trainer_deployment(self) return self From cc367f7c27a9183c48e0cc9996a1ba48c64380ba Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:16:56 +0200 Subject: [PATCH 025/132] fix(pretrain): snapshots clone the shards on the GPU and write them in a thread --- src/prime_rl/trainer/pretrain/snapshots.py | 41 +++++++++++----------- 1 file changed, 20 insertions(+), 21 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/snapshots.py b/src/prime_rl/trainer/pretrain/snapshots.py index d89487fd959..512aa10fbb6 100644 --- a/src/prime_rl/trainer/pretrain/snapshots.py +++ b/src/prime_rl/trainer/pretrain/snapshots.py @@ -1,8 +1,11 @@ import shutil +from concurrent.futures import Future, ThreadPoolExecutor from pathlib import Path -import torch.distributed.checkpoint as dcp +import torch from torch import nn +from torch.distributed.checkpoint.state_dict_saver import save as dcp_save +from torch.utils._pytree import tree_map_only from prime_rl.trainer.ckpt import AppState from prime_rl.trainer.utils import get_dcp_group @@ -28,34 +31,30 @@ def delete_snapshots_up_to(run_dir: Path, step: int) -> None: class SnapshotWriter: - """Model-only DCP snapshots for the eval worker (`prime_rl.trainer.pretrain.eval_worker`), written in the - background: the trainer only waits for the copy of its shards to CPU memory. A snapshot is complete once its - `.finished` marker exists; the eval worker deletes it after scoring it.""" + """Model-only DCP snapshots for the eval worker (`prime_rl.trainer.pretrain.eval_worker`), in the layout of a + model-only checkpoint (`{"app": AppState}`). The trainer only waits for a copy of its weight shards on the GPU; + a background thread writes them. A snapshot is complete once its `.finished` marker exists; the eval worker + deletes it after scoring it.""" def __init__(self, run_dir: Path): self.run_dir = run_dir - self.pending = None + self.executor = ThreadPoolExecutor(max_workers=1) + self.pending: Future | None = None + # Created here, on every rank at the same point; the writer thread only uses it. + self.group = get_dcp_group() def save(self, model: nn.Module, step: int) -> None: self.wait() - path = snapshot_dir(self.run_dir) / f"step_{step}" - future = dcp.async_save( - {"app": AppState(model, [], None, None)}, checkpoint_id=path, process_group=get_dcp_group() - ) + state = tree_map_only(torch.Tensor, torch.Tensor.clone, AppState(model, [], None, None).state_dict()) + self.pending = self.executor.submit(self.write, state, snapshot_dir(self.run_dir) / f"step_{step}") + + def write(self, state: dict, path: Path) -> None: + dcp_save({"app": state}, checkpoint_id=path, process_group=self.group) if get_world().is_master: - # The coordinator's save finishes after every rank's shards and the metadata are written. - future.add_done_callback(lambda done: mark_finished(path) if done.exception() is None else None) - self.pending = (future, path) + (path / FINISHED_MARKER).touch() def wait(self) -> None: - """Block until the last snapshot is on disk (at the end of training).""" + """Block until the previous snapshot is on disk.""" if self.pending is not None: - future, path = self.pending - future.result() - if get_world().is_master: - mark_finished(path) + self.pending.result() self.pending = None - - -def mark_finished(path: Path) -> None: - (path / FINISHED_MARKER).touch() From 224df6a91cde518a3c422e1cb0c944555d3804a7 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:33:52 +0200 Subject: [PATCH 026/132] feat(pretrain): bf16 snapshots (router kept in fp32) --- src/prime_rl/trainer/pretrain/snapshots.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/snapshots.py b/src/prime_rl/trainer/pretrain/snapshots.py index 512aa10fbb6..c9e94f30b1f 100644 --- a/src/prime_rl/trainer/pretrain/snapshots.py +++ b/src/prime_rl/trainer/pretrain/snapshots.py @@ -5,7 +5,6 @@ import torch from torch import nn from torch.distributed.checkpoint.state_dict_saver import save as dcp_save -from torch.utils._pytree import tree_map_only from prime_rl.trainer.ckpt import AppState from prime_rl.trainer.utils import get_dcp_group @@ -30,9 +29,19 @@ def delete_snapshots_up_to(run_dir: Path, step: int) -> None: shutil.rmtree(path, ignore_errors=True) +def snapshot_tensor(name: str, tensor: torch.Tensor) -> torch.Tensor: + """A copy of one model tensor for the snapshot, in bf16 (the dtype the forward computes in, so scoring is + unchanged and the snapshot is half the size), except the router (gate weights and selection bias), which + runs in fp32 and decides the expert selection.""" + if tensor.is_floating_point() and "router" not in name: + return tensor.to(torch.bfloat16, copy=True) + return tensor.clone() + + class SnapshotWriter: """Model-only DCP snapshots for the eval worker (`prime_rl.trainer.pretrain.eval_worker`), in the layout of a - model-only checkpoint (`{"app": AppState}`). The trainer only waits for a copy of its weight shards on the GPU; + model-only checkpoint (`{"app": AppState}`), in bf16 except the router (`snapshot_tensor`); loading casts back + to the model's dtypes. The trainer only waits for a copy of its weight shards on the GPU; a background thread writes them. A snapshot is complete once its `.finished` marker exists; the eval worker deletes it after scoring it.""" @@ -45,7 +54,8 @@ def __init__(self, run_dir: Path): def save(self, model: nn.Module, step: int) -> None: self.wait() - state = tree_map_only(torch.Tensor, torch.Tensor.clone, AppState(model, [], None, None).state_dict()) + state = AppState(model, [], None, None).state_dict() + state["model"] = {name: snapshot_tensor(name, tensor) for name, tensor in state["model"].items()} self.pending = self.executor.submit(self.write, state, snapshot_dir(self.run_dir) / f"step_{step}") def write(self, state: dict, path: Path) -> None: From 993df687c75e78399e04896cec1e4e5f3dc1d2e4 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:34:33 +0200 Subject: [PATCH 027/132] refactor(pretrain): MoE load balancing in [model.moe.balance] The balancing knobs were spread over the glm4_moe HF config (load_balance_coeff, seq_aux_loss_coeff, which also ended up in exported config.json files) and a top-level moe_bias_freeze. They are now [model.moe.balance] in the pretrain model config: bias_update_rate (gamma, 3e-3), bias_freeze (0.1) and seq_aux_coeff (alpha, 1e-3, per packed row). The MoE layer only keeps its routing when the loop asks for it (keep_routing), and update_selection_biases takes the rate, so every MoE family with a selection bias is balanced the same way. glm4_moe goes back to main's config. --- docs/pretraining.md | 12 ++++++++ .../src/prime_rl/configs/pretrain.py | 28 ++++++++++++++++--- src/prime_rl/trainer/model.py | 8 +++--- .../models/glm4_moe/configuration_glm4_moe.py | 8 ------ .../models/glm4_moe/modeling_glm4_moe.py | 3 +- src/prime_rl/trainer/models/layers/moe.py | 11 +++----- src/prime_rl/trainer/pretrain/train.py | 21 ++++++++------ tests/unit/train/models/test_moe.py | 6 ++-- 8 files changed, 61 insertions(+), 36 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index ba0ba23d529..b6cfe2571f1 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -75,6 +75,18 @@ decay_shape = "1-sqrt" # or "linear" (default), "cosine" `f(0) = 0`, the warmup reaches 1 at its last update (`warmup - 1`; a warmup of 0 or 1 starts at 1), and the decay reaches `end_ratio` at the last update (`max_steps - 1`). AdamW and Muon apply weight decay as `lr * weight_decay`, so the schedule also scales the decay. The multiplier is recomputed from the step, so checkpoints hold no schedule state. Resuming with `ckpt.skip_scheduler = true` restarts the schedule at the resumed step, spanning the remaining steps. The multiplier is logged as `schedule/lr`. +## MoE load balancing + +`[model.moe.balance]` balances the routed experts as in DeepSeek-V3, for every MoE family whose router has a selection bias (e.g. `glm4_moe`): + +| Field | Meaning | +|---|---| +| `bias_update_rate` | Every step, each layer's selection bias moves by this step size (gamma) towards underloaded experts, using the step's global expert counts. The update is mean-centred. Default 3e-3; 0 turns it off. The bias starts at 0. | +| `bias_freeze` | Fraction of `max_steps` at the end during which the biases stop updating. Default 0.1. | +| `seq_aux_coeff` | Weight (alpha) of the sequence-wise balance loss, computed per packed row of `data.seq_len` tokens and added to the training loss. Default 1e-3; 0 turns it off. | + +The metrics are `expert_load/*` (global balance of the step), `expert_load/seq_balance` (1 is balanced) and `expert_load/selection_bias_absmax`. + ## Data `data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of a source are read in name order as one stream; documents continue across rows, and attention restarts at every document and every row. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 90868956170..9bf049ed21d 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -9,6 +9,7 @@ from prime_rl.configs.trainer import ( CheckpointConfig, ModelConfig, + MoERuntimeConfig, MuonConfig, OptimizerConfig, TokenizerConfig, @@ -40,6 +41,27 @@ class InitConfig(BaseConfig): """Seed for the initialization. The values depend only on the seed, not on the number of GPUs.""" +class MoEBalanceConfig(BaseConfig): + """MoE load balancing as in DeepSeek-V3: an aux-loss-free selection bias plus a small sequence-wise balance + loss.""" + + bias_update_rate: float = Field(3e-3, ge=0) + """Step size of the selection-bias update (gamma): every step, each layer's bias moves by this towards + underloaded experts, mean-centred. 0 turns the update off.""" + + bias_freeze: float = Field(0.1, ge=0, lt=1) + """Fraction of ``max_steps`` at the end of training during which the biases stop updating, as DeepSeek-V3, + GLM-4.5 and Moonlight do in their final stage.""" + + seq_aux_coeff: float = Field(1e-3, ge=0) + """Weight of the sequence-wise balance loss (alpha), computed per packed row. 0 turns it off.""" + + +class PretrainMoEConfig(MoERuntimeConfig): + balance: MoEBalanceConfig = MoEBalanceConfig() + """Load balancing of the routed experts.""" + + class PretrainModelConfig(ModelConfig): arch: dict[str, Any] """The architecture: ``model_type`` (a model in ``prime_rl.trainer.models``, e.g. ``glm4_moe``) and the fields of its config class. ``vocab_size`` and the special token ids come from the tokenizer. The launcher writes the config and the tokenizer to ``/model``, which ``name`` points at.""" @@ -47,6 +69,8 @@ class PretrainModelConfig(ModelConfig): init: InitConfig = InitConfig() """Random weight initialization. Used unless the run resumes from a checkpoint.""" + moe: PretrainMoEConfig = PretrainMoEConfig() + @model_validator(mode="after") def validate_arch(self): if "model_type" not in self.arch: @@ -212,10 +236,6 @@ class PretrainConfig(TrainRunConfig): max_steps: int = Field(ge=1) """Training steps.""" - moe_bias_freeze: float = Field(0.1, ge=0, lt=1) - """Fraction of ``max_steps`` at the end of training during which the MoE selection biases stop updating, as - DeepSeek-V3, GLM-4.5 and Moonlight do in their final stage.""" - def default_run_name(self) -> str: """``--``.""" return f"{self.model.arch['model_type']}--{uuid.uuid4().hex[:8]}" diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 78c4a5cba58..4ff32e3bb5d 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -330,15 +330,15 @@ def get_expert_load_stats(tokens_per_expert: Tensor, group: dist.ProcessGroup) - @torch.no_grad() -def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor) -> None: +def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor, rate: float) -> None: """Aux-loss-free load balancing (https://arxiv.org/abs/2408.15664): move each layer's selection bias by - `load_balance_coeff` towards underloaded experts, given the step's global `[num_moe_layers, num_experts]` counts. + ``rate`` towards underloaded experts, given the step's global ``[num_moe_layers, num_experts]`` counts. The update is mean-centred (as in Moonlight and torchtitan), which leaves the selection unchanged but keeps the biases from drifting together.""" for moe, load in zip(iter_moe_blocks(model), tokens_per_expert): - if moe.load_balance_coeff is not None: + if moe.router.selection_bias is not None: delta = torch.sign(load.mean() - load) - moe.router.selection_bias += moe.load_balance_coeff * (delta - delta.mean()) + moe.router.selection_bias += rate * (delta - delta.mean()) def get_model( diff --git a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py index 5c9c19cc0b1..30dce92cd34 100644 --- a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py @@ -104,10 +104,6 @@ class Glm4MoeConfig(PretrainedConfig): Whether to normalize the topk probabilities. use_qk_norm (`bool`, *optional*, defaults to `False`): Whether to use query-key normalization in the attention - load_balance_coeff (`float`, *optional*, defaults to 1e-3): - Update speed of the aux-loss-free load-balancing bias; `None` disables the bias. - seq_aux_loss_coeff (`float`, *optional*): - Weight of the sequence-wise balance loss (DeepSeek-V3); `None` disables it, 0 only logs it. ```python >>> from transformers import Glm4MoeModel, Glm4MoeConfig >>> # Initializing a Glm4Moe style configuration @@ -169,8 +165,6 @@ def __init__( first_k_dense_replace=1, norm_topk_prob=True, use_qk_norm=False, - load_balance_coeff=1e-3, - seq_aux_loss_coeff=None, pad_token_id=None, **kwargs, ): @@ -203,8 +197,6 @@ def __init__( self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob self.use_qk_norm = use_qk_norm - self.load_balance_coeff = load_balance_coeff - self.seq_aux_loss_coeff = seq_aux_loss_coeff super().__init__( tie_word_embeddings=tie_word_embeddings, diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py index f50f329c1ef..fcc0af1a201 100644 --- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py @@ -61,8 +61,7 @@ def __init__(self, config: Glm4MoeConfig, layer_idx: int): route_scale=config.routed_scaling_factor, score_before_experts=False, top_k=config.num_experts_per_tok, - load_balance_coeff=config.load_balance_coeff, - seq_aux_loss_coeff=config.seq_aux_loss_coeff, + load_balance_coeff=1e-3, ) if layer_idx >= config.first_k_dense_replace: shared_expert = None diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index c8195cf1320..551713de8cf 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -61,7 +61,6 @@ class MoEArgs: # token-choice top_k: int = 1 load_balance_coeff: float | None = 1e-3 - seq_aux_loss_coeff: float | None = None def __post_init__(self) -> None: ActivationDispatch[self.activation] @@ -330,7 +329,6 @@ def from_args( shared_expert=shared_expert, score_before_experts=args.score_before_experts, load_balance_coeff=args.load_balance_coeff, - seq_aux_loss_coeff=args.seq_aux_loss_coeff, ) def __init__( @@ -341,7 +339,6 @@ def __init__( shared_expert: FeedForward | None, score_before_experts: bool, load_balance_coeff: float | None, - seq_aux_loss_coeff: float | None = None, ) -> None: super().__init__() self.router = router @@ -368,9 +365,9 @@ def __init__( persistent=False, ) self.register_buffer("routing_confidence_sum", torch.tensor(0.0, dtype=torch.float32), persistent=False) - # Sequence-wise balance loss (DeepSeek-V3 eq. 17-20): with a coefficient set, every forward with grad - # keeps the scores and selections, for the training loop, which knows the packed rows, to compute it. - self.seq_aux_loss_coeff = seq_aux_loss_coeff + # Sequence-wise balance loss (DeepSeek-V3 eq. 17-20): with keep_routing, every forward with grad keeps + # the scores and selections, for the training loop, which knows the packed rows, to compute it. + self.keep_routing = False self.routing: tuple[torch.Tensor, torch.Tensor] | None = None def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None: @@ -414,7 +411,7 @@ def forward( scores, ) = self.router(x, routed_experts=routed_experts) - if self.seq_aux_loss_coeff is not None and torch.is_grad_enabled(): + if self.keep_routing and torch.is_grad_enabled(): self.routing = (scores, selected_experts_indices) # Accumulate expert usage for selection-bias updates and metrics. diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 1c14e9e22ce..1e4fedfa0fa 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -202,8 +202,11 @@ def train(config: PretrainConfig): cp_size = parallel_dims.cp dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() is_moe_model = is_tt_moe_model(model) - moes_with_seq_balance = [moe for moe in iter_moe_blocks(model) if moe.seq_aux_loss_coeff is not None] - bias_freeze_step = config.max_steps - round(config.moe_bias_freeze * config.max_steps) + balance = config.model.moe.balance + moes_with_seq_balance = list(iter_moe_blocks(model)) if balance.seq_aux_coeff > 0 else [] + for moe in moes_with_seq_balance: + moe.keep_routing = True + bias_freeze_step = config.max_steps - round(balance.bias_freeze * config.max_steps) # Per-row balance of every MoE layer and micro-batch of the current step, for the metrics. step_seq_balances: list[torch.Tensor] = [] @@ -232,9 +235,9 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.T # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. num_rows = config.data.micro_batch_size for moe in moes_with_seq_balance: - balance = sequence_balance(*moe.routing, num_rows) - aux_loss_sum += moe.seq_aux_loss_coeff * balance.sum() * (input_ids.shape[1] // num_rows) - step_seq_balances.append(balance.detach().mean()) + row_balance = sequence_balance(*moe.routing, num_rows) + aux_loss_sum += balance.seq_aux_coeff * row_balance.sum() * (input_ids.shape[1] // num_rows) + step_seq_balances.append(row_balance.detach().mean()) return out["loss"], loss_mask.sum(dtype=torch.int64), aux_loss_sum @torch.no_grad() @@ -285,9 +288,11 @@ def run_evals(step: int) -> None: if step_seq_balances: expert_load_stats["expert_load/seq_balance"] = torch.stack(step_seq_balances).mean().item() step_seq_balances.clear() - if progress.step <= bias_freeze_step: - update_selection_biases(model, result.tokens_per_expert) - biases = [moe.router.selection_bias for moe in iter_moe_blocks(model) if moe.load_balance_coeff is not None] + if balance.bias_update_rate > 0 and progress.step <= bias_freeze_step: + update_selection_biases(model, result.tokens_per_expert, balance.bias_update_rate) + biases = [ + moe.router.selection_bias for moe in iter_moe_blocks(model) if moe.router.selection_bias is not None + ] if biases: expert_load_stats["expert_load/selection_bias_absmax"] = torch.stack(biases).abs().max().item() diff --git a/tests/unit/train/models/test_moe.py b/tests/unit/train/models/test_moe.py index 3a3341cd0b9..9ad8b10f0af 100644 --- a/tests/unit/train/models/test_moe.py +++ b/tests/unit/train/models/test_moe.py @@ -255,15 +255,15 @@ def test_expert_load_stats(): def test_update_selection_biases_moves_towards_underloaded_experts(): - moe = MoE.from_args(MoEArgs(num_experts=4, load_balance_coeff=0.1), dim=4, hidden_dim=8, shared_expert=None) + moe = MoE.from_args(MoEArgs(num_experts=4), dim=4, hidden_dim=8, shared_expert=None) layer = torch.nn.Module() layer.mlp = moe model = torch.nn.Module() model.model = torch.nn.Module() model.model.layers = torch.nn.ModuleList([layer]) - update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]])) - update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]])) + update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]]), rate=0.1) + update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]]), rate=0.1) # sign(mean - load) = [-1, 1, 1, 1], centred to [-1.5, 0.5, 0.5, 0.5] per step torch.testing.assert_close(moe.router.selection_bias, torch.tensor([-0.3, 0.1, 0.1, 0.1])) From 5f4f1f955f0a3eb0f83a3d036e3f4756e77daf81 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:35:05 +0200 Subject: [PATCH 028/132] refactor(pretrain): drop the scaled init scheme and the truncated normal No experiment uses model.init.scheme = 'scaled', and the truncated normal was only used by the Marin init arm, which did not beat the plain normal init. model.init.scheme is normal or completep. --- docs/pretraining.md | 3 +-- .../src/prime_rl/configs/pretrain.py | 7 ++----- src/prime_rl/trainer/init.py | 15 ++------------- tests/unit/train/pretrain/test_init.py | 7 ------- 4 files changed, 5 insertions(+), 27 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index b6cfe2571f1..e1e97980762 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -30,9 +30,8 @@ Weights start from a random init, unless the run resumes from a checkpoint. | Field | Meaning | |---|---| -| `model.init.scheme` | `normal`: every weight matrix and the embedding from N(0, std). `scaled`: additionally divides the std of the residual output projections (`o_proj`, `down_proj`) by sqrt(2 * num_layers). `completep`: CompleteP, see below. Norm weights start at 1, biases at 0. | +| `model.init.scheme` | `normal`: every weight matrix and the embedding from N(0, std). `completep`: CompleteP, see below. Norm weights start at 1, biases at 0. | | `model.init.std` | Base standard deviation (default 0.02). | -| `model.init.truncate` | Draw from the normal truncated at +-3 std. | | `model.init.seed` | Seed. Every rank initializes only its own shard, and the values depend only on the seed and the parameter name, not on the number of GPUs. | ### CompleteP diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 9bf049ed21d..e1fd0a6727a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -22,15 +22,12 @@ class InitConfig(BaseConfig): """Random weight initialization.""" - scheme: Literal["normal", "scaled", "completep"] = "normal" - """``normal`` draws every weight matrix and the embedding from N(0, std). ``scaled`` additionally divides the std of the residual output projections (``o_proj``, ``down_proj``) by sqrt(2 * num_layers). ``completep`` follows CompleteP (muP extended to depth, Dey et al. 2025) relative to the base shape: the embedding keeps ``std``, hidden matrices get ``std / sqrt(m_N)``, residual output projections ``std / (sqrt(m_N) * m_L)`` and the LM head ``std / m_N``, where ``m_N = hidden_size / base_hidden_size`` and ``m_L = num_layers / base_num_layers``. Norm weights start at 1, biases at 0.""" + scheme: Literal["normal", "completep"] = "normal" + """``normal`` draws every weight matrix and the embedding from N(0, std). ``completep`` follows CompleteP (muP extended to depth, Dey et al. 2025) relative to the base shape: the embedding keeps ``std``, hidden matrices get ``std / sqrt(m_N)``, residual output projections ``std / (sqrt(m_N) * m_L)`` and the LM head ``std / m_N``, where ``m_N = hidden_size / base_hidden_size`` and ``m_L = num_layers / base_num_layers``. Norm weights start at 1, biases at 0.""" std: float = Field(0.02, gt=0) """Base standard deviation.""" - truncate: bool = False - """Draw from the normal distribution truncated at +-3 std instead of the full normal.""" - base_hidden_size: int = Field(256, ge=1) """Hidden size of the base shape. Only used by ``completep``.""" diff --git a/src/prime_rl/trainer/init.py b/src/prime_rl/trainer/init.py index e2dfa64c431..e603a6edd8a 100644 --- a/src/prime_rl/trainer/init.py +++ b/src/prime_rl/trainer/init.py @@ -47,8 +47,6 @@ def role_std(role: ParamRole, config: InitConfig, hidden_size: int, num_layers: m_n, m_l = completep_multipliers(config, hidden_size, num_layers) scale = {"embedding": 1.0, "hidden": m_n**-0.5, "residual_out": m_n**-0.5 / m_l, "output": 1 / m_n} return config.std * scale[role] - if config.scheme == "scaled" and role == "residual_out": - return config.std / math.sqrt(2 * num_layers) return config.std @@ -64,12 +62,10 @@ def normal_shard( global_offset: tuple[int, ...], seed: int, device: torch.device, - truncate: bool = False, ) -> torch.Tensor: """The ``local_shape`` slice at ``global_offset`` of a standard-normal tensor of ``global_shape`` whose values depend only on ``(seed, name)``. The tensor is drawn in - blocks of rows along dim 0; only the blocks overlapping the slice are generated. With - ``truncate``, the distribution is truncated at +-3.""" + blocks of rows along dim 0; only the blocks overlapping the slice are generated.""" row_numel = math.prod(global_shape[1:]) rows_per_block = max(1, BLOCK_NUMEL // max(row_numel, 1)) start, stop = global_offset[0], global_offset[0] + local_shape[0] @@ -78,13 +74,7 @@ def normal_shard( block_start = block * rows_per_block block_rows = min(rows_per_block, global_shape[0] - block_start) generator = torch.Generator(device=device).manual_seed(block_seed(seed, name, block)) - if truncate: - # Inverse-CDF sampling: uniform on (erf(-3/sqrt(2)), erf(3/sqrt(2))) mapped through erfinv. - bound = math.erf(3 / math.sqrt(2)) - uniform = torch.rand((block_rows, *global_shape[1:]), generator=generator, device=device) - values = torch.erfinv(uniform.mul_(2 * bound).sub_(bound)).mul_(math.sqrt(2)) - else: - values = torch.randn((block_rows, *global_shape[1:]), generator=generator, device=device) + values = torch.randn((block_rows, *global_shape[1:]), generator=generator, device=device) lo, hi = max(start, block_start) - block_start, min(stop, block_start + block_rows) - block_start pieces.append(values[lo:hi]) out = torch.cat(pieces) if pieces else torch.empty((0, *global_shape[1:]), device=device) @@ -122,6 +112,5 @@ def init_weights(model: nn.Module, config: InitConfig) -> None: tuple(global_offset), config.seed, local.device, - config.truncate, ) local.copy_(values.mul_(role_std(role, config, hidden_size, num_layers))) diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index 089133a887e..81095f25f61 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -28,13 +28,6 @@ def test_normal_shard_is_independent_of_sharding(): assert not torch.equal(normal_shard("w", shape, shape, (0, 0, 0), seed=1, device=torch.device("cpu")), full) -def test_truncated_normal_shard_stays_within_three_std(): - values = normal_shard("w", (4096, 64), (4096, 64), (0, 0), seed=0, device=torch.device("cpu"), truncate=True) - assert values.abs().max() <= 3 - # Variance of the standard normal truncated at +-3. - assert abs(values.std().item() - 0.9866) < 0.01 - - def test_completep_folded_multipliers_match_explicit_multipliers(): """CompleteP folds its 1/m_L residual and 1/m_N logit multipliers into the weights. Training the folded model must match the multiplier form of CompleteP (Dey et al. 2025, Table 1).""" From 0bc6bc36f18680931916158098b6cf4eeb9bdda3 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:43:09 +0200 Subject: [PATCH 029/132] refactor(moe): the MoE layer computes the per-row balance itself With balance_rows set, each forward with grad stores only the [rows] sequence-wise balance instead of every layer's fp32 router scores and selections, which the loop used to read back after the forward. --- src/prime_rl/trainer/models/layers/moe.py | 12 ++++++------ src/prime_rl/trainer/pretrain/train.py | 12 +++++------- 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index 551713de8cf..11e90bfd4ce 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -365,10 +365,10 @@ def __init__( persistent=False, ) self.register_buffer("routing_confidence_sum", torch.tensor(0.0, dtype=torch.float32), persistent=False) - # Sequence-wise balance loss (DeepSeek-V3 eq. 17-20): with keep_routing, every forward with grad keeps - # the scores and selections, for the training loop, which knows the packed rows, to compute it. - self.keep_routing = False - self.routing: tuple[torch.Tensor, torch.Tensor] | None = None + # Sequence-wise balance (DeepSeek-V3 eq. 17-20): with balance_rows set by the training loop, which knows + # how many packed rows a sample holds, every forward with grad computes the per-row balance for its loss. + self.balance_rows: int | None = None + self.seq_balance: torch.Tensor | None = None def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None: self.token_dispatcher = token_dispatcher @@ -411,8 +411,8 @@ def forward( scores, ) = self.router(x, routed_experts=routed_experts) - if self.keep_routing and torch.is_grad_enabled(): - self.routing = (scores, selected_experts_indices) + if self.balance_rows is not None and torch.is_grad_enabled(): + self.seq_balance = sequence_balance(scores, selected_experts_indices, self.balance_rows) # Accumulate expert usage for selection-bias updates and metrics. with torch.no_grad(): diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 1e4fedfa0fa..d8758f8bbc7 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -29,7 +29,6 @@ update_selection_biases, ) from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX -from prime_rl.trainer.models.layers.moe import sequence_balance from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter @@ -205,7 +204,8 @@ def train(config: PretrainConfig): balance = config.model.moe.balance moes_with_seq_balance = list(iter_moe_blocks(model)) if balance.seq_aux_coeff > 0 else [] for moe in moes_with_seq_balance: - moe.keep_routing = True + # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. + moe.balance_rows = config.data.micro_batch_size bias_freeze_step = config.max_steps - round(balance.bias_freeze * config.max_steps) # Per-row balance of every MoE layer and micro-batch of the current step, for the metrics. step_seq_balances: list[torch.Tensor] = [] @@ -232,12 +232,10 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.T ) aux_loss_sum = torch.zeros((), device="cuda") if torch.is_grad_enabled(): - # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. - num_rows = config.data.micro_batch_size + row_len = input_ids.shape[1] // config.data.micro_batch_size for moe in moes_with_seq_balance: - row_balance = sequence_balance(*moe.routing, num_rows) - aux_loss_sum += balance.seq_aux_coeff * row_balance.sum() * (input_ids.shape[1] // num_rows) - step_seq_balances.append(row_balance.detach().mean()) + aux_loss_sum += balance.seq_aux_coeff * moe.seq_balance.sum() * row_len + step_seq_balances.append(moe.seq_balance.detach().mean()) return out["loss"], loss_mask.sum(dtype=torch.int64), aux_loss_sum @torch.no_grad() From 4de892c7e489bec731dd069ce1aeb5b977ba7586 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:44:36 +0200 Subject: [PATCH 030/132] refactor(pretrain): one per-role parametrization; [parametrization] turns on CompleteP model.init.scheme switched between the shared default optimizer groups and CompleteP's per-role groups, so the SP and CompleteP arms also differed in weight decay on norms. Now both use the per-role groups: the standard parametrization is CompleteP with every multiplier at 1, which equals the old SP except that norms and biases get no weight decay. [parametrization] (base_hidden_size, base_num_layers) turns on CompleteP; model.init keeps std and seed. --- docs/pretraining.md | 11 +++--- .../src/prime_rl/configs/pretrain.py | 34 +++++++++++-------- src/prime_rl/trainer/init.py | 28 ++++++--------- .../{completep.py => parametrization.py} | 22 ++++++++---- src/prime_rl/trainer/pretrain/train.py | 11 +++--- tests/unit/train/pretrain/test_init.py | 10 +++--- 6 files changed, 59 insertions(+), 57 deletions(-) rename src/prime_rl/trainer/pretrain/{completep.py => parametrization.py} (76%) diff --git a/docs/pretraining.md b/docs/pretraining.md index e1e97980762..e82ba55eadf 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -30,13 +30,14 @@ Weights start from a random init, unless the run resumes from a checkpoint. | Field | Meaning | |---|---| -| `model.init.scheme` | `normal`: every weight matrix and the embedding from N(0, std). `completep`: CompleteP, see below. Norm weights start at 1, biases at 0. | -| `model.init.std` | Base standard deviation (default 0.02). | +| `model.init.std` | Every weight matrix and the embedding start from N(0, std) (default 0.02), scaled per role under CompleteP (below). Norm weights start at 1, biases at 0. | | `model.init.seed` | Seed. Every rank initializes only its own shard, and the values depend only on the seed and the parameter name, not on the number of GPUs. | -### CompleteP +## Parametrization -`model.init.scheme = "completep"` trains with [CompleteP](https://arxiv.org/abs/2505.01618) (muP extended to depth, alpha = 1), so that hyperparameters tuned on a small base shape carry over to wider and deeper models. `optim.lr`, `optim.eps`, `optim.weight_decay` and `model.init.std` are the base-shape values, and the shape is set by `model.init.base_hidden_size` (default 256) and `model.init.base_num_layers` (default 2). With `m_N = hidden_size / base_hidden_size` and `m_L = num_layers / base_num_layers`: +The optimizer gets one parameter group per role (embedding, hidden matrices, `o_proj`/`down_proj`, LM head, norms, biases; with MoE also experts and the router). By default (standard parametrization) every role uses `optim.lr`, `optim.eps` and `optim.weight_decay`, except norms and biases, which get no weight decay. + +`[parametrization]` trains with [CompleteP](https://arxiv.org/abs/2505.01618) (muP extended to depth, alpha = 1), so that hyperparameters tuned on a small base shape carry over to wider and deeper models. `optim.lr`, `optim.eps`, `optim.weight_decay` and `model.init.std` are the base-shape values, and the base shape is set by `parametrization.base_hidden_size` and `parametrization.base_num_layers`. With `m_N = hidden_size / base_hidden_size` and `m_L = num_layers / base_num_layers`: | Parameters | Init std | Learning rate | Adam eps | Weight decay | |---|---|---|---|---| @@ -58,7 +59,7 @@ With `optim.type = "muon"`, the hidden matrices and `o_proj`/`down_proj` are tra In MoE models, the expert matrices follow the hidden and `down_proj` rules per expert. The router maps the width to a fixed number of experts, so under Muon its learning rate falls as 1/m_N (with AdamW it follows the hidden rule). -CompleteP needs `model.offload = "none"`. Keep `head_dim` and the ratios `intermediate_size / hidden_size` (and `moe_intermediate_size / hidden_size`) and `num_attention_heads / num_key_value_heads` fixed across shapes. +The per-role groups need AdamW or Muon and `model.offload = "none"`. The standard parametrization is CompleteP with the base shape equal to the model shape. Keep `head_dim` and the ratios `intermediate_size / hidden_size` (and `moe_intermediate_size / hidden_size`) and `num_attention_heads / num_key_value_heads` fixed across shapes. ## Learning-rate schedule diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index e1fd0a6727a..afbbefc232f 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -20,24 +20,28 @@ class InitConfig(BaseConfig): - """Random weight initialization.""" - - scheme: Literal["normal", "completep"] = "normal" - """``normal`` draws every weight matrix and the embedding from N(0, std). ``completep`` follows CompleteP (muP extended to depth, Dey et al. 2025) relative to the base shape: the embedding keeps ``std``, hidden matrices get ``std / sqrt(m_N)``, residual output projections ``std / (sqrt(m_N) * m_L)`` and the LM head ``std / m_N``, where ``m_N = hidden_size / base_hidden_size`` and ``m_L = num_layers / base_num_layers``. Norm weights start at 1, biases at 0.""" + """Random weight initialization: every weight matrix and the embedding from N(0, std), scaled per role under + CompleteP (see ``[parametrization]``). Norm weights start at 1, biases at 0.""" std: float = Field(0.02, gt=0) """Base standard deviation.""" - base_hidden_size: int = Field(256, ge=1) - """Hidden size of the base shape. Only used by ``completep``.""" - - base_num_layers: int = Field(2, ge=1) - """Number of layers of the base shape. Only used by ``completep``.""" - seed: int = 0 """Seed for the initialization. The values depend only on the seed, not on the number of GPUs.""" +class ParametrizationConfig(BaseConfig): + """CompleteP (muP extended to depth, Dey et al. 2025, alpha = 1) relative to a base shape: init std, learning + rate, Adam eps and weight decay scale per parameter role with ``m_N = hidden_size / base_hidden_size`` and + ``m_L = num_layers / base_num_layers``, so hyperparameters tuned on the base shape carry over.""" + + base_hidden_size: int = Field(ge=1) + """Hidden size of the base shape.""" + + base_num_layers: int = Field(ge=1) + """Number of layers of the base shape.""" + + class MoEBalanceConfig(BaseConfig): """MoE load balancing as in DeepSeek-V3: an aux-loss-free selection bias plus a small sequence-wise balance loss.""" @@ -233,6 +237,10 @@ class PretrainConfig(TrainRunConfig): max_steps: int = Field(ge=1) """Training steps.""" + parametrization: ParametrizationConfig | None = None + """CompleteP relative to a base shape. If None, the standard parametrization: the same per-role parameter + groups with every multiplier at 1 (norms and biases without weight decay).""" + def default_run_name(self) -> str: """``--``.""" return f"{self.model.arch['model_type']}--{uuid.uuid4().hex[:8]}" @@ -267,11 +275,9 @@ def validate_tokenizer(self): return self @model_validator(mode="after") - def validate_completep(self): - if self.model.init.scheme != "completep": - return self + def validate_param_groups(self): if self.optim.type not in ("adamw", "muon") or self.model.offload != "none": - raise ValueError("model.init.scheme = 'completep' needs AdamW or Muon and model.offload = 'none'.") + raise ValueError("Pretraining's per-role parameter groups need AdamW or Muon and model.offload = 'none'.") return self @model_validator(mode="after") diff --git a/src/prime_rl/trainer/init.py b/src/prime_rl/trainer/init.py index e603a6edd8a..4ee1608facb 100644 --- a/src/prime_rl/trainer/init.py +++ b/src/prime_rl/trainer/init.py @@ -35,19 +35,12 @@ def param_role(name: str, param: torch.Tensor) -> ParamRole: return "hidden" -def completep_multipliers(config: InitConfig, hidden_size: int, num_layers: int) -> tuple[float, float]: - """Width and depth multipliers ``(m_N, m_L)`` of the model relative to the CompleteP base shape.""" - return hidden_size / config.base_hidden_size, num_layers / config.base_num_layers - - -def role_std(role: ParamRole, config: InitConfig, hidden_size: int, num_layers: int) -> float: - if config.scheme == "completep": - # CompleteP's 1/m_L residual-branch and 1/m_N logit multipliers are folded into the - # residual output projections and the LM head, so the forward pass stays unchanged. - m_n, m_l = completep_multipliers(config, hidden_size, num_layers) - scale = {"embedding": 1.0, "hidden": m_n**-0.5, "residual_out": m_n**-0.5 / m_l, "output": 1 / m_n} - return config.std * scale[role] - return config.std +def role_std(role: ParamRole, std: float, m_n: float, m_l: float) -> float: + """Init std of a role under CompleteP with width and depth multipliers ``m_N``, ``m_L`` (both 1 in the + standard parametrization). CompleteP's 1/m_L residual-branch and 1/m_N logit multipliers are folded into the + residual output projections and the LM head, so the forward pass stays unchanged.""" + scale = {"embedding": 1.0, "hidden": m_n**-0.5, "residual_out": m_n**-0.5 / m_l, "output": 1 / m_n} + return std * scale[role] def block_seed(seed: int, name: str, block: int) -> int: @@ -84,10 +77,9 @@ def normal_shard( @torch.no_grad() -def init_weights(model: nn.Module, config: InitConfig) -> None: - """Initialize every parameter in place. Each rank fills only its local shard, without - communication, and the values do not depend on the world size or the sharding.""" - hidden_size, num_layers = model.config.hidden_size, model.config.num_hidden_layers +def init_weights(model: nn.Module, config: InitConfig, m_n: float = 1.0, m_l: float = 1.0) -> None: + """Initialize every parameter in place, with the CompleteP multipliers ``m_N``, ``m_L``. Each rank fills only + its local shard, without communication, and the values do not depend on the world size or the sharding.""" for name, param in model.named_parameters(): # Canonical names, so wrappers such as activation checkpointing don't change the values. name = resolve_fqn(model, name) @@ -113,4 +105,4 @@ def init_weights(model: nn.Module, config: InitConfig) -> None: config.seed, local.device, ) - local.copy_(values.mul_(role_std(role, config, hidden_size, num_layers))) + local.copy_(values.mul_(role_std(role, config.std, m_n, m_l))) diff --git a/src/prime_rl/trainer/pretrain/completep.py b/src/prime_rl/trainer/pretrain/parametrization.py similarity index 76% rename from src/prime_rl/trainer/pretrain/completep.py rename to src/prime_rl/trainer/pretrain/parametrization.py index 6a2b94af1e7..345d6295b08 100644 --- a/src/prime_rl/trainer/pretrain/completep.py +++ b/src/prime_rl/trainer/pretrain/parametrization.py @@ -2,18 +2,27 @@ from torch import nn -from prime_rl.configs.pretrain import InitConfig +from prime_rl.configs.pretrain import ParametrizationConfig from prime_rl.configs.trainer import AdamWConfig, MuonConfig -from prime_rl.trainer.init import completep_multipliers, param_role +from prime_rl.trainer.init import param_role from prime_rl.utils.weights import resolve_fqn -def completep_param_groups( - model: nn.Module, init: InitConfig, optim: AdamWConfig | MuonConfig, ep_enabled: bool +def multipliers(config: ParametrizationConfig | None, model: nn.Module) -> tuple[float, float]: + """Width and depth multipliers ``(m_N, m_L)`` of the model relative to the CompleteP base shape; 1 and 1 in + the standard parametrization.""" + if config is None: + return 1.0, 1.0 + return model.config.hidden_size / config.base_hidden_size, model.config.num_hidden_layers / config.base_num_layers + + +def param_groups( + model: nn.Module, optim: AdamWConfig | MuonConfig, m_n: float, m_l: float, ep_enabled: bool ) -> list[dict]: - """Optimizer parameter groups for CompleteP with alpha = 1 (Dey et al. 2025). ``optim.lr``, + """Optimizer parameter groups per role for CompleteP with alpha = 1 (Dey et al. 2025). ``optim.lr``, ``optim.eps`` and ``optim.weight_decay`` are the values for the base shape; each role scales - them by the width and depth multipliers. The residual output projections and the LM head + them by the width and depth multipliers. With both multipliers at 1 this is the standard + parametrization, where only norms and biases differ (no weight decay). The residual output projections and the LM head carry CompleteP's 1/m_L and 1/m_N forward multipliers folded into their weights (see ``role_std``), which turns their rules into the ones below. @@ -27,7 +36,6 @@ def completep_param_groups( sqrt(num_experts / d) shrinks as m_N^-1/2: with Muon its learning rate is scaled by 1/m_N; with AdamW it follows the hidden rule. """ - m_n, m_l = completep_multipliers(init, model.config.hidden_size, model.config.num_hidden_layers) # (lr, eps, weight decay) scale per role for AdamW. PyTorch AdamW and Dion decay by lr * weight_decay. adam_scales = { "embedding": (1.0, 1 / m_n, 1.0), diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index d8758f8bbc7..83ef80f3588 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -33,7 +33,7 @@ from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter from prime_rl.trainer.init import init_weights -from prime_rl.trainer.pretrain.completep import completep_param_groups +from prime_rl.trainer.pretrain.parametrization import multipliers, param_groups from prime_rl.trainer.pretrain.bytes import token_byte_lengths from prime_rl.trainer.pretrain.data import ( TokenShards, @@ -121,8 +121,9 @@ def train(config: PretrainConfig): logger.info(f"Initializing model ({config.model})") # Weights come from the checkpoint when resuming and from the random init otherwise. model = setup_model(config.model, parallel_dims, loading_from_checkpoint_later=True) + m_n, m_l = multipliers(config.parametrization, model) if checkpoint_step is None: - init_weights(model, config.model.init) + init_weights(model, config.model.init, m_n, m_l) if parallel_dims.cp_enabled: setup_context_parallel(model, config.model, parallel_dims) num_params = count_parameters(model) @@ -142,11 +143,7 @@ def train(config: PretrainConfig): full_offload_dtype_policy=( get_full_offload_dtype_policy(model, config.model) if config.model.offload == "full" else None ), - param_groups=( - completep_param_groups(model, config.model.init, config.optim, parallel_dims.ep_enabled) - if config.model.init.scheme == "completep" - else None - ), + param_groups=param_groups(model, config.optim, m_n, m_l, parallel_dims.ep_enabled), ) # With ckpt.skip_scheduler, the schedule restarts at the resumed step and spans the remaining steps. diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index 81095f25f61..9d0fef5588a 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -9,7 +9,7 @@ from prime_rl.configs.trainer import AdamWConfig, MuonConfig from prime_rl.trainer.init import init_weights, normal_shard from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM -from prime_rl.trainer.pretrain.completep import completep_param_groups +from prime_rl.trainer.pretrain.parametrization import param_groups def test_normal_shard_is_independent_of_sharding(): @@ -42,12 +42,11 @@ def test_completep_folded_multipliers_match_explicit_multipliers(): head_dim=32, tie_word_embeddings=False, ) - init = InitConfig(scheme="completep", base_hidden_size=32, base_num_layers=2) optim = AdamWConfig(lr=1e-2, eps=1e-10, weight_decay=0.1) m_n, m_l = 2.0, 2.0 folded = LlamaForCausalLM(config) - init_weights(folded, init) + init_weights(folded, InitConfig(), m_n, m_l) explicit = copy.deepcopy(folded) with torch.no_grad(): for name, param in explicit.named_parameters(): @@ -75,7 +74,7 @@ def explicit_scales(name: str) -> tuple[float, float, float]: ) optimizers = { "explicit": AdamW(explicit_groups, betas=(0.9, 0.95)), - "folded": AdamW(completep_param_groups(folded, init, optim, ep_enabled=False), betas=(0.9, 0.95)), + "folded": AdamW(param_groups(folded, optim, m_n, m_l, ep_enabled=False), betas=(0.9, 0.95)), } models = {"explicit": explicit, "folded": folded} @@ -108,8 +107,7 @@ def test_completep_moe_groups_cover_experts_and_scale_the_router(): config._attn_implementation = "flash_attention_2" with torch.device("meta"): model = Glm4MoeForCausalLM._from_config(config) - init = InitConfig(scheme="completep", base_hidden_size=256, base_num_layers=2) - groups = completep_param_groups(model, init, MuonConfig(lr=1.0, weight_decay=0.1), ep_enabled=False) + groups = param_groups(model, MuonConfig(lr=1.0, weight_decay=0.1), m_n=2.0, m_l=1.0, ep_enabled=False) lr = {param: group["lr"] for group in groups for param in group["params"]} assert sum(len(group["params"]) for group in groups) == len(lr) == len(list(model.parameters())) moe = model.model.layers[1].mlp From fbe37d9e68614d264f593382cee02561a9f12562 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:45:09 +0200 Subject: [PATCH 031/132] feat(pretrain): count parameters by role, including active non-embedding parameters count_parameters uses param_role instead of its own name match, and counts the routed experts of each MoE layer at top_k / num_experts for the active non-embedding parameters that MoE scaling laws use. The counts are logged at startup and as params/* at step 0. --- src/prime_rl/trainer/pretrain/train.py | 27 ++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 83ef80f3588..d7527ed9734 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -32,7 +32,7 @@ from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter -from prime_rl.trainer.init import init_weights +from prime_rl.trainer.init import init_weights, param_role from prime_rl.trainer.pretrain.parametrization import multipliers, param_groups from prime_rl.trainer.pretrain.bytes import token_byte_lengths from prime_rl.trainer.pretrain.data import ( @@ -60,16 +60,28 @@ from prime_rl.utils.pathing import resolve_latest_ckpt_step from prime_rl.utils.process import set_proc_title from prime_rl.utils.utils import clean_exit +from prime_rl.utils.weights import resolve_fqn def count_parameters(model: nn.Module) -> dict[str, int]: - """Total, embedding (input embedding and LM head) and non-embedding parameter counts.""" + """Total, embedding (input embedding and LM head), non-embedding and active non-embedding parameter counts. + Active counts the routed experts of each MoE layer at ``top_k / num_experts``, the share each token uses.""" total = embedding = 0 for name, param in model.named_parameters(): total += param.numel() - if "embed_tokens" in name or "lm_head" in name: + if param_role(resolve_fqn(model, name), param) in ("embedding", "output"): embedding += param.numel() - return {"total": total, "embedding": embedding, "non_embedding": total - embedding} + inactive = 0 + for moe in iter_moe_blocks(model): + experts = sum(param.numel() for param in moe.experts.parameters()) + inactive += experts - experts * moe.router.top_k // moe.router.num_experts + non_embedding = total - embedding + return { + "total": total, + "embedding": embedding, + "non_embedding": non_embedding, + "active_non_embedding": non_embedding - inactive, + } @clean_exit @@ -128,9 +140,12 @@ def train(config: PretrainConfig): setup_context_parallel(model, config.model, parallel_dims) num_params = count_parameters(model) logger.info( - f"Parameters: {num_params['non_embedding'] / 1e6:,.1f}M non-embedding, " - f"{num_params['embedding'] / 1e6:,.1f}M embedding, {num_params['total'] / 1e6:,.1f}M total" + f"Parameters: {num_params['active_non_embedding'] / 1e6:,.1f}M active non-embedding, " + f"{num_params['non_embedding'] / 1e6:,.1f}M non-embedding, {num_params['embedding'] / 1e6:,.1f}M embedding, " + f"{num_params['total'] / 1e6:,.1f}M total" ) + if checkpoint_step is None: + asyncio.run(monitors.log({f"params/{name}": count for name, count in num_params.items()} | {"step": 0}, step=0)) tokenizer = setup_tokenizer(config.tokenizer) From 80cb4711df2258e0b22869df1ab4f2ba6176214d Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:48:10 +0200 Subject: [PATCH 032/132] feat(pretrain): the launcher gives the trainer and the eval worker one W&B run --- docs/pretraining.md | 2 +- src/prime_rl/entrypoints/pretrain.py | 4 ++++ src/prime_rl/trainer/pretrain/eval_worker.py | 9 ++++++++- src/prime_rl/trainer/pretrain/snapshots.py | 13 +++++++++++++ src/prime_rl/trainer/pretrain/train.py | 5 ++++- src/prime_rl/utils/pathing.py | 5 +++++ 6 files changed, 35 insertions(+), 3 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 90ce86f52e2..c5ce9f5521d 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -138,7 +138,7 @@ torchrun --nproc-per-node 8 -m prime_rl.trainer.pretrain.eval_worker \ --run-dir / --data-dir ~/evals-data --sets fwedu-holdout lambada math ``` -To log into the same W&B run as the trainer, start both in W&B shared mode with the same run id: `WANDB_SHARED_MODE=1 WANDB_RUN_ID= WANDB_SHARED_PRIMARY=trainer WANDB_SHARED_FINISHER=eval` for both, plus `WANDB_SHARED_LABEL=trainer` for the trainer and `WANDB_SHARED_LABEL=eval` for the worker. +The trainer and the worker log into one W&B run (W&B shared mode): the launcher writes the run id to `/wandb_run_id`, and both processes read it. ## Multi-node diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py index f4a07a96ebf..4ecd8611d4e 100644 --- a/src/prime_rl/entrypoints/pretrain.py +++ b/src/prime_rl/entrypoints/pretrain.py @@ -11,6 +11,7 @@ format_config_message, format_log_message, get_launcher_dir, + get_wandb_run_id_path, prepare_attempt_dirs, validate_run_dir, write_launch_artifacts, @@ -91,6 +92,9 @@ def pretrain(config: PretrainConfig): if not config.dry_run: write_model_dir(config) + if config.eval_snapshots is not None and not get_wandb_run_id_path(config.run_dir).exists(): + # The trainer and the eval worker (a separate job) log into one W&B run; a resumed run keeps its id. + get_wandb_run_id_path(config.run_dir).write_text(os.environ["PRL_RUN_ID"]) config_path = config_dir / PRETRAIN_CONFIG exclude = {"dry_run", "clean"} diff --git a/src/prime_rl/trainer/pretrain/eval_worker.py b/src/prime_rl/trainer/pretrain/eval_worker.py index 5a5e1f7077e..4cefce335fd 100644 --- a/src/prime_rl/trainer/pretrain/eval_worker.py +++ b/src/prime_rl/trainer/pretrain/eval_worker.py @@ -10,6 +10,7 @@ import asyncio import json +import os import time from datetime import timedelta @@ -24,7 +25,12 @@ from prime_rl.trainer.ckpt import AppState from prime_rl.trainer.model import setup_model from prime_rl.trainer.parallel_dims import get_parallel_dims -from prime_rl.trainer.pretrain.snapshots import delete_snapshots_up_to, finished_snapshot_steps, snapshot_dir +from prime_rl.trainer.pretrain.snapshots import ( + delete_snapshots_up_to, + finished_snapshot_steps, + shared_wandb_env, + snapshot_dir, +) from prime_rl.trainer.pretrain.suite import prepare_set, score_requests, summarize, use_scoring_head from prime_rl.trainer.utils import get_dcp_group, setup_torch_distributed from prime_rl.trainer.world import get_world @@ -62,6 +68,7 @@ def main_loop(config: EvalWorkerConfig) -> None: setup_torch_distributed(timeout=timedelta(seconds=run.dist_timeout_seconds)) torch.set_float32_matmul_precision(run.matmul_precision) + os.environ.update(shared_wandb_env(config.run_dir, "eval")) asyncio.run( monitors.setup( producer="eval", diff --git a/src/prime_rl/trainer/pretrain/snapshots.py b/src/prime_rl/trainer/pretrain/snapshots.py index c9e94f30b1f..232d59d12e7 100644 --- a/src/prime_rl/trainer/pretrain/snapshots.py +++ b/src/prime_rl/trainer/pretrain/snapshots.py @@ -9,10 +9,23 @@ from prime_rl.trainer.ckpt import AppState from prime_rl.trainer.utils import get_dcp_group from prime_rl.trainer.world import get_world +from prime_rl.utils.pathing import get_wandb_run_id_path FINISHED_MARKER = ".finished" +def shared_wandb_env(run_dir: Path, label: str) -> dict[str, str]: + """W&B shared-mode settings that put the trainer (``label = "trainer"``, creates the run) and the eval worker + (``"eval"``, finishes it) into one W&B run, whose id the launcher writes to ``/wandb_run_id``.""" + return { + "WANDB_SHARED_MODE": "1", + "WANDB_RUN_ID": get_wandb_run_id_path(run_dir).read_text().strip(), + "WANDB_SHARED_LABEL": label, + "WANDB_SHARED_PRIMARY": "trainer", + "WANDB_SHARED_FINISHER": "eval", + } + + def snapshot_dir(run_dir: Path) -> Path: return run_dir / "snapshots" diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index ef68e1b2f09..f559038dc05 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -1,6 +1,7 @@ import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before import import asyncio +import os import math import time from datetime import timedelta @@ -44,7 +45,7 @@ ) from prime_rl.trainer.pretrain.schedules import LRSchedule from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state -from prime_rl.trainer.pretrain.snapshots import SnapshotWriter +from prime_rl.trainer.pretrain.snapshots import SnapshotWriter, shared_wandb_env from prime_rl.trainer.step import run_micro_batches from prime_rl.trainer.metrics import DeadExperts, SpikeDetector, grad_norms_by_group from prime_rl.trainer.utils import ( @@ -80,6 +81,8 @@ def train(config: PretrainConfig): logger = setup_logger(config.log.level, json_logging=config.log.json_logging) logger.info(f"Starting pretrain trainer in {world} (output_dir={config.run_dir})") + if config.eval_snapshots is not None: + os.environ.update(shared_wandb_env(config.run_dir, "trainer")) asyncio.run( monitors.setup( producer="trainer", diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index f660684c7a1..03109ab0597 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -284,6 +284,11 @@ def get_eval_dir(output_dir: Path) -> Path: return output_dir / "eval" +def get_wandb_run_id_path(output_dir: Path) -> Path: + """The W&B run id that a pretraining run's trainer and its eval worker (separate jobs) share.""" + return output_dir / "wandb_run_id" + + def get_broadcast_dir(output_dir: Path) -> Path: return output_dir / "broadcasts" From f56c5a21ffc673695c2f1f59df2864732b206a43 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:51:24 +0200 Subject: [PATCH 033/132] test: unique basename for the step metrics test --- tests/unit/train/{test_metrics.py => test_step_metrics.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/unit/train/{test_metrics.py => test_step_metrics.py} (100%) diff --git a/tests/unit/train/test_metrics.py b/tests/unit/train/test_step_metrics.py similarity index 100% rename from tests/unit/train/test_metrics.py rename to tests/unit/train/test_step_metrics.py From a87bfa77506153fdf820b3ed68f21153f957ad0a Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:51:43 +0200 Subject: [PATCH 034/132] feat(pretrain): pretraining defaults to model.offload = none Custom parameter groups don't support optimizer offloading, which is the shared default. --- packages/prime-rl-configs/src/prime_rl/configs/pretrain.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index afbbefc232f..7d6c4e6dbe1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -11,6 +11,7 @@ ModelConfig, MoERuntimeConfig, MuonConfig, + OffloadMode, OptimizerConfig, TokenizerConfig, ) @@ -72,6 +73,9 @@ class PretrainModelConfig(ModelConfig): moe: PretrainMoEConfig = PretrainMoEConfig() + offload: OffloadMode = "none" + """CPU offloading. Pretraining's per-role parameter groups need ``none``.""" + @model_validator(mode="after") def validate_arch(self): if "model_type" not in self.arch: From c18d9c7e582ab3fce223439cac5883a8ea173d57 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 18:51:53 +0200 Subject: [PATCH 035/132] fix(pretrain): round fractional schedule phases down round(0.1 * N) + round(0.9 * N) can exceed N (N = 21095 gives 21096), so the default schedule failed validation for some max_steps. Fractions now round down, so fractions that sum to at most 1 always fit. --- packages/prime-rl-configs/src/prime_rl/configs/pretrain.py | 4 ++-- tests/unit/train/pretrain/test_schedules.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index de4a6092981..7d3695f91a0 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -158,8 +158,8 @@ class ValConfig(BaseConfig): def resolve_steps(value: int | float, total_steps: int) -> int: - """Steps as given, or a fraction of ``total_steps``.""" - return round(value * total_steps) if isinstance(value, float) else value + """Steps as given, or a fraction of ``total_steps`` rounded down, so fractions that sum to at most 1 fit.""" + return math.floor(value * total_steps) if isinstance(value, float) else value class ScheduleConfig(BaseConfig): diff --git a/tests/unit/train/pretrain/test_schedules.py b/tests/unit/train/pretrain/test_schedules.py index fa6dc486e4a..dc867d73d11 100644 --- a/tests/unit/train/pretrain/test_schedules.py +++ b/tests/unit/train/pretrain/test_schedules.py @@ -39,6 +39,8 @@ def test_schedule_phases_must_fit_max_steps(): ) with pytest.raises(ValidationError, match="exceed max_steps"): PretrainConfig(**{**config.model_dump(), "schedule": {"warmup": 50, "decay": 51}}) + # Fractions summing to 1 fit for any max_steps (rounding both up used to exceed it). + PretrainConfig(**{**config.model_dump(), "max_steps": 21095}) def test_apply_keeps_group_ratios(): From 3275e0ad26bd235e56071b060c4193df2accb310 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 19:04:01 +0200 Subject: [PATCH 036/132] fix(pretrain): wait for the snapshot copies before the background write --- src/prime_rl/trainer/pretrain/snapshots.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/prime_rl/trainer/pretrain/snapshots.py b/src/prime_rl/trainer/pretrain/snapshots.py index 232d59d12e7..cbae952ae41 100644 --- a/src/prime_rl/trainer/pretrain/snapshots.py +++ b/src/prime_rl/trainer/pretrain/snapshots.py @@ -69,6 +69,8 @@ def save(self, model: nn.Module, step: int) -> None: self.wait() state = AppState(model, [], None, None).state_dict() state["model"] = {name: snapshot_tensor(name, tensor) for name, tensor in state["model"].items()} + # DCP copies the tensors to the CPU on its own CUDA stream, which does not wait for the copies above. + torch.cuda.current_stream().synchronize() self.pending = self.executor.submit(self.write, state, snapshot_dir(self.run_dir) / f"step_{step}") def write(self, state: dict, path: Path) -> None: From 489f963d5c040774c24836a83c6f0ac981c93507 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 19:19:59 +0200 Subject: [PATCH 037/132] fix(pretrain): the snapshot writer thread uses the rank's CUDA device --- src/prime_rl/trainer/pretrain/snapshots.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/pretrain/snapshots.py b/src/prime_rl/trainer/pretrain/snapshots.py index cbae952ae41..e977067b460 100644 --- a/src/prime_rl/trainer/pretrain/snapshots.py +++ b/src/prime_rl/trainer/pretrain/snapshots.py @@ -60,7 +60,10 @@ class SnapshotWriter: def __init__(self, run_dir: Path): self.run_dir = run_dir - self.executor = ThreadPoolExecutor(max_workers=1) + # A new thread starts on CUDA device 0; DCP's GPU-to-CPU copies must run on this rank's device. + self.executor = ThreadPoolExecutor( + max_workers=1, initializer=torch.cuda.set_device, initargs=(torch.cuda.current_device(),) + ) self.pending: Future | None = None # Created here, on every rank at the same point; the writer thread only uses it. self.group = get_dcp_group() From 7247c32bd7bca0b75ee7accba2bc604d834ec007 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 19:41:32 +0200 Subject: [PATCH 038/132] fix(pretrain): finish the snapshot write before a checkpoint save --- src/prime_rl/trainer/pretrain/train.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 64831c8eaea..65b9a36c91c 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -346,6 +346,9 @@ def run_evals(step: int) -> None: disk_metrics = {} if config.ckpt and config.ckpt.interval and progress.step % config.ckpt.interval == 0 and not is_last_step: save_ckpt_start_time = time.perf_counter() + if snapshots is not None: + # Both write through the same DCP process group; their collectives must not interleave. + snapshots.wait() ckpt_manager.save(progress.step, model, [optimizer], None, progress, dataloader=dataloader) ckpt_manager.maybe_clean() save_ckpt_time = time.perf_counter() - save_ckpt_start_time From 9ee2f65ac689b855e72cd9e47c789878c733096b Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 19:59:57 +0200 Subject: [PATCH 039/132] fix(pretrain): leave the SFT step and its metrics as they are; the once-per-step reduction is pretrain-only --- .../trainer/{ => pretrain}/metrics.py | 0 src/prime_rl/trainer/pretrain/step.py | 145 ++++++++++++++++++ src/prime_rl/trainer/pretrain/train.py | 4 +- src/prime_rl/trainer/sft/train.py | 12 +- src/prime_rl/trainer/step.py | 121 ++++++--------- .../test_metrics.py} | 2 +- 6 files changed, 205 insertions(+), 79 deletions(-) rename src/prime_rl/trainer/{ => pretrain}/metrics.py (100%) create mode 100644 src/prime_rl/trainer/pretrain/step.py rename tests/unit/train/{test_step_metrics.py => pretrain/test_metrics.py} (93%) diff --git a/src/prime_rl/trainer/metrics.py b/src/prime_rl/trainer/pretrain/metrics.py similarity index 100% rename from src/prime_rl/trainer/metrics.py rename to src/prime_rl/trainer/pretrain/metrics.py diff --git a/src/prime_rl/trainer/pretrain/step.py b/src/prime_rl/trainer/pretrain/step.py new file mode 100644 index 00000000000..60bd169ffd0 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/step.py @@ -0,0 +1,145 @@ +"""The pretrain loop's forward and backward passes over one step's micro-batches (the SFT loop uses +`prime_rl.trainer.step`): the same loss and gradient scaling, with the step's statistics reduced once.""" + +import time +from collections.abc import Callable, Iterator +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch import Tensor, nn + +from prime_rl.trainer.model import is_tt_moe_model, iter_moe_blocks +from prime_rl.trainer.parallel_dims import ParallelDims +from prime_rl.trainer.pretrain.metrics import expert_load_metrics, reduce_stats, take_routing_counts +from prime_rl.trainer.utils import begin_backward, finish_backward, prepare_gradient_offload, scale_gradients_ + + +@dataclass +class StepResult: + loss: float + """Mean loss per unmasked token over the global batch.""" + + nan_count: int + """Micro-batches with a NaN loss over all ranks. They are left out of the backward pass.""" + + num_tokens: int + """Input tokens of the global batch.""" + + tokens_per_expert: Tensor | None + """Global ``[num_moe_layers, num_experts]`` token counts of the step (fp32, on the GPU). None for dense models.""" + + metrics: dict[str, float] + """Statistics of the step over all ranks: data wait and peak memory (max and mean over ranks), and for MoE + models the expert balance of the global counts (`expert_load_metrics`), the busiest expert of one rank's + micro-batch over the mean (``expert_load/micro_max_mean``, mean and max over ranks and micro-batches) and the + routing confidence.""" + + sums: dict[str, float] + """The caller's ``extra_sums``, summed over the ranks.""" + + +def run_micro_batches( + model: nn.Module, + dataiter: Iterator[dict], + compute_loss: Callable[[dict], tuple[Tensor, Tensor, Tensor | float]], + grad_accum_steps: int, + parallel_dims: ParallelDims, + gradient_manager=None, + *, + extra_sums: dict[str, Tensor] | None = None, +) -> StepResult: + """Forward and backward passes over the ``grad_accum_steps`` micro-batches of one step, leaving the + gradients of the summed token loss divided by the global token count (the mean loss per token, as on one + GPU without accumulation). ``compute_loss`` returns the summed loss and the number of unmasked tokens of a + micro-batch, and an auxiliary loss on the same scale as the summed loss that is trained but not logged. + + The step's statistics stay on the GPU during the micro-batches (no collective, no host sync) and are reduced + with one SUM and one MAX all-reduce at the end, together with ``extra_sums`` (GPU tensors the caller adds to + in ``compute_loss``).""" + dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() + is_moe_model = is_tt_moe_model(model) + + if gradient_manager is not None: + # The offloaded gradients are scaled as they stream out, so the token count is needed up front. + micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] + global_token_count = torch.tensor( + sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches), device="cuda" + ) + dist.all_reduce(global_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count = global_token_count.item() // parallel_dims.cp + grad_scale = ( + parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count + if global_token_count > 0 + else 1.0 + ) + prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=True) + + zero = torch.zeros((), device="cuda", dtype=torch.float64) + sums = {"loss": zero.clone(), "loss_tokens": zero.clone(), "nan_count": zero.clone(), "tokens": zero.clone()} + maxes = {} + data_time = 0.0 + for micro_step in range(grad_accum_steps): + data_start_time = time.perf_counter() + micro_batch = next(dataiter) if gradient_manager is None else micro_batches[micro_step] + data_time += time.perf_counter() - data_start_time + sums["tokens"] += micro_batch["input_ids"].shape[1] + micro_loss_sum, micro_token_count, aux_loss_sum = compute_loss(micro_batch) + sums["loss_tokens"] += micro_token_count + # A NaN micro-batch is left out of the loss and the backward, without a host sync. + is_nan = torch.isnan(micro_loss_sum.detach()) + sums["nan_count"] += is_nan + sums["loss"] += torch.where(is_nan, 0.0, micro_loss_sum.detach()) + scaled_loss = torch.where(is_nan, torch.nan_to_num(micro_loss_sum, nan=0.0), micro_loss_sum + aux_loss_sum) + begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) + (scaled_loss / grad_accum_steps).backward() + finish_backward(gradient_manager) + + if is_moe_model: + counts, confidence = take_routing_counts(model) + sums["expert_counts"] = sums.get("expert_counts", zero) + counts + sums["routing_confidence"] = sums.get("routing_confidence", zero) + confidence + micro_max_mean = counts.amax(dim=1) / counts.mean(dim=1) + sums["micro_max_mean"] = sums.get("micro_max_mean", zero) + micro_max_mean.mean() + maxes["micro_max_mean_max"] = torch.maximum(maxes.get("micro_max_mean_max", zero), micro_max_mean.amax()) + + # Peak memory since the caller last reset the counter (at the start of the step: the forward and backward passes). + peak_memory = torch.cuda.max_memory_reserved() / 1024**3 + sums |= {"data_time": zero + data_time, "peak_memory": zero + peak_memory} + maxes |= {"data_time_max": zero + data_time, "peak_memory_max": zero + peak_memory} + extra_sums = extra_sums or {} + stats = reduce_stats(sums | {f"extra/{name}": value for name, value in extra_sums.items()}, maxes, dp_cp_group) + + if gradient_manager is None: + global_token_count = int(stats["loss_tokens"]) + if global_token_count > 0: + scale_gradients_( + None, model, parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count + ) + + num_ranks = dist.get_world_size(dp_cp_group) + metrics = { + "time/data/max": stats["data_time_max"].item(), + "time/data/mean": stats["data_time"].item() / num_ranks, + "perf/peak_memory/max": stats["peak_memory_max"].item(), + "perf/peak_memory/mean": stats["peak_memory"].item() / num_ranks, + } + tokens_per_expert = None + if is_moe_model: + counts = stats["expert_counts"] + routed_tokens = counts.sum(dim=1) / next(iter_moe_blocks(model)).router.top_k + metrics |= expert_load_metrics(counts) | { + "expert_load/micro_max_mean/mean": stats["micro_max_mean"].item() / (grad_accum_steps * num_ranks), + "expert_load/micro_max_mean/max": stats["micro_max_mean_max"].item(), + "router/confidence": (stats["routing_confidence"] / routed_tokens).mean().item(), + } + tokens_per_expert = counts.to("cuda", torch.float32) + return StepResult( + loss=stats["loss"].item() / global_token_count if global_token_count > 0 else 0.0, + nan_count=int(stats["nan_count"]), + # Every CP rank counts its rows in full. + num_tokens=int(stats["tokens"]) // parallel_dims.cp, + tokens_per_expert=tokens_per_expert, + metrics=metrics, + sums={name: stats[f"extra/{name}"].item() for name in extra_sums}, + ) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 65b9a36c91c..af9883eef6a 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -45,8 +45,8 @@ from prime_rl.trainer.pretrain.schedules import LRSchedule from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state from prime_rl.trainer.pretrain.snapshots import SnapshotWriter, shared_wandb_env -from prime_rl.trainer.step import run_micro_batches -from prime_rl.trainer.metrics import DeadExperts, SpikeDetector, grad_norms_by_group +from prime_rl.trainer.pretrain.step import run_micro_batches +from prime_rl.trainer.pretrain.metrics import DeadExperts, SpikeDetector, grad_norms_by_group from prime_rl.trainer.utils import ( GarbageCollection, clip_grad_norm_, diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index f19f065a53a..22caeb5e2a7 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -27,6 +27,7 @@ from prime_rl.trainer.model import ( forward, get_full_offload_dtype_policy, + get_expert_load_stats, get_load_balance_stats, is_tt_moe_model, reshard_module, @@ -430,6 +431,7 @@ def is_online_eval_step(step: int) -> bool: record_function=maybe_record_function, ) forward_backward_time = time.perf_counter() - forward_backward_start_time + expert_load_stats = get_expert_load_stats(result.tokens_per_expert, dp_cp_group) if is_moe_model else {} # Run validation after forward-backward (so torch.compile sees training graph first) but before # optimizer step (so eval_on_start evaluates untrained weights) @@ -496,8 +498,10 @@ def is_online_eval_step(step: int) -> bool: step_message += f" | Grad. Norm {grad_norm:.4f}" step_message += f" | LR {current_lr:.2e} | Throughput {throughput:.0f} tokens/s | MFU {mfu:.1f}% | Peak Mem. {peak_memory:.1f}/{max_memory:.1f} GiB ({peak_memory / max_memory * 100:.1f}%)" if is_moe_model: - for name, label in (("expert_load/max_mean/mean", "Max/Mean Load"), ("router/confidence", "Routing Conf.")): - step_message += f" | {label} {result.metrics[name]:.4f}" + for name, label in (("max_vio/mean", "Max Vio"), ("routing_confidence/mean", "Routing Conf.")): + value = result.moe_stats[name] + if value > 0: + step_message += f" | {label} {value:.4f}" logger.success(step_message) # Log progress metrics @@ -568,7 +572,9 @@ def is_online_eval_step(step: int) -> bool: disk_metrics["step"] = progress.step asyncio.run(monitors.log(disk_metrics, step=progress.step)) - asyncio.run(monitors.log({**result.metrics, "step": progress.step}, step=progress.step)) + moe_log_metrics = result.moe_stats | expert_load_stats + if moe_log_metrics: + asyncio.run(monitors.log({**moe_log_metrics, "step": progress.step}, step=progress.step)) is_first_step = False diff --git a/src/prime_rl/trainer/step.py b/src/prime_rl/trainer/step.py index 6b660839e82..dead2636d59 100644 --- a/src/prime_rl/trainer/step.py +++ b/src/prime_rl/trainer/step.py @@ -1,4 +1,3 @@ -import time from collections.abc import Callable, Iterator from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass @@ -7,10 +6,10 @@ import torch.distributed as dist from torch import Tensor, nn -from prime_rl.trainer.metrics import expert_load_metrics, reduce_stats, take_routing_counts -from prime_rl.trainer.model import is_tt_moe_model, iter_moe_blocks +from prime_rl.trainer.model import get_global_moe_stats, is_tt_moe_model from prime_rl.trainer.parallel_dims import ParallelDims from prime_rl.trainer.utils import begin_backward, finish_backward, prepare_gradient_offload, scale_gradients_ +from prime_rl.utils.logger import get_logger @dataclass @@ -24,17 +23,12 @@ class StepResult: num_tokens: int """Input tokens of the global batch.""" - tokens_per_expert: Tensor | None - """Global ``[num_moe_layers, num_experts]`` token counts of the step (fp32, on the GPU). None for dense models.""" - - metrics: dict[str, float] - """Statistics of the step over all ranks: data wait and peak memory (max and mean over ranks), and for MoE - models the expert balance of the global counts (`expert_load_metrics`), the busiest expert of one rank's - micro-batch over the mean (``expert_load/micro_max_mean``, mean and max over ranks and micro-batches) and the - routing confidence.""" + moe_stats: dict[str, float] + """Routing statistics averaged over the micro-batches (``max_vio/mean``, ``max_vio/max``, + ``routing_confidence/mean``). Empty for dense models.""" - sums: dict[str, float] - """The caller's ``extra_sums``, summed over the ranks.""" + tokens_per_expert: Tensor | None + """Global ``[num_moe_layers, num_experts]`` token counts of the step. None for dense models.""" def run_micro_batches( @@ -47,20 +41,20 @@ def run_micro_batches( *, overlap_optimizer: bool = True, record_function: Callable[[str], AbstractContextManager] = nullcontext, - extra_sums: dict[str, Tensor] | None = None, ) -> StepResult: """Forward and backward passes over the ``grad_accum_steps`` micro-batches of one step, leaving the gradients of the summed token loss divided by the global token count (the mean loss per token, as on one GPU without accumulation). ``compute_loss`` returns the summed loss and the number of unmasked tokens of a - micro-batch, and an auxiliary loss on the same scale as the summed loss that is trained but not logged. - - The step's statistics stay on the GPU during the micro-batches (no collective, no host sync) and are reduced - with one SUM and one MAX all-reduce at the end, together with ``extra_sums`` (GPU tensors the caller adds to - in ``compute_loss``).""" + micro-batch, and an auxiliary loss on the same scale as the summed loss that is trained but not logged.""" + logger = get_logger() dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() + ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None is_moe_model = is_tt_moe_model(model) - if gradient_manager is not None: + if gradient_manager is None: + micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) + local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") + else: # The offloaded gradients are scaled as they stream out, so the token count is needed up front. micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] global_token_count = torch.tensor( @@ -75,73 +69,54 @@ def run_micro_batches( ) prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=overlap_optimizer) - zero = torch.zeros((), device="cuda", dtype=torch.float64) - sums = {"loss": zero.clone(), "loss_tokens": zero.clone(), "nan_count": zero.clone(), "tokens": zero.clone()} - maxes = {} - data_time = 0.0 - for micro_step in range(grad_accum_steps): - data_start_time = time.perf_counter() - micro_batch = next(dataiter) if gradient_manager is None else micro_batches[micro_step] - data_time += time.perf_counter() - data_start_time - sums["tokens"] += micro_batch["input_ids"].shape[1] + loss_sum = torch.tensor(0.0, device="cuda") + nan_count = torch.tensor(0, device="cuda") + moe_stats = {name: torch.tensor(0.0) for name in ("max_vio/mean", "max_vio/max", "routing_confidence/mean")} + tokens_per_expert = None + local_num_tokens = 0 + for micro_step, micro_batch in enumerate(micro_batches): + local_num_tokens += micro_batch["input_ids"].shape[1] with record_function("forward"): micro_loss_sum, micro_token_count, aux_loss_sum = compute_loss(micro_batch) - sums["loss_tokens"] += micro_token_count - # A NaN micro-batch is left out of the loss and the backward, without a host sync. - is_nan = torch.isnan(micro_loss_sum.detach()) - sums["nan_count"] += is_nan - sums["loss"] += torch.where(is_nan, 0.0, micro_loss_sum.detach()) - scaled_loss = torch.where(is_nan, torch.nan_to_num(micro_loss_sum, nan=0.0), micro_loss_sum + aux_loss_sum) + if gradient_manager is None: + local_token_count += micro_token_count + if torch.isnan(micro_loss_sum.detach()): + nan_count += 1 + logger.warning("Local loss is nan, excluding this micro step from backward") + scaled_loss = torch.nan_to_num(micro_loss_sum, nan=0.0) / grad_accum_steps + else: + loss_sum += micro_loss_sum.detach() + scaled_loss = (micro_loss_sum + aux_loss_sum) / grad_accum_steps with record_function("backward"): begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) - (scaled_loss / grad_accum_steps).backward() + scaled_loss.backward() finish_backward(gradient_manager) if is_moe_model: - counts, confidence = take_routing_counts(model) - sums["expert_counts"] = sums.get("expert_counts", zero) + counts - sums["routing_confidence"] = sums.get("routing_confidence", zero) + confidence - micro_max_mean = counts.amax(dim=1) / counts.mean(dim=1) - sums["micro_max_mean"] = sums.get("micro_max_mean", zero) + micro_max_mean.mean() - maxes["micro_max_mean_max"] = torch.maximum(maxes.get("micro_max_mean_max", zero), micro_max_mean.amax()) - - # Peak memory since the caller last reset the counter (at the start of the step: the forward and backward passes). - peak_memory = torch.cuda.max_memory_reserved() / 1024**3 - sums |= {"data_time": zero + data_time, "peak_memory": zero + peak_memory} - maxes |= {"data_time_max": zero + data_time, "peak_memory_max": zero + peak_memory} - extra_sums = extra_sums or {} - stats = reduce_stats(sums | {f"extra/{name}": value for name, value in extra_sums.items()}, maxes, dp_cp_group) + micro_moe_stats, micro_tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) + for name, value in micro_moe_stats.items(): + moe_stats[f"{name}/mean"] += value / grad_accum_steps + if name == "max_vio": + moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) + tokens_per_expert = micro_tokens_per_expert + (0 if tokens_per_expert is None else tokens_per_expert) if gradient_manager is None: - global_token_count = int(stats["loss_tokens"]) + dist.all_reduce(local_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count = local_token_count.item() if global_token_count > 0: scale_gradients_( None, model, parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count ) - num_ranks = dist.get_world_size(dp_cp_group) - metrics = { - "time/data/max": stats["data_time_max"].item(), - "time/data/mean": stats["data_time"].item() / num_ranks, - "perf/peak_memory/max": stats["peak_memory_max"].item(), - "perf/peak_memory/mean": stats["peak_memory"].item() / num_ranks, - } - tokens_per_expert = None - if is_moe_model: - counts = stats["expert_counts"] - routed_tokens = counts.sum(dim=1) / next(iter_moe_blocks(model)).router.top_k - metrics |= expert_load_metrics(counts) | { - "expert_load/micro_max_mean/mean": stats["micro_max_mean"].item() / (grad_accum_steps * num_ranks), - "expert_load/micro_max_mean/max": stats["micro_max_mean_max"].item(), - "router/confidence": (stats["routing_confidence"] / routed_tokens).mean().item(), - } - tokens_per_expert = counts.to("cuda", torch.float32) + dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) + dist.all_reduce(nan_count, op=dist.ReduceOp.SUM) + # The dp mesh excludes cp, whose ranks hold the same rows. + num_tokens = torch.tensor(local_num_tokens, dtype=torch.int64, device="cuda") + dist.all_reduce(num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) return StepResult( - loss=stats["loss"].item() / global_token_count if global_token_count > 0 else 0.0, - nan_count=int(stats["nan_count"]), - # Every CP rank counts its rows in full. - num_tokens=int(stats["tokens"]) // parallel_dims.cp, + loss=(loss_sum / global_token_count).item() if global_token_count > 0 else 0.0, + nan_count=int(nan_count.item()), + num_tokens=int(num_tokens.item()), + moe_stats={name: value.item() for name, value in moe_stats.items()} if is_moe_model else {}, tokens_per_expert=tokens_per_expert, - metrics=metrics, - sums={name: stats[f"extra/{name}"].item() for name in extra_sums}, ) diff --git a/tests/unit/train/test_step_metrics.py b/tests/unit/train/pretrain/test_metrics.py similarity index 93% rename from tests/unit/train/test_step_metrics.py rename to tests/unit/train/pretrain/test_metrics.py index 00ffb606cbe..55e8de6b3ba 100644 --- a/tests/unit/train/test_step_metrics.py +++ b/tests/unit/train/pretrain/test_metrics.py @@ -1,7 +1,7 @@ import torch from torch import nn -from prime_rl.trainer.metrics import expert_load_metrics, grad_norms_by_group, param_group +from prime_rl.trainer.pretrain.metrics import expert_load_metrics, grad_norms_by_group, param_group def test_expert_load_metrics_per_layer_summaries(): From aa398a47711f63cd465a853b6638d330c789c4b0 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 20:01:48 +0200 Subject: [PATCH 040/132] refactor(pretrain): write the model dir to /configs/model configs/ is already a launcher artifact, so the model dir doesn't count as an artifact of a previous run and the shared run-dir check doesn't need to know about it. --- docs/pretraining.md | 2 +- packages/prime-rl-configs/src/prime_rl/configs/pretrain.py | 6 +++--- src/prime_rl/entrypoints/pretrain.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 358f39c4cad..ba644b97b5c 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -22,7 +22,7 @@ num_experts_per_tok = 6 name = "tokenizers/deepseek-v4-flash-base" ``` -The launcher writes the config and the tokenizer to `/model` and trains from there, so the trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all see an ordinary model directory. The trainer logs the non-embedding, embedding (input embedding and LM head) and total parameter counts at startup. +The launcher writes the config and the tokenizer to `/configs/model` and trains from there, so the trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all see an ordinary model directory. The trainer logs the non-embedding, embedding (input embedding and LM head) and total parameter counts at startup. ## Initialization diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 7d3695f91a0..9cfc747f397 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -35,7 +35,7 @@ class InitConfig(BaseConfig): class PretrainModelConfig(ModelConfig): arch: dict[str, Any] - """The architecture: ``model_type`` (a model in ``prime_rl.trainer.models``, e.g. ``glm4_moe``) and the fields of its config class. ``vocab_size`` and the special token ids come from the tokenizer. The launcher writes the config and the tokenizer to ``/model``, which ``name`` points at.""" + """The architecture: ``model_type`` (a model in ``prime_rl.trainer.models``, e.g. ``glm4_moe``) and the fields of its config class. ``vocab_size`` and the special token ids come from the tokenizer. The launcher writes the config and the tokenizer to ``/configs/model``, which ``name`` points at.""" init: InitConfig = InitConfig() """Random weight initialization. Used unless the run resumes from a checkpoint.""" @@ -281,8 +281,8 @@ def auto_setup_run_identity(self): @model_validator(mode="after") def auto_setup_model_dir(self): - """The model is trained from ``/model``, written by the launcher from ``model.arch``.""" - self.model.name = (self.run_dir / "model").as_posix() + """The model is trained from ``/configs/model``, written by the launcher from ``model.arch``.""" + self.model.name = (self.run_dir / "configs" / "model").as_posix() return self @model_validator(mode="after") diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py index 7c2c503bf12..f3749b0ea4c 100644 --- a/src/prime_rl/entrypoints/pretrain.py +++ b/src/prime_rl/entrypoints/pretrain.py @@ -25,7 +25,7 @@ def write_model_dir(config: PretrainConfig) -> None: - """Write the config built from ``[model.arch]`` and the tokenizer to ``/model``. The vocabulary size + """Write the config built from ``[model.arch]`` and the tokenizer to ``/configs/model``. The vocabulary size (rounded up to a multiple of 128) and the special token ids come from the tokenizer.""" from transformers import AutoTokenizer from transformers.models.auto.configuration_auto import CONFIG_MAPPING From 1ecbb195ea2530fbfa9d4ef7d5ed4c15ecf26b81 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 20:02:32 +0200 Subject: [PATCH 041/132] revert(configs): keep SFTConfig as it was; PretrainConfig is standalone again New rule: the pretrain work must not change SFT or RL behavior. TrainRunConfig changed SFTConfig's field order (and so its resolved config dumps) and validator order. This reverts 20578b007 and the 'model' launcher artifact in utils/pathing.py. PretrainConfig keeps the DeepEP and full-offload validators it gained, as its own copies. The SFT loop's forward profiler annotation wraps only the forward again, as before run_micro_batches. --- .../src/prime_rl/configs/pretrain.py | 105 ++++++- .../src/prime_rl/configs/sft.py | 264 +++++++++--------- src/prime_rl/trainer/sft/train.py | 3 +- src/prime_rl/trainer/step.py | 3 +- src/prime_rl/utils/pathing.py | 2 +- 5 files changed, 231 insertions(+), 146 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 9c2a80fca93..3d2d4c05e64 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -1,19 +1,23 @@ import math import uuid +import warnings from pathlib import Path from typing import Annotated, Any, Literal, TypeAlias from pydantic import Field, model_validator -from prime_rl.configs.sft import TrainRunConfig +from prime_rl.configs.monitors import TrainMonitorsConfig +from prime_rl.configs.sft import SFTDeploymentConfig, SingleNodeDeploymentConfig, setup_trainer_deployment +from prime_rl.configs.shared import EnvVars, HeartbeatConfig, ResumeConfig, RunConfig, SlurmConfig, TrainerLogConfig from prime_rl.configs.trainer import ( CheckpointConfig, + GCConfig, ModelConfig, MuonConfig, OptimizerConfig, TokenizerConfig, ) -from prime_rl.utils.config import BaseConfig +from prime_rl.utils.config import BaseConfig, default_output_dir TOKENIZER_FIELDS = ("vocab_size", "bos_token_id", "eos_token_id", "pad_token_id") @@ -180,9 +184,12 @@ class ScheduleConfig(BaseConfig): """Shape of the decay from 1 to ``end_ratio``.""" -class PretrainConfig(TrainRunConfig): +class PretrainConfig(BaseConfig): model: PretrainModelConfig + env_vars: EnvVars = {} + """Extra environment variables for the trainer processes.""" + tokenizer: TokenizerConfig """The tokenizer (``tokenizer.name`` is required). It sets the vocabulary and special tokens of the model.""" @@ -200,12 +207,51 @@ class PretrainConfig(TrainRunConfig): ckpt: CheckpointConfig | None = CheckpointConfig() """Checkpointing. Defaults to a checkpoint at the end of training.""" + resume: ResumeConfig | None = None + """Resume the run from a checkpoint. If None, does not resume.""" + + log: TrainerLogConfig = TrainerLogConfig() + + monitors: TrainMonitorsConfig = TrainMonitorsConfig() + """Metric monitors (``monitors.wandb``, ``monitors.file``).""" + + run: RunConfig = Field(default_factory=RunConfig) + """Run metadata. ``run.name`` names the run directory under ``output_dir``.""" + + output_dir: Path = Field(default_factory=default_output_dir) + """Directory that groups related runs. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``.""" + + clean: bool = False + """Delete the run directory before starting. Required to overwrite a previous run's artifacts when not resuming.""" + + matmul_precision: Literal["highest", "high", "medium"] = "high" + """Precision for float32 matrix multiplications (``torch.set_float32_matmul_precision``).""" + max_steps: int = Field(ge=1) """Training steps.""" - def default_run_name(self) -> str: - """``--``.""" - return f"{self.model.arch['model_type']}--{uuid.uuid4().hex[:8]}" + gc: GCConfig | None = GCConfig() + """Deterministic garbage collection every N steps. Set to null for Python's default GC.""" + + dist_timeout_seconds: int = 900 + """Timeout in seconds for torch distributed ops.""" + + heartbeat: HeartbeatConfig | None = None + """BetterStack heartbeat configuration.""" + + deployment: SFTDeploymentConfig = SingleNodeDeploymentConfig() + """``single_node`` (``num_train_gpus`` on this node) or ``multi_node`` (``num_train_nodes`` x ``gpus_per_node``, needs ``[slurm]``). The online-eval fields (``num_infer_*``) are not used.""" + + slurm: SlurmConfig | None = None + """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" + + dry_run: bool = False + """Only validate and dump resolved configs, then exit.""" + + @property + def run_dir(self) -> Path: + assert self.run.dir is not None # resolved at construction + return self.output_dir / self.run.dir @model_validator(mode="after") def resolve_autos(self): @@ -224,6 +270,16 @@ def validate_lr(self): raise ValueError("optim.lr is required") return self + @model_validator(mode="after") + def auto_setup_run_identity(self): + if self.run.name is None: + self.run.name = f"{self.model.arch['model_type']}--{uuid.uuid4().hex[:8]}" + if self.run.dir is None: + self.run.dir = self.run.name + if self.monitors.wandb is not None and self.monitors.wandb.name is None: + self.monitors.wandb.name = self.run.name + return self + @model_validator(mode="after") def auto_setup_model_dir(self): """The model is trained from ``/configs/model``, written by the launcher from ``model.arch``.""" @@ -231,9 +287,44 @@ def auto_setup_model_dir(self): return self @model_validator(mode="after") - def validate_tokenizer(self): + def auto_setup_tokenizer(self): if self.tokenizer.name is None: raise ValueError("tokenizer.name is required") + if self.tokenizer.trust_remote_code is None: + self.tokenizer.trust_remote_code = self.model.trust_remote_code + return self + + @model_validator(mode="after") + def auto_setup_deployment(self): + setup_trainer_deployment(self) + return self + + @model_validator(mode="after") + def deepep_disables_grad_clipping(self): + if self.model.ep != 1 and self.model.moe.dispatch.type == "deepep" and self.optim.max_norm is not None: + warnings.warn( + "Gradient clipping is not compatible with DeepEP. " + "Automatically setting optim.max_norm to None (disabled).", + stacklevel=1, + ) + self.optim.max_norm = None + return self + + @model_validator(mode="after") + def full_optimizer_offload_requires_supported_optimizer(self): + if self.model.offload == "full" and self.optim.type not in ("adamw", "sign_sgd"): + raise ValueError("Full optimizer offload only supports AdamW and SignSGD") + return self + + @model_validator(mode="after") + def full_optimizer_offload_disables_grad_clipping(self): + if self.model.offload == "full" and self.optim.max_norm is not None: + warnings.warn( + "Gradient clipping prevents optimizer-in-backward overlap with CPU optimizer offload. " + "Automatically setting optim.max_norm to None (disabled).", + stacklevel=1, + ) + self.optim.max_norm = None return self @model_validator(mode="after") diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 05c709d1415..682974a3770 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -1,7 +1,7 @@ import uuid import warnings from pathlib import Path -from typing import Annotated, Literal, TypeAlias +from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias from urllib.parse import urlparse from pydantic import AliasChoices, Field, model_validator @@ -36,6 +36,9 @@ ) from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource +if TYPE_CHECKING: + from prime_rl.configs.pretrain import PretrainConfig + class BaseDataConfig(BaseConfig): batch_size: int = Field(128, ge=1) @@ -216,18 +219,62 @@ class MultiNodeDeploymentConfig(BaseDeploymentConfig): ] -class TrainRunConfig(BaseConfig): - """Fields and validators shared by the single-process trainers (``sft`` and ``pretrain``).""" +def setup_trainer_deployment(config: "SFTConfig | PretrainConfig") -> None: + """Shared by ``sft`` and ``pretrain``: a multi-node deployment needs SLURM and sets + ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the SLURM job name defaults to the + run name and the template to the bundled single- or multi-node trainer template.""" + deployment = config.deployment + if deployment.type == "multi_node": + if config.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + if deployment.nodes_per_fsdp_group is not None: + if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" + ) + config.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group + if config.slurm is not None: + if config.slurm.job_name is None: + config.slurm.job_name = config.run.name + if config.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + config.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" + +class SFTConfig(BaseConfig): model: ModelConfig = ModelConfig() env_vars: EnvVars = {} - """Extra environment variables for the trainer process(es). Merged on top of the launcher defaults.""" + """Extra environment variables for the SFT trainer process(es). Merged on top of the launcher defaults.""" tokenizer: TokenizerConfig = TokenizerConfig() + renderer: RendererConfig = AutoRendererConfig() + """Renderer config. Defaults to auto-selecting from the tokenizer model name.""" + + data: DataConfig = SFTDataConfig() + + val: SFTValConfig | None = None + """Validation configuration. If None, no validation runs.""" + + eval: SFTOnlineEvalConfig | None = None + """Online evaluation configuration: rollout-based evals against a live inference + server that receives the trainer's weight broadcasts. If None, no online evals run.""" + + inference: InferenceConfig | None = None + """Inference server for online evals. If None (with ``eval`` set), the launcher + does not start a server and evals connect to ``eval.client.base_url``.""" + + weight_broadcast: WeightBroadcastConfig | None = None + """Trainer-to-inference weight transport for online evals. Defaults to NCCL. + LoRA and external inference use filesystem broadcast.""" + optim: OptimizerConfig = AdamWConfig() + scheduler: SchedulerConfig = ConstantSchedulerConfig() + ckpt: CheckpointConfig | None = None resume: ResumeConfig | None = None @@ -247,15 +294,65 @@ class TrainRunConfig(BaseConfig): clean: bool = False """Delete the run directory (``output_dir / run.name``) before starting training. Required to overwrite a run directory that contains artifacts from a previous run when not resuming.""" + @property + def run_dir(self) -> Path: + assert self.run.dir is not None # resolved at construction + return self.output_dir / self.run.dir + + @model_validator(mode="after") + def resolve_moe_router_dtype_auto(self): + """Resolve ``model.moe_router_dtype='auto'``: SFT defaults to bf16 routing, skipping the fp32 gate GEMM and its fp32 FSDP unit.""" + if self.model.moe_router_dtype == "auto": + self.model.moe_router_dtype = "bfloat16" + return self + + @model_validator(mode="after") + def resolve_weight_decay_auto(self): + """Resolve ``optim.weight_decay='auto'``: SFT keeps the historical 0.01 default — standard L2 regularization for supervised training.""" + if self.optim.weight_decay == "auto": + self.optim.weight_decay = 0.01 + return self + + @model_validator(mode="after") + def resolve_freeze_moe_router_auto(self): + """Resolve ``model.freeze_moe_router='auto'``: SFT trains the router.""" + if self.model.freeze_moe_router == "auto": + self.model.freeze_moe_router = False + return self + + @model_validator(mode="after") + def auto_setup_run_identity(self): + """Auto-generate the run name (``----``) when unset and + default the run directory, W&B run name and platform run name to it when not + set explicitly.""" + if self.run.name is None: + dataset = str(getattr(self.data, "name", "")).split("/")[-1] + model = self.model.name.split("/")[-1] + parts = [part for part in (dataset, model) if part] + self.run.name = "--".join([*parts, uuid.uuid4().hex[:8]]).lower() + if self.run.dir is None: + self.run.dir = self.run.name + if self.monitors.wandb is not None and self.monitors.wandb.name is None: + self.monitors.wandb.name = self.run.name + if self.monitors.prime is not None and self.monitors.prime.name is None: + self.monitors.prime.name = self.run.name + return self + matmul_precision: Literal["highest", "high", "medium"] = "high" """Precision for float32 matrix multiplications. ``highest`` is full FP32 (required on ROCm/AMD GPUs to avoid catastrophic precision loss in softmax over large vocabularies). ``high`` enables TF32 on NVIDIA GPUs for a speedup with minor precision tradeoff. See ``torch.set_float32_matmul_precision``.""" max_steps: int | None = None """Maximum training steps. If None, runs indefinitely.""" + memory_profiler_path: Path | None = None + """Path to write the memory profile to.""" + gc: GCConfig | None = GCConfig() """Garbage collection config. Disables automatic GC and runs deterministic collections every N steps to avoid stragglers. Set to null to use Python's default GC behavior.""" + trace_path: Path | None = None + """Path to write the PyTorch profiler trace to.""" + dist_timeout_seconds: int = 900 """Timeout in seconds for torch distributed ops: a hung collective (e.g. a dead rank) fails the run after this long.""" @@ -267,13 +364,14 @@ class TrainRunConfig(BaseConfig): slurm: SlurmConfig | None = None """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" + dashboard: bool = True + """Make sure a local dashboard daemon serves this run's output dir (started on + demand in interactive sessions; an already-running daemon's URL is logged).""" + dry_run: bool = False """Only validate and dump resolved configs, then exit early.""" - @property - def run_dir(self) -> Path: - assert self.run.dir is not None # resolved at construction - return self.output_dir / self.run.dir + ### Pre-validation normalization @model_validator(mode="before") @classmethod @@ -286,22 +384,23 @@ def normalize_deployment(cls, data): deployment.pop(key, None) return data - def default_run_name(self) -> str: - raise NotImplementedError + @model_validator(mode="before") + @classmethod + def propagate_model_to_inference(cls, data): + """Fill ``inference.vllm.model`` from ``model.name`` before InferenceConfig + construction, so its parser auto-resolution sees the right model name.""" + if not isinstance(data, dict): + return data + inference = data.get("inference") + model = data.get("model") + name = model.get("name") if isinstance(model, dict) else None + if isinstance(inference, dict) and name: + vllm = inference.setdefault("vllm", {}) + if isinstance(vllm, dict): + vllm.setdefault("model", name) + return data - @model_validator(mode="after") - def auto_setup_run_identity(self): - """Default the run name (see ``default_run_name``) and default the run directory, W&B run name and - platform run name to it when not set explicitly.""" - if self.run.name is None: - self.run.name = self.default_run_name() - if self.run.dir is None: - self.run.dir = self.run.name - if self.monitors.wandb is not None and self.monitors.wandb.name is None: - self.monitors.wandb.name = self.run.name - if self.monitors.prime is not None and self.monitors.prime.name is None: - self.monitors.prime.name = self.run.name - return self + ### Validate configs (e.g. raise for unsupported (combinations of) configs) @model_validator(mode="after") def deepep_disables_grad_clipping(self): @@ -331,108 +430,6 @@ def full_optimizer_offload_disables_grad_clipping(self): self.optim.max_norm = None return self - @model_validator(mode="after") - def auto_setup_deployment(self): - """A multi-node deployment needs SLURM and sets ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the - SLURM job name defaults to the run name and the template to the bundled single- or multi-node template.""" - deployment = self.deployment - if deployment.type == "multi_node": - if self.slurm is None: - raise ValueError("Must use SLURM for multi-node deployment.") - if deployment.nodes_per_fsdp_group is not None: - if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: - raise ValueError( - f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " - f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" - ) - self.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group - if self.slurm is not None: - if self.slurm.job_name is None: - self.slurm.job_name = self.run.name - if self.slurm.template_path is None: - templates_dir = find_package_resource("templates") - if templates_dir is not None: - self.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" - return self - - @model_validator(mode="after") - def auto_setup_tokenizer_trust_remote_code(self): - if self.tokenizer.trust_remote_code is None: - self.tokenizer.trust_remote_code = self.model.trust_remote_code - return self - - -class SFTConfig(TrainRunConfig): - def default_run_name(self) -> str: - """``----``.""" - dataset = str(getattr(self.data, "name", "")).split("/")[-1] - model = self.model.name.split("/")[-1] - parts = [part for part in (dataset, model) if part] - return "--".join([*parts, uuid.uuid4().hex[:8]]).lower() - - renderer: RendererConfig = AutoRendererConfig() - """Renderer config. Defaults to auto-selecting from the tokenizer model name.""" - - data: DataConfig = SFTDataConfig() - - val: SFTValConfig | None = None - """Validation configuration. If None, no validation runs.""" - - eval: SFTOnlineEvalConfig | None = None - """Online evaluation configuration: rollout-based evals against a live inference - server that receives the trainer's weight broadcasts. If None, no online evals run.""" - - inference: InferenceConfig | None = None - """Inference server for online evals. If None (with ``eval`` set), the launcher - does not start a server and evals connect to ``eval.client.base_url``.""" - - weight_broadcast: WeightBroadcastConfig | None = None - """Trainer-to-inference weight transport for online evals. Defaults to NCCL. - LoRA and external inference use filesystem broadcast.""" - - scheduler: SchedulerConfig = ConstantSchedulerConfig() - - @model_validator(mode="after") - def resolve_moe_router_dtype_auto(self): - """Resolve ``model.moe_router_dtype='auto'``: SFT defaults to bf16 routing, skipping the fp32 gate GEMM and its fp32 FSDP unit.""" - if self.model.moe_router_dtype == "auto": - self.model.moe_router_dtype = "bfloat16" - return self - - @model_validator(mode="after") - def resolve_weight_decay_auto(self): - """Resolve ``optim.weight_decay='auto'``: SFT keeps the historical 0.01 default — standard L2 regularization for supervised training.""" - if self.optim.weight_decay == "auto": - self.optim.weight_decay = 0.01 - return self - - @model_validator(mode="after") - def resolve_freeze_moe_router_auto(self): - """Resolve ``model.freeze_moe_router='auto'``: SFT trains the router.""" - if self.model.freeze_moe_router == "auto": - self.model.freeze_moe_router = False - return self - - ### Pre-validation normalization - - @model_validator(mode="before") - @classmethod - def propagate_model_to_inference(cls, data): - """Fill ``inference.vllm.model`` from ``model.name`` before InferenceConfig - construction, so its parser auto-resolution sees the right model name.""" - if not isinstance(data, dict): - return data - inference = data.get("inference") - model = data.get("model") - name = model.get("name") if isinstance(model, dict) else None - if isinstance(inference, dict) and name: - vllm = inference.setdefault("vllm", {}) - if isinstance(vllm, dict): - vllm.setdefault("model", name) - return data - - ### Validate configs (e.g. raise for unsupported (combinations of) configs) - @model_validator(mode="after") def auto_setup_online_eval(self): """Wire online evals and the trainer-to-inference weight transport.""" @@ -645,18 +642,15 @@ def validate_scheduler_steps(self): ### Auto-setup and validate shared configs - memory_profiler_path: Path | None = None - """Path to write the memory profile to.""" - - trace_path: Path | None = None - """Path to write the PyTorch profiler trace to.""" - - dashboard: bool = True - """Make sure a local dashboard daemon serves this run's output dir (started on - demand in interactive sessions; an already-running daemon's URL is logged).""" - @model_validator(mode="after") def auto_setup_tokenizer(self): if self.tokenizer.name is None: self.tokenizer.name = self.model.name + if self.tokenizer.trust_remote_code is None: + self.tokenizer.trust_remote_code = self.model.trust_remote_code + return self + + @model_validator(mode="after") + def auto_setup_deployment(self): + setup_trainer_deployment(self) return self diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 8aa104721a5..87ee38b7b84 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -299,7 +299,8 @@ def compute_train_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: print_sample( micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer ) - return compute_loss(micro_batch) + with maybe_record_function("forward"): + return compute_loss(micro_batch) maybe_record_function = nullcontext diff --git a/src/prime_rl/trainer/step.py b/src/prime_rl/trainer/step.py index c94b656536f..7ea93626ab6 100644 --- a/src/prime_rl/trainer/step.py +++ b/src/prime_rl/trainer/step.py @@ -76,8 +76,7 @@ def run_micro_batches( local_num_tokens = 0 for micro_step, micro_batch in enumerate(micro_batches): local_num_tokens += micro_batch["input_ids"].shape[1] - with record_function("forward"): - micro_loss_sum, micro_token_count = compute_loss(micro_batch) + micro_loss_sum, micro_token_count = compute_loss(micro_batch) if gradient_manager is None: local_token_count += micro_token_count if torch.isnan(micro_loss_sum.detach()): diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index f660684c7a1..7dc00a77993 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -319,7 +319,7 @@ def has_checkpoints(output_dir: Path) -> bool: # Launcher artifacts may exist before training starts. Everything else is treated as # artifacts of a previous run. -LAUNCHER_ARTIFACTS = ("configs", "launcher", "model") +LAUNCHER_ARTIFACTS = ("configs", "launcher") def has_run_artifacts(run_dir: Path) -> bool: From a000c2397d67768fd413ffa01849ad0e05883aef Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 20:04:16 +0200 Subject: [PATCH 042/132] revert(trainer): no SFT/RL changes in shared optimizer, perf and glm4_moe code New rule: the pretrain work must not change SFT or RL behavior. - optim.eps leaves the shared AdamW/Muon configs and the shared optimizer setup; pretrain has its own PretrainAdamWConfig/PretrainMuonConfig with eps, which its per-role param groups use. - perf.py: MFU no longer counts glm4_moe's n_shared_experts (back to main). - glm4_moe: the selection bias is no longer zeroed in init_buffers_post_meta; pretrain's init_weights zeroes the MoE selection biases instead. setup_optimizer(param_groups=None) keeps SFT/RL's default groups unchanged. --- .../src/prime_rl/configs/pretrain.py | 23 +++++++++++++++---- .../src/prime_rl/configs/trainer.py | 6 ----- src/prime_rl/trainer/init.py | 7 +++++- .../models/glm4_moe/modeling_glm4_moe.py | 5 ++-- src/prime_rl/trainer/optim/__init__.py | 5 +--- src/prime_rl/trainer/perf.py | 5 ++-- .../trainer/pretrain/parametrization.py | 5 ++-- tests/unit/train/pretrain/test_init.py | 7 +++--- 8 files changed, 34 insertions(+), 29 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index cfa1739b21b..c85291b97f1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -10,13 +10,13 @@ from prime_rl.configs.sft import SFTDeploymentConfig, SingleNodeDeploymentConfig, setup_trainer_deployment from prime_rl.configs.shared import EnvVars, HeartbeatConfig, ResumeConfig, RunConfig, SlurmConfig, TrainerLogConfig from prime_rl.configs.trainer import ( + AdamWConfig, CheckpointConfig, GCConfig, ModelConfig, MoERuntimeConfig, MuonConfig, OffloadMode, - OptimizerConfig, TokenizerConfig, ) from prime_rl.utils.config import BaseConfig, default_output_dir @@ -222,6 +222,19 @@ class ScheduleConfig(BaseConfig): """Shape of the decay from 1 to ``end_ratio``.""" +class PretrainAdamWConfig(AdamWConfig): + eps: float = Field(1e-8, gt=0) + """Adam epsilon of the base shape (scaled per role under CompleteP).""" + + +class PretrainMuonConfig(MuonConfig): + eps: float = Field(1e-8, gt=0) + """Adam epsilon of the base shape for the AdamW groups (scaled per role under CompleteP).""" + + +PretrainOptimizerConfig: TypeAlias = Annotated[PretrainAdamWConfig | PretrainMuonConfig, Field(discriminator="type")] + + class PretrainConfig(BaseConfig): model: PretrainModelConfig @@ -236,7 +249,7 @@ class PretrainConfig(BaseConfig): val: ValConfig | None = None """Held-out evals. If None, none run.""" - optim: OptimizerConfig = MuonConfig() + optim: PretrainOptimizerConfig = PretrainMuonConfig() """Optimizer. ``optim.lr`` is required.""" schedule: ScheduleConfig = ScheduleConfig() @@ -353,9 +366,9 @@ def deepep_disables_grad_clipping(self): return self @model_validator(mode="after") - def validate_param_groups(self): - if self.optim.type not in ("adamw", "muon") or self.model.offload != "none": - raise ValueError("Pretraining's per-role parameter groups need AdamW or Muon and model.offload = 'none'.") + def validate_offload(self): + if self.model.offload != "none": + raise ValueError("Pretraining's per-role parameter groups need model.offload = 'none'.") return self @model_validator(mode="after") diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 0643ff3a498..f7035c4e949 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -419,9 +419,6 @@ class AdamWConfig(BaseOptimizerConfig): betas2: float = Field(0.999, ge=0) """Adam second-moment (β2) decay.""" - eps: float = Field(1e-8, gt=0) - """Adam epsilon.""" - class MuonConfig(BaseOptimizerConfig): type: Literal["muon"] = "muon" @@ -435,9 +432,6 @@ class MuonConfig(BaseOptimizerConfig): betas2: float = Field(0.95, ge=0) """β2 for the AdamW/Lion sub-optimizer used on non-Muon params.""" - eps: float = Field(1e-8, gt=0) - """Adam epsilon for the AdamW sub-optimizer used on non-Muon params.""" - class SignSGDConfig(BaseOptimizerConfig): type: Literal["sign_sgd"] = "sign_sgd" diff --git a/src/prime_rl/trainer/init.py b/src/prime_rl/trainer/init.py index 4ee1608facb..dbd7d83689b 100644 --- a/src/prime_rl/trainer/init.py +++ b/src/prime_rl/trainer/init.py @@ -8,6 +8,7 @@ from torch.distributed.tensor._utils import compute_local_shape_and_global_offset from prime_rl.configs.pretrain import InitConfig +from prime_rl.trainer.model import iter_moe_blocks from prime_rl.utils.weights import resolve_fqn ParamRole = Literal["embedding", "output", "residual_out", "hidden", "norm", "bias"] @@ -78,7 +79,7 @@ def normal_shard( @torch.no_grad() def init_weights(model: nn.Module, config: InitConfig, m_n: float = 1.0, m_l: float = 1.0) -> None: - """Initialize every parameter in place, with the CompleteP multipliers ``m_N``, ``m_L``. Each rank fills only + """Initialize every parameter and the MoE selection biases in place, with the CompleteP multipliers ``m_N``, ``m_L``. Each rank fills only its local shard, without communication, and the values do not depend on the world size or the sharding.""" for name, param in model.named_parameters(): # Canonical names, so wrappers such as activation checkpointing don't change the values. @@ -106,3 +107,7 @@ def init_weights(model: nn.Module, config: InitConfig, m_n: float = 1.0, m_l: fl local.device, ) local.copy_(values.mul_(role_std(role, config.std, m_n, m_l))) + # MoE selection biases are buffers: they start at 0, so routing starts from the router scores alone. + for moe in iter_moe_blocks(model): + if moe.router.selection_bias is not None: + moe.router.selection_bias.zero_() diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py index fcc0af1a201..19cda1c6259 100644 --- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py @@ -339,9 +339,8 @@ def init_buffers_post_meta(self): ) rotary_emb.inv_freq.copy_(inv_freq) - for module in self.modules(): - if isinstance(module, MoE) and module.router.selection_bias is not None: - module.router.selection_bias.zero_() + # TODO: Init TT MoE buffers + # I think .to_empty() on gpu by default fills 0 so we are ok but this might not be guaranteed behavior __all__ = ["Glm4MoeConfig", "Glm4MoePreTrainedModel", "Glm4MoeModel", "Glm4MoeForCausalLM"] diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index a47913ff7bb..3d94d596ee2 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/src/prime_rl/trainer/optim/__init__.py @@ -128,7 +128,6 @@ def _create_optimizer( lr=lr, weight_decay=config.weight_decay, betas=(config.betas1, config.betas2), - eps=config.eps, fused=fused_adamw, ) case "muon": @@ -201,9 +200,7 @@ def muon_enabled(n, p): ) ) - param_groups.append( - dict(params=adamw_params, algorithm="adamw", lr=lr, weight_decay=config.weight_decay, epsilon=config.eps) - ) + param_groups.append(dict(params=adamw_params, algorithm="adamw", lr=lr, weight_decay=config.weight_decay)) return param_groups diff --git a/src/prime_rl/trainer/perf.py b/src/prime_rl/trainer/perf.py index 17b25ff3c0f..13ba65d6b99 100644 --- a/src/prime_rl/trainer/perf.py +++ b/src/prime_rl/trainer/perf.py @@ -156,9 +156,8 @@ def get_active_mm_params(config: PretrainedConfig) -> float: # Some MoE models (e.g. DeepSeek) use moe_intermediate_size, others (e.g. Granite) just use intermediate_size moe_intermediate_size = getattr(config, "moe_intermediate_size", None) or intermediate_size - num_shared_experts = getattr(config, "num_shared_experts", None) or getattr(config, "n_shared_experts", 0) - if num_shared_experts: # Shared experts - sparse_mlp_params += num_sparse_layers * num_shared_experts * 3 * moe_intermediate_size * hidden_size + if hasattr(config, "num_shared_experts") and config.num_shared_experts: # Shared experts + sparse_mlp_params += num_sparse_layers * config.num_shared_experts * 3 * moe_intermediate_size * hidden_size if hasattr(config, "num_experts_per_tok") and config.num_experts_per_tok: # Routed experts sparse_mlp_params += ( num_sparse_layers * config.num_experts_per_tok * 3 * moe_intermediate_size * hidden_size diff --git a/src/prime_rl/trainer/pretrain/parametrization.py b/src/prime_rl/trainer/pretrain/parametrization.py index 345d6295b08..37357f08403 100644 --- a/src/prime_rl/trainer/pretrain/parametrization.py +++ b/src/prime_rl/trainer/pretrain/parametrization.py @@ -2,8 +2,7 @@ from torch import nn -from prime_rl.configs.pretrain import ParametrizationConfig -from prime_rl.configs.trainer import AdamWConfig, MuonConfig +from prime_rl.configs.pretrain import ParametrizationConfig, PretrainAdamWConfig, PretrainMuonConfig from prime_rl.trainer.init import param_role from prime_rl.utils.weights import resolve_fqn @@ -17,7 +16,7 @@ def multipliers(config: ParametrizationConfig | None, model: nn.Module) -> tuple def param_groups( - model: nn.Module, optim: AdamWConfig | MuonConfig, m_n: float, m_l: float, ep_enabled: bool + model: nn.Module, optim: PretrainAdamWConfig | PretrainMuonConfig, m_n: float, m_l: float, ep_enabled: bool ) -> list[dict]: """Optimizer parameter groups per role for CompleteP with alpha = 1 (Dey et al. 2025). ``optim.lr``, ``optim.eps`` and ``optim.weight_decay`` are the values for the base shape; each role scales diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index 9d0fef5588a..61e85997be3 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -5,8 +5,7 @@ from torch.optim import AdamW from transformers import LlamaConfig, LlamaForCausalLM -from prime_rl.configs.pretrain import InitConfig -from prime_rl.configs.trainer import AdamWConfig, MuonConfig +from prime_rl.configs.pretrain import InitConfig, PretrainAdamWConfig, PretrainMuonConfig from prime_rl.trainer.init import init_weights, normal_shard from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM from prime_rl.trainer.pretrain.parametrization import param_groups @@ -42,7 +41,7 @@ def test_completep_folded_multipliers_match_explicit_multipliers(): head_dim=32, tie_word_embeddings=False, ) - optim = AdamWConfig(lr=1e-2, eps=1e-10, weight_decay=0.1) + optim = PretrainAdamWConfig(lr=1e-2, eps=1e-10, weight_decay=0.1) m_n, m_l = 2.0, 2.0 folded = LlamaForCausalLM(config) @@ -107,7 +106,7 @@ def test_completep_moe_groups_cover_experts_and_scale_the_router(): config._attn_implementation = "flash_attention_2" with torch.device("meta"): model = Glm4MoeForCausalLM._from_config(config) - groups = param_groups(model, MuonConfig(lr=1.0, weight_decay=0.1), m_n=2.0, m_l=1.0, ep_enabled=False) + groups = param_groups(model, PretrainMuonConfig(lr=1.0, weight_decay=0.1), m_n=2.0, m_l=1.0, ep_enabled=False) lr = {param: group["lr"] for group in groups for param in group["params"]} assert sum(len(group["params"]) for group in groups) == len(lr) == len(list(model.parameters())) moe = model.model.layers[1].mlp From 6d478fc0168fb01687dd0ff5546f86cf7e670322 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 20:06:21 +0200 Subject: [PATCH 043/132] refactor: move the trainer deployment configs and the stateful dataset helpers to neutral modules Pure moves, so pretrain doesn't import from SFT modules: the single-/multi-node trainer deployment configs and setup_trainer_deployment go from configs/sft.py to configs/shared.py, and StatefulIterableDataset, get_dataset_state and get_dataset_progress go from trainer/sft/data.py to trainer/data.py. Names and code are unchanged. --- .../src/prime_rl/configs/pretrain.py | 13 +++- .../src/prime_rl/configs/sft.py | 76 ++----------------- .../src/prime_rl/configs/shared.py | 74 +++++++++++++++++- src/prime_rl/trainer/data.py | 73 ++++++++++++++++++ src/prime_rl/trainer/pretrain/data.py | 2 +- src/prime_rl/trainer/pretrain/train.py | 2 +- src/prime_rl/trainer/sft/data.py | 70 +---------------- src/prime_rl/trainer/sft/train.py | 3 +- tests/unit/train/sft/test_dataloader.py | 3 +- 9 files changed, 167 insertions(+), 149 deletions(-) create mode 100644 src/prime_rl/trainer/data.py diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 3d2d4c05e64..3e2e686f2f6 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -7,8 +7,17 @@ from pydantic import Field, model_validator from prime_rl.configs.monitors import TrainMonitorsConfig -from prime_rl.configs.sft import SFTDeploymentConfig, SingleNodeDeploymentConfig, setup_trainer_deployment -from prime_rl.configs.shared import EnvVars, HeartbeatConfig, ResumeConfig, RunConfig, SlurmConfig, TrainerLogConfig +from prime_rl.configs.shared import ( + EnvVars, + HeartbeatConfig, + ResumeConfig, + RunConfig, + SFTDeploymentConfig, + SingleNodeDeploymentConfig, + SlurmConfig, + TrainerLogConfig, + setup_trainer_deployment, +) from prime_rl.configs.trainer import ( CheckpointConfig, GCConfig, diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 682974a3770..88866b04003 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -1,10 +1,10 @@ import uuid import warnings from pathlib import Path -from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias +from typing import Annotated, Literal, TypeAlias from urllib.parse import urlparse -from pydantic import AliasChoices, Field, model_validator +from pydantic import Field, model_validator from renderers import AutoRendererConfig, DefaultRendererConfig, RendererConfig from renderers.base import MODEL_RENDERER_MAP @@ -19,9 +19,12 @@ NCCLWeightBroadcastConfig, ResumeConfig, RunConfig, + SFTDeploymentConfig, + SingleNodeDeploymentConfig, SlurmConfig, TrainerLogConfig, WeightBroadcastConfig, + setup_trainer_deployment, ) from prime_rl.configs.trainer import ( AdamWConfig, @@ -34,10 +37,7 @@ TokenizerConfig, validate_scheduler, ) -from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource - -if TYPE_CHECKING: - from prime_rl.configs.pretrain import PretrainConfig +from prime_rl.utils.config import BaseConfig, default_output_dir class BaseDataConfig(BaseConfig): @@ -179,70 +179,6 @@ class SFTValConfig(BaseConfig): DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig, Field(discriminator="type")] -class BaseDeploymentConfig(BaseConfig): - gpus_per_node: int = 8 - """GPUs per node.""" - - -class SingleNodeDeploymentConfig(BaseDeploymentConfig): - type: Literal["single_node"] = "single_node" - - num_train_gpus: int = 1 - """GPUs allocated to the trainer.""" - - num_infer_gpus: int = Field(1, validation_alias=AliasChoices("num_infer_gpus", "num_eval_gpus")) - """GPUs allocated to inference for online evals (alias: ``num_eval_gpus``). Only used when an ``[inference]`` block is configured.""" - - @model_validator(mode="after") - def validate_gpu_count(self): - if self.num_train_gpus > self.gpus_per_node: - raise ValueError(f"num_train_gpus ({self.num_train_gpus}) exceeds gpus_per_node ({self.gpus_per_node}).") - return self - - -class MultiNodeDeploymentConfig(BaseDeploymentConfig): - type: Literal["multi_node"] = "multi_node" - - num_train_nodes: int = Field(2, ge=1) - """Training nodes.""" - - num_infer_nodes: int = Field(0, ge=0, validation_alias=AliasChoices("num_infer_nodes", "num_eval_nodes")) - """Inference nodes for online evals (alias: ``num_eval_nodes``). These nodes share - one SLURM allocation with the trainer nodes.""" - - nodes_per_fsdp_group: int | None = None - """Nodes per FSDP island. Auto-sets ``model.dp_replicate = num_train_nodes / nodes_per_fsdp_group``.""" - - -SFTDeploymentConfig: TypeAlias = Annotated[ - SingleNodeDeploymentConfig | MultiNodeDeploymentConfig, Field(discriminator="type") -] - - -def setup_trainer_deployment(config: "SFTConfig | PretrainConfig") -> None: - """Shared by ``sft`` and ``pretrain``: a multi-node deployment needs SLURM and sets - ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the SLURM job name defaults to the - run name and the template to the bundled single- or multi-node trainer template.""" - deployment = config.deployment - if deployment.type == "multi_node": - if config.slurm is None: - raise ValueError("Must use SLURM for multi-node deployment.") - if deployment.nodes_per_fsdp_group is not None: - if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: - raise ValueError( - f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " - f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" - ) - config.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group - if config.slurm is not None: - if config.slurm.job_name is None: - config.slurm.job_name = config.run.name - if config.slurm.template_path is None: - templates_dir = find_package_resource("templates") - if templates_dir is not None: - config.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" - - class SFTConfig(BaseConfig): model: ModelConfig = ModelConfig() diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index d9dd63b9e01..52b2883207d 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -1,12 +1,16 @@ import os import re from pathlib import Path -from typing import Annotated, Literal, TypeAlias +from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias -from pydantic import AfterValidator, Field, model_validator +from pydantic import AfterValidator, AliasChoices, Field, model_validator from verifiers.v1.configs.client import BaseClientConfig -from prime_rl.utils.config import BaseConfig +from prime_rl.utils.config import BaseConfig, find_package_resource + +if TYPE_CHECKING: + from prime_rl.configs.pretrain import PretrainConfig + from prime_rl.configs.sft import SFTConfig # Launcher-managed env vars that a component's `env_vars` must not set: GPU partitioning # and the single shared W&B run. The launcher always sets these last, so allowing them in @@ -293,3 +297,67 @@ class ZMQTransportConfig(BaseTransportConfig): TransportConfig: TypeAlias = Annotated[FileSystemTransportConfig | ZMQTransportConfig, Field(discriminator="type")] + + +class BaseDeploymentConfig(BaseConfig): + gpus_per_node: int = 8 + """GPUs per node.""" + + +class SingleNodeDeploymentConfig(BaseDeploymentConfig): + type: Literal["single_node"] = "single_node" + + num_train_gpus: int = 1 + """GPUs allocated to the trainer.""" + + num_infer_gpus: int = Field(1, validation_alias=AliasChoices("num_infer_gpus", "num_eval_gpus")) + """GPUs allocated to inference for online evals (alias: ``num_eval_gpus``). Only used when an ``[inference]`` block is configured.""" + + @model_validator(mode="after") + def validate_gpu_count(self): + if self.num_train_gpus > self.gpus_per_node: + raise ValueError(f"num_train_gpus ({self.num_train_gpus}) exceeds gpus_per_node ({self.gpus_per_node}).") + return self + + +class MultiNodeDeploymentConfig(BaseDeploymentConfig): + type: Literal["multi_node"] = "multi_node" + + num_train_nodes: int = Field(2, ge=1) + """Training nodes.""" + + num_infer_nodes: int = Field(0, ge=0, validation_alias=AliasChoices("num_infer_nodes", "num_eval_nodes")) + """Inference nodes for online evals (alias: ``num_eval_nodes``). These nodes share + one SLURM allocation with the trainer nodes.""" + + nodes_per_fsdp_group: int | None = None + """Nodes per FSDP island. Auto-sets ``model.dp_replicate = num_train_nodes / nodes_per_fsdp_group``.""" + + +SFTDeploymentConfig: TypeAlias = Annotated[ + SingleNodeDeploymentConfig | MultiNodeDeploymentConfig, Field(discriminator="type") +] + + +def setup_trainer_deployment(config: "SFTConfig | PretrainConfig") -> None: + """Shared by ``sft`` and ``pretrain``: a multi-node deployment needs SLURM and sets + ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the SLURM job name defaults to the + run name and the template to the bundled single- or multi-node trainer template.""" + deployment = config.deployment + if deployment.type == "multi_node": + if config.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + if deployment.nodes_per_fsdp_group is not None: + if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" + ) + config.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group + if config.slurm is not None: + if config.slurm.job_name is None: + config.slurm.job_name = config.run.name + if config.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + config.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" diff --git a/src/prime_rl/trainer/data.py b/src/prime_rl/trainer/data.py new file mode 100644 index 00000000000..dd0e725b084 --- /dev/null +++ b/src/prime_rl/trainer/data.py @@ -0,0 +1,73 @@ +from collections import defaultdict + +from torch.distributed.checkpoint.stateful import Stateful +from torch.utils.data import IterableDataset, get_worker_info +from torchdata.stateful_dataloader import StatefulDataLoader + +from prime_rl.trainer.world import get_world + + +class StatefulIterableDataset(Stateful, IterableDataset): + """SFT dataset are iterable (infinite) and stateful (can be checkpointed).""" + + def __init__(self, non_dp_size: int = 1): + self.step, self.epoch = 0, 0 + self.num_samples = defaultdict(int) + self.num_tokens = defaultdict(int) + self.fast_forward = False + self.non_dp_size = non_dp_size + self._setup_world_info() + + def state_dict(self) -> dict: + return {"step": self.step, "epoch": self.epoch} + + def load_state_dict(self, state_dict: dict): + assert "step" in state_dict and "epoch" in state_dict + self.fast_forward = True + self.step = state_dict["step"] + self.epoch = state_dict["epoch"] + + def _setup_world_info(self): + worker_info = get_worker_info() + if worker_info is not None: + worker_id = worker_info.id + num_workers = worker_info.num_workers + else: + worker_id, num_workers = 0, 1 + world = get_world() + assert world.world_size % self.non_dp_size == 0, "world_size must be divisible by non_dp_size" + self.data_rank = world.rank // self.non_dp_size * num_workers + worker_id + self.data_world_size = world.world_size // self.non_dp_size * num_workers + + +def get_dataset_state(dataloader: StatefulDataLoader) -> dict: + """Dataset position per worker for the startup log, parsed from ``StatefulDataLoader.state_dict()``. + + The loader is the only source that is correct in every case: after a resume's + ``load_state_dict`` the restored position exists solely in the loader's stashed + state (it reaches the dataset copies inside workers when the iterator forks them; + the main-process dataset object stays at position zero). The keys are torchdata's + private worker-snapshot layout.""" + snapshots = dataloader.state_dict()["_snapshot"]["_worker_snapshots"] + return {wid: snap["dataset_state"]["dataset"] for wid, snap in sorted(snapshots.items())} + + +def get_dataset_progress(dataloader: StatefulDataLoader) -> dict: + """Dataset position and aggregate counters from dataloader workers.""" + snapshot = dataloader.state_dict()["_snapshot"] + worker_snapshots = snapshot["_worker_snapshots"] + positions = [worker_snapshot["dataset_state"]["dataset"] for worker_snapshot in worker_snapshots.values()] + furthest = max(positions, key=lambda position: position["step"]) + num_samples = defaultdict(int) + num_tokens = defaultdict(int) + for worker_snapshot in worker_snapshots.values(): + progress = worker_snapshot["dataset_state"].get("progress", {}) + for name, count in progress.get("num_samples", {}).items(): + num_samples[name] += count + for name, count in progress.get("num_tokens", {}).items(): + num_tokens[name] += count + return { + **furthest, + "num_samples": dict(num_samples), + "num_tokens": dict(num_tokens), + } diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index 2ca6e4bb4dc..fac27e4a289 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -6,7 +6,7 @@ from torchdata.stateful_dataloader import StatefulDataLoader from prime_rl.configs.pretrain import MixturePhaseConfig, PretrainDataConfig, TokenSourceConfig -from prime_rl.trainer.sft.data import StatefulIterableDataset +from prime_rl.trainer.data import StatefulIterableDataset from prime_rl.utils.logger import get_logger diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 4d680c49f83..0f896487cb2 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -40,7 +40,7 @@ setup_val_dataset, ) from prime_rl.trainer.pretrain.schedules import LRSchedule -from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state +from prime_rl.trainer.data import get_dataset_progress, get_dataset_state from prime_rl.trainer.step import run_micro_batches from prime_rl.trainer.utils import ( GarbageCollection, diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index 6fd0a5dc7da..d8bb479faa8 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -13,13 +13,11 @@ from renderers import AutoRendererConfig, RendererConfig, merge_chat_template_kwargs from renderers.base import MultiModalData, PlaceholderRange, Renderer, build_training_sample, create_renderer from torch import Tensor -from torch.distributed.checkpoint.stateful import Stateful -from torch.utils.data import IterableDataset, get_worker_info from torchdata.stateful_dataloader import StatefulDataLoader from transformers.tokenization_utils import PreTrainedTokenizer from prime_rl.configs.sft import DataConfig, LossMaskConfig, SFTColumnsConfig, SFTDataConfig -from prime_rl.trainer.world import get_world +from prime_rl.trainer.data import StatefulIterableDataset from prime_rl.utils.chat_template import deserialize_tool_calls, normalize_messages from prime_rl.utils.logger import get_logger from prime_rl.utils.utils import format_time @@ -45,39 +43,6 @@ class Batch(TypedDict): mm_token_type_ids: Int[Tensor, "batch seq"] | None -class StatefulIterableDataset(Stateful, IterableDataset): - """SFT dataset are iterable (infinite) and stateful (can be checkpointed).""" - - def __init__(self, non_dp_size: int = 1): - self.step, self.epoch = 0, 0 - self.num_samples = defaultdict(int) - self.num_tokens = defaultdict(int) - self.fast_forward = False - self.non_dp_size = non_dp_size - self._setup_world_info() - - def state_dict(self) -> dict: - return {"step": self.step, "epoch": self.epoch} - - def load_state_dict(self, state_dict: dict): - assert "step" in state_dict and "epoch" in state_dict - self.fast_forward = True - self.step = state_dict["step"] - self.epoch = state_dict["epoch"] - - def _setup_world_info(self): - worker_info = get_worker_info() - if worker_info is not None: - worker_id = worker_info.id - num_workers = worker_info.num_workers - else: - worker_id, num_workers = 0, 1 - world = get_world() - assert world.world_size % self.non_dp_size == 0, "world_size must be divisible by non_dp_size" - self.data_rank = world.rank // self.non_dp_size * num_workers + worker_id - self.data_world_size = world.world_size // self.non_dp_size * num_workers - - class FakeDataset(StatefulIterableDataset): """A dataset of fake tokens""" @@ -769,36 +734,3 @@ def setup_dataloader(dataset: StatefulIterableDataset, config: DataConfig, cp: i num_workers=config.num_workers, pin_memory=True, ) - - -def get_dataset_state(dataloader: StatefulDataLoader) -> dict: - """Dataset position per worker for the startup log, parsed from ``StatefulDataLoader.state_dict()``. - - The loader is the only source that is correct in every case: after a resume's - ``load_state_dict`` the restored position exists solely in the loader's stashed - state (it reaches the dataset copies inside workers when the iterator forks them; - the main-process dataset object stays at position zero). The keys are torchdata's - private worker-snapshot layout.""" - snapshots = dataloader.state_dict()["_snapshot"]["_worker_snapshots"] - return {wid: snap["dataset_state"]["dataset"] for wid, snap in sorted(snapshots.items())} - - -def get_dataset_progress(dataloader: StatefulDataLoader) -> dict: - """Dataset position and aggregate counters from dataloader workers.""" - snapshot = dataloader.state_dict()["_snapshot"] - worker_snapshots = snapshot["_worker_snapshots"] - positions = [worker_snapshot["dataset_state"]["dataset"] for worker_snapshot in worker_snapshots.values()] - furthest = max(positions, key=lambda position: position["step"]) - num_samples = defaultdict(int) - num_tokens = defaultdict(int) - for worker_snapshot in worker_snapshots.values(): - progress = worker_snapshot["dataset_state"].get("progress", {}) - for name, count in progress.get("num_samples", {}).items(): - num_samples[name] += count - for name, count in progress.get("num_tokens", {}).items(): - num_tokens[name] += count - return { - **furthest, - "num_samples": dict(num_samples), - "num_tokens": dict(num_tokens), - } diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 87ee38b7b84..234074d354a 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -37,9 +37,8 @@ ) from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter +from prime_rl.trainer.data import get_dataset_progress, get_dataset_state from prime_rl.trainer.sft.data import ( - get_dataset_progress, - get_dataset_state, load_sft_dataset, setup_dataloader, setup_dataset, diff --git a/tests/unit/train/sft/test_dataloader.py b/tests/unit/train/sft/test_dataloader.py index 7f2c7d47839..9615cf05104 100644 --- a/tests/unit/train/sft/test_dataloader.py +++ b/tests/unit/train/sft/test_dataloader.py @@ -6,7 +6,8 @@ from datasets import Dataset from prime_rl.configs.sft import FakeDataConfig, SFTDataConfig -from prime_rl.trainer.sft.data import FakeDataset, SFTDataset, get_dataset_progress, get_dataset_state, setup_dataloader +from prime_rl.trainer.data import get_dataset_progress, get_dataset_state +from prime_rl.trainer.sft.data import FakeDataset, SFTDataset, setup_dataloader from prime_rl.trainer.world import reset_world From ee0f8894c6e119d1eecfe5ba835201b5cd1d4496 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 20:11:53 +0200 Subject: [PATCH 044/132] feat(pretrain): platform monitor, pretrain overview in W&B and the dashboard monitors.prime is accepted for pretraining: the trainer registers a platform run (no environments, the architecture as model) and the eval worker attaches to it. The W&B overview view and the dashboard get a pretrain flavor (loss and held-out bpb, stability incl. grad norms per group, MoE balance, performance); the dashboard recognizes runs by their pretrain.json. SFT and RL views are unchanged. --- .../src/prime_rl/configs/pretrain.py | 6 --- src/prime_rl/dashboard/server.py | 4 +- src/prime_rl/dashboard/static/app.js | 12 ++++++ src/prime_rl/monitors/__init__.py | 2 +- src/prime_rl/monitors/prime.py | 15 ++++++++ src/prime_rl/monitors/wandb/monitor.py | 2 +- src/prime_rl/monitors/wandb/overview.py | 38 ++++++++++++++++++- src/prime_rl/trainer/pretrain/eval_worker.py | 9 ++++- src/prime_rl/trainer/pretrain/train.py | 4 +- 9 files changed, 78 insertions(+), 14 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 7efd740187d..670dc62f288 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -324,12 +324,6 @@ def validate_param_groups(self): raise ValueError("Pretraining's per-role parameter groups need AdamW or Muon and model.offload = 'none'.") return self - @model_validator(mode="after") - def validate_monitors(self): - if self.monitors.prime is not None: - raise ValueError("monitors.prime is not supported for pretraining yet.") - return self - @model_validator(mode="after") def validate_cp(self): if self.model.cp > 1: diff --git a/src/prime_rl/dashboard/server.py b/src/prime_rl/dashboard/server.py index 5b8d0861e1d..8129ede31a7 100755 --- a/src/prime_rl/dashboard/server.py +++ b/src/prime_rl/dashboard/server.py @@ -246,6 +246,8 @@ def main_config(run_dir: Path) -> tuple[str, dict]: return "rl", read_json(configs / "orchestrator.json") or read_json(configs / "trainer.json") if (configs / "eval.json").exists(): return "eval", read_json(configs / "eval.json") + if (configs / "pretrain.json").exists(): + return "pretrain", read_json(configs / "pretrain.json") return "other", {} @@ -471,7 +473,7 @@ def read_log(run: str, file: str, start: int | None = None, end: int | None = No # ------------------------------------------------------------------------- configs -CONFIG_ORDER = ["rl", "sft", "eval", "orchestrator", "trainer", "inference"] +CONFIG_ORDER = ["rl", "sft", "eval", "pretrain", "orchestrator", "trainer", "inference"] def config_rank(name: str) -> tuple[int, str]: diff --git a/src/prime_rl/dashboard/static/app.js b/src/prime_rl/dashboard/static/app.js index f63243df761..c08f30e5d74 100644 --- a/src/prime_rl/dashboard/static/app.js +++ b/src/prime_rl/dashboard/static/app.js @@ -446,6 +446,11 @@ const SFT_TRAIN_METRICS = ["loss/mean", "loss/perplexity", "val/loss", "val/perp // max_vio only exists on MoE models: a regex panel stays away on dense ones const SFT_STABILITY_PANELS = [{ metric: "optim/grad_norm" }, { metric: "optim/lr" }, { metric: "loss/nan_count" }, { regex: "max_vio/(mean|max)" }]; const SFT_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memory", "time/forward_backward"]; +// Pretrain flavor (overview.py PRETRAIN_*): held-out bpb from the trainer (val/) and the eval worker (eval/) +const PRETRAIN_TRAIN_PANELS = [{ metric: "loss/mean" }, { metric: "loss/spike" }, { regex: "val/.*/bpb" }, { regex: "eval/.*/bpb" }]; +const PRETRAIN_STABILITY_PANELS = [{ metric: "optim/grad_norm" }, { metric: "optim/grad_norm_spike" }, { metric: "optim/lr" }, { metric: "loss/nan_count" }, { regex: "optim/grad_norm/.*" }]; +const PRETRAIN_MOE_METRICS = ["expert_load/max_mean/mean", "expert_load/max_mean/max", "expert_load/micro_max_mean/mean", "expert_load/micro_max_mean/max", "expert_load/entropy/min", "expert_load/dead", "router/confidence", "router/bias/absmax", "loss/seq_balance"]; +const PRETRAIN_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memory/max", "time/step", "time/data/max", "time/eval_round"]; // Multi-series inference panels (overview.py INFERENCE_PANELS): fleet aggregate // paired with the cross-engine tail that flags a single sick engine. @@ -1455,6 +1460,13 @@ function buildSections(meta) { }); const sections = []; const evalEnvs = meta.eval_envs || []; + if (meta.type === "pretrain") { + sections.push({ name: "train", panels: PRETRAIN_TRAIN_PANELS }); + sections.push({ name: "stability", panels: PRETRAIN_STABILITY_PANELS }); + sections.push({ name: "moe", panels: PRETRAIN_MOE_METRICS.map((m) => ({ metric: m })) }); + sections.push({ name: "performance", panels: PRETRAIN_PERFORMANCE_METRICS.map((m) => ({ metric: m })) }); + return sections; + } if (meta.type === "sft") { // the val panes show once a validation value is logged, not on the config alone sections.push({ name: "train", panels: SFT_TRAIN_METRICS.map((m) => (m.startsWith("val/") ? { regex: escRe(m) } : { metric: m })) }); diff --git a/src/prime_rl/monitors/__init__.py b/src/prime_rl/monitors/__init__.py index ad08e113091..a3775c048ee 100644 --- a/src/prime_rl/monitors/__init__.py +++ b/src/prime_rl/monitors/__init__.py @@ -55,7 +55,7 @@ async def setup( run_config: BaseConfig | None = None, train_env_names: list[str] | None = None, eval_env_names: list[str] | None = None, - overview_flavor: Literal["rl", "sft", "eval"] = "rl", + overview_flavor: Literal["rl", "sft", "eval", "pretrain"] = "rl", ) -> None: """Construct, initialize, and register one monitor per non-None config. diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 9558db6efe5..313d0e9a3a6 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -14,6 +14,7 @@ from prime_sandboxes import Config as PrimeConfig from prime_rl.configs.monitors import PrimeEvalMonitorConfig, PrimeTrainMonitorConfig +from prime_rl.configs.pretrain import PretrainConfig from prime_rl.configs.sft import SFTConfig from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.utils.config import BaseConfig @@ -171,6 +172,20 @@ async def init(self, config: BaseConfig | None = None, output_dir: Path | None = ), config=config.model_dump(exclude_none=True, mode="json"), ) + elif isinstance(config, PretrainConfig): + # Pretraining trains on token data from scratch: no environments, the model is its architecture. + init_kwargs = dict( + name=self.config.name, + model=config.model.arch["model_type"], + environments=[], + training=pr.TrainingSpec( + max_steps=config.max_steps, + batch_size=config.data.batch_size, + seq_len=config.data.seq_len, + wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + ), + config=config.model_dump(exclude_none=True, mode="json"), + ) elif config is not None: init_kwargs = dict( name=self.config.name, diff --git a/src/prime_rl/monitors/wandb/monitor.py b/src/prime_rl/monitors/wandb/monitor.py index a3b731763df..5a6d9efd331 100644 --- a/src/prime_rl/monitors/wandb/monitor.py +++ b/src/prime_rl/monitors/wandb/monitor.py @@ -33,7 +33,7 @@ async def init( config: BaseConfig | None = None, train_env_names: list[str] | None = None, eval_env_names: list[str] | None = None, - overview_flavor: Literal["rl", "sft", "eval"] = "rl", + overview_flavor: Literal["rl", "sft", "eval", "pretrain"] = "rl", ) -> None: # W&B reads the start command off sys.argv; the launcher passes the original # command to subprocesses via $WANDB_ARGS. diff --git a/src/prime_rl/monitors/wandb/overview.py b/src/prime_rl/monitors/wandb/overview.py index 83012193bbf..6962828bb29 100644 --- a/src/prime_rl/monitors/wandb/overview.py +++ b/src/prime_rl/monitors/wandb/overview.py @@ -91,6 +91,31 @@ "time/save_ckpt", ] +# Pretrain flavor: loss and held-out evals, training stability, MoE balance, performance. +PRETRAIN_TRAIN_METRICS = ["loss/mean", "loss/spike"] +PRETRAIN_TRAIN_REGEXES = ["val/.*/bpb", "eval/.*/bpb"] +PRETRAIN_STABILITY_METRICS = ["optim/grad_norm", "optim/grad_norm_spike", "optim/lr", "loss/nan_count"] +PRETRAIN_STABILITY_REGEXES = ["optim/grad_norm/.*"] +PRETRAIN_MOE_METRICS = [ + "expert_load/max_mean/mean", + "expert_load/max_mean/max", + "expert_load/micro_max_mean/mean", + "expert_load/micro_max_mean/max", + "expert_load/entropy/min", + "expert_load/dead", + "router/confidence", + "router/bias/absmax", + "loss/seq_balance", +] +PRETRAIN_PERFORMANCE_METRICS = [ + "perf/mfu", + "perf/throughput", + "perf/peak_memory/max", + "time/step", + "time/data/max", + "time/eval_round", +] + # Dense grid: more, smaller panels per row and enough rows that sections don't paginate. COLUMNS = 4 ROWS = 6 @@ -150,8 +175,17 @@ def eval_section(name: str, env_pattern: str) -> ws.Section: def build_sections( - train_envs: Sequence[str] = (), eval_envs: Sequence[str] = (), flavor: Literal["rl", "sft", "eval"] = "rl" + train_envs: Sequence[str] = (), + eval_envs: Sequence[str] = (), + flavor: Literal["rl", "sft", "eval", "pretrain"] = "rl", ) -> list[ws.Section]: + if flavor == "pretrain": + return [ + section("train", metrics=PRETRAIN_TRAIN_METRICS, regexes=PRETRAIN_TRAIN_REGEXES), + section("stability", metrics=PRETRAIN_STABILITY_METRICS, regexes=PRETRAIN_STABILITY_REGEXES), + section("moe", metrics=PRETRAIN_MOE_METRICS), + section("performance", metrics=PRETRAIN_PERFORMANCE_METRICS), + ] # A standalone eval run has no training signal: one section per eval env. if flavor == "eval": if eval_envs: @@ -231,7 +265,7 @@ def ensure_overview_view( name: str = OVERVIEW_NAME, train_envs: Sequence[str] = (), eval_envs: Sequence[str] = (), - flavor: Literal["rl", "sft", "eval"] = "rl", + flavor: Literal["rl", "sft", "eval", "pretrain"] = "rl", ) -> str | None: """Ensure an overview saved view exists for this run's env set. Reuses an existing overview built for the same envs; when the env set is new, creates a fresh versioned view (``overview`` → diff --git a/src/prime_rl/trainer/pretrain/eval_worker.py b/src/prime_rl/trainer/pretrain/eval_worker.py index 4cefce335fd..7446d65a8e3 100644 --- a/src/prime_rl/trainer/pretrain/eval_worker.py +++ b/src/prime_rl/trainer/pretrain/eval_worker.py @@ -22,6 +22,7 @@ from prime_rl import monitors from prime_rl.configs.pretrain import EvalWorkerConfig, PretrainConfig +from prime_rl.monitors.prime import read_platform_record from prime_rl.trainer.ckpt import AppState from prime_rl.trainer.model import setup_model from prime_rl.trainer.parallel_dims import get_parallel_dims @@ -69,14 +70,20 @@ def main_loop(config: EvalWorkerConfig) -> None: torch.set_float32_matmul_precision(run.matmul_precision) os.environ.update(shared_wandb_env(config.run_dir, "eval")) + if run.monitors.prime is not None: + # Attach to the trainer's platform run, whose id the trainer's monitor leaves in the run dir. + while (record := read_platform_record(config.run_dir)) is None or "id" not in record: + time.sleep(config.poll_seconds) + os.environ["RUN_ID"] = record["id"] asyncio.run( monitors.setup( producer="eval", wandb=run.monitors.wandb, + prime=run.monitors.prime, file=run.monitors.file, output_dir=config.run_dir, run_config=run, - overview_flavor="sft", + overview_flavor="pretrain", ) ) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index af9883eef6a..9242868d19d 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -98,11 +98,11 @@ def train(config: PretrainConfig): monitors.setup( producer="trainer", wandb=config.monitors.wandb, - prime=None, + prime=config.monitors.prime, file=config.monitors.file, output_dir=config.run_dir, run_config=config, - overview_flavor="sft", + overview_flavor="pretrain", ) ) From 6e453d79556c8fe77a4ca6516d35db68a3937833 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 20:22:36 +0200 Subject: [PATCH 045/132] feat(pretrain): the launcher starts the eval worker; evals at several context lengths [eval_worker] reserves GPUs (single node) or nodes (multi-node, after the trainer nodes in the shared template; SFT renders unchanged) and starts the worker with a written eval_worker.json. The worker scores every set at each of seq_lens and logs eval//@; requests that are the same at several lengths are scored once. --- docs/pretraining.md | 14 ++--- .../src/prime_rl/configs/pretrain.py | 51 +++++++++++++++---- src/prime_rl/entrypoints/pretrain.py | 50 ++++++++++++++++-- .../templates/multi_node_sft.sbatch.j2 | 42 +++++++++++++-- src/prime_rl/trainer/pretrain/eval_worker.py | 44 ++++++++++------ src/prime_rl/trainer/pretrain/suite.py | 30 +++++++++++ tests/unit/train/pretrain/test_launcher.py | 30 ++++++++++- tests/unit/train/pretrain/test_suite.py | 23 ++++++++- 8 files changed, 245 insertions(+), 39 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index d5ea7d23ed1..b499a8e160b 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -127,19 +127,21 @@ interval = 500 web = { path = "data/web-val" } ``` -**Asynchronous evals on dedicated GPUs.** With `[eval_snapshots]`, the trainer writes a model-only snapshot every `interval` steps and at the last step to `/snapshots/step_N`, in the background (it never waits for the eval side). The eval worker runs on its own GPUs (e.g. one node next to the training nodes), takes the newest finished snapshot, scores it on the eval suite (`ppl`, `ctxcont` and `pairs` sets in `//{test.parquet,meta.json}`), logs `eval/{set}/{metric}` at the snapshot's step, deletes that snapshot and every older one, and repeats until it has scored the last step. The eval frequency therefore follows from the cost of a round and the number of eval GPUs. +**Asynchronous evals on dedicated GPUs.** With `[eval_snapshots]`, the trainer writes a model-only snapshot (bf16, router in fp32) every `interval` steps and at the last step to `/snapshots/step_N`, in the background: it never waits for the eval side. With `[eval_worker]`, the launcher starts the eval worker on GPUs of its own (`num_gpus` of the node in a single-node run, `num_nodes` extra nodes in a multi-node run). The worker takes the newest finished snapshot, scores it on the eval suite (`ppl`, `ctxcont` and `pairs` sets in `//{test.parquet,meta.json}`) at every context length in `seq_lens`, logs `eval/{set}/{metric}@{length}` (e.g. `eval/holdout/bpb@4k`) at the snapshot's step, deletes that snapshot and every older one, and repeats until it has scored the last step. The eval frequency therefore follows from the cost of a round and the number of eval GPUs. Requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. ```toml [eval_snapshots] interval = 50 -``` -```bash -torchrun --nproc-per-node 8 -m prime_rl.trainer.pretrain.eval_worker \ - --run-dir / --data-dir ~/evals-data --sets fwedu-holdout lambada math +[eval_worker] +data_dir = "/shared/evals-data" +sets = ["fwedu-holdout", "lambada", "math"] +seq_lens = [2048, 4096, 16384] +num_gpus = 1 # single-node runs +num_nodes = 1 # multi-node runs ``` -The trainer and the worker log into one W&B run (W&B shared mode): the launcher writes the run id to `/wandb_run_id`, and both processes read it. +The trainer and the worker log into one W&B run (W&B shared mode, run id in `/wandb_run_id`) and, with `[monitors.prime]`, into one platform run. The worker can also be started by hand on any GPUs: `torchrun --nproc-per-node 8 -m prime_rl.trainer.pretrain.eval_worker --run-dir --data-dir --sets `. ## Multi-node diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 670dc62f288..a3faae72532 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -227,12 +227,8 @@ class EvalSnapshotsConfig(BaseConfig): """Write a snapshot every N steps, and at the final step.""" -class EvalWorkerConfig(BaseConfig): - """The asynchronous eval worker: scores the newest snapshot of a training run on the eval suite, logs the - results into the run's W&B run and ``metrics.jsonl`` at the snapshot's step, and deletes scored snapshots.""" - - run_dir: Path - """The training run's directory (with ``configs/latest/resolved/pretrain.json`` and ``snapshots/``).""" +class EvalSuiteConfig(BaseConfig): + """What the asynchronous eval worker scores.""" data_dir: Path """Eval sets, one directory per set with ``test.parquet`` and ``meta.json``.""" @@ -240,11 +236,13 @@ class EvalWorkerConfig(BaseConfig): sets: list[str] """Names of the sets to score.""" - seq_len: int = Field(2048, ge=1) - """Window length for ``ppl`` sets and the context limit for ``ctxcont`` sets.""" + seq_lens: list[int] = [2048] + """Context lengths to score every set at: the window length for ``ppl`` sets and the context limit for + ``ctxcont`` sets. Metrics are named ``eval//@`` (e.g. ``bpb@4k``). Requests that are the + same at several lengths are scored once.""" token_budget: int = Field(32768, ge=1) - """Tokens per packed row (one forward).""" + """Tokens per packed row (one forward). At least the longest context length plus one.""" poll_seconds: int = Field(30, ge=1) """Seconds to wait before looking for a new snapshot again.""" @@ -252,9 +250,34 @@ class EvalWorkerConfig(BaseConfig): save_arrays: bool = True """Write per-document arrays (npz) and summaries to ``/evals/step_N/``.""" + @model_validator(mode="after") + def validate_token_budget(self): + if self.token_budget <= max(self.seq_lens): + raise ValueError("token_budget must exceed the longest of seq_lens") + return self + + +class EvalWorkerConfig(EvalSuiteConfig): + """The asynchronous eval worker (``python -m prime_rl.trainer.pretrain.eval_worker``): scores the newest snapshot + of a training run on the eval suite, logs the results into the run's monitors at the snapshot's step, and deletes + scored snapshots. The pretrain launcher writes this config and starts the worker for ``[eval_worker]``.""" + + run_dir: Path + """The training run's directory (with ``configs/latest/resolved/pretrain.json`` and ``snapshots/``).""" + log: LogConfig = LogConfig() +class PretrainEvalWorkerConfig(EvalSuiteConfig): + """The eval worker, started by the launcher next to the trainer on GPUs of its own.""" + + num_gpus: int = Field(1, ge=1) + """Single-node runs: GPUs of the node for the worker, next to the trainer's ``deployment.num_train_gpus``.""" + + num_nodes: int = Field(1, ge=1) + """Multi-node runs: nodes for the worker, next to the trainer's ``deployment.num_train_nodes``.""" + + class PretrainConfig(TrainRunConfig): model: PretrainModelConfig @@ -269,6 +292,10 @@ class PretrainConfig(TrainRunConfig): eval_snapshots: EvalSnapshotsConfig | None = None """Model-only snapshots for the asynchronous eval worker. If None, none are written.""" + eval_worker: PretrainEvalWorkerConfig | None = None + """The asynchronous eval worker, started by the launcher on GPUs of its own. Needs ``[eval_snapshots]``. If None, + none is started (a worker can still be started by hand).""" + optim: OptimizerConfig = MuonConfig() """Optimizer. ``optim.lr`` is required.""" @@ -324,6 +351,12 @@ def validate_param_groups(self): raise ValueError("Pretraining's per-role parameter groups need AdamW or Muon and model.offload = 'none'.") return self + @model_validator(mode="after") + def validate_eval_worker(self): + if self.eval_worker is not None and self.eval_snapshots is None: + raise ValueError("[eval_worker] needs [eval_snapshots]") + return self + @model_validator(mode="after") def validate_cp(self): if self.model.cp > 1: diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py index 4ecd8611d4e..24692d31f34 100644 --- a/src/prime_rl/entrypoints/pretrain.py +++ b/src/prime_rl/entrypoints/pretrain.py @@ -3,7 +3,8 @@ import uuid from pathlib import Path -from prime_rl.configs.pretrain import PretrainConfig +from prime_rl.configs.pretrain import EvalWorkerConfig, PretrainConfig +from prime_rl.configs.shared import LogConfig from prime_rl.entrypoints.sft import submit_slurm_script, write_trainer_slurm_script from prime_rl.utils.config import cli, dump_resolved_config from prime_rl.utils.logger import get_logger, setup_logger @@ -20,12 +21,14 @@ DEFAULT_COMMON_ENV_VARS, DEFAULT_TRAINER_ENV_VARS, ProcessGroup, + partition_gpus, set_proc_title, torchrun_cmd, ) PRETRAIN_CONFIG = "pretrain.json" PRETRAIN_SBATCH = "pretrain.sbatch" +EVAL_WORKER_CONFIG = "eval_worker.json" def write_model_dir(config: PretrainConfig) -> None: @@ -56,10 +59,25 @@ def write_model_dir(config: PretrainConfig) -> None: tokenizer.save_pretrained(model_dir) +def write_eval_worker_config(config: PretrainConfig, config_dir: Path) -> Path: + """The eval worker's config: the ``[eval_worker]`` suite for this run's directory.""" + assert config.eval_worker is not None + suite = config.eval_worker.model_dump(exclude={"num_gpus", "num_nodes"}) + worker = EvalWorkerConfig(**suite, run_dir=config.run_dir, log=LogConfig(level=config.log.level)) + path = config_dir / EVAL_WORKER_CONFIG + path.write_text(json.dumps(dump_resolved_config(worker), indent=2)) + return path + + def pretrain_slurm(config: PretrainConfig, config_path: Path, log_dir: Path) -> None: logger = get_logger() script_path = get_launcher_dir(config.run_dir) / PRETRAIN_SBATCH - write_trainer_slurm_script(config, "pretrain", config_path, log_dir, script_path) + eval_worker_vars = {} + if config.eval_worker is not None and config.deployment.type == "multi_node": + eval_worker_vars = dict( + num_eval_nodes=config.eval_worker.num_nodes, eval_worker_config_path=config_path.parent / EVAL_WORKER_CONFIG + ) + write_trainer_slurm_script(config, "pretrain", config_path, log_dir, script_path, **eval_worker_vars) logger.info(f"Wrote SLURM script to {script_path}") num_nodes = config.deployment.num_train_nodes if config.deployment.type == "multi_node" else 1 log_message = format_log_message(log_dir=log_dir, trainer=True, num_train_nodes=num_nodes) @@ -103,7 +121,10 @@ def pretrain(config: PretrainConfig): # multi-node job read the config directly. exclude |= {"slurm", "deployment"} if config.deployment.type == "multi_node" else {"slurm"} config_path.write_text(json.dumps(dump_resolved_config(config, exclude=exclude), indent=2)) - logger.info(f"Configs:\n{format_config_message(config_dir, 'pretrain', [('Pretrain', config_path)])}") + components = [("Pretrain", config_path)] + if config.eval_worker is not None: + components.append(("Eval worker", write_eval_worker_config(config, config_dir))) + logger.info(f"Configs:\n{format_config_message(config_dir, 'pretrain', components)}") if config.slurm is not None: pretrain_slurm(config, config_path, log_dir) return @@ -112,7 +133,25 @@ def pretrain(config: PretrainConfig): return logger.info(format_log_message(log_dir=log_dir, trainer=True)) + trainer_env = {**os.environ, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, **config.env_vars} with ProcessGroup() as processes: + if config.eval_worker is not None: + # The worker's GPUs first, then the trainer's. + eval_gpu_ids, train_gpu_ids = partition_gpus(config.eval_worker.num_gpus, config.deployment.num_train_gpus) + trainer_env["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, train_gpu_ids)) + processes.start( + "eval-worker", + torchrun_cmd( + "prime_rl.trainer.pretrain.eval_worker", + config_dir / EVAL_WORKER_CONFIG, + config.eval_worker.num_gpus, + log_dir, + config.log.ranks_filter, + ), + env={**trainer_env, "CUDA_VISIBLE_DEVICES": ",".join(map(str, eval_gpu_ids))}, + log_path=log_dir / "eval_worker.log", + ) + logger.success(f"Started the eval worker on GPU(s) {' '.join(map(str, eval_gpu_ids))}") processes.start( "trainer", torchrun_cmd( @@ -122,11 +161,14 @@ def pretrain(config: PretrainConfig): log_dir, config.log.ranks_filter, ), - env={**os.environ, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, **config.env_vars}, + env=trainer_env, log_path=log_dir / "trainer.log", ) logger.success(f"Started trainer with {config.deployment.num_train_gpus} GPU(s)") processes.wait("trainer") + if config.eval_worker is not None: + # The worker scores the final snapshot after the trainer has written it. + processes.wait("eval-worker") logger.success("Pretraining finished!") diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 index 6dcb0ca9e44..c494970ec95 100644 --- a/src/prime_rl/templates/multi_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH --job-name={{ job_name }} -#SBATCH --nodes={{ num_train_nodes + num_infer_nodes }} +#SBATCH --nodes={{ num_train_nodes + num_infer_nodes + (num_eval_nodes or 0) }} #SBATCH --ntasks-per-node=1 #SBATCH --gres=gpu:{{ gpus_per_node }} {%- if partition %} @@ -59,8 +59,11 @@ export TRAINER_DONE=$LAUNCHER_DIR/.sft_trainer_done_$SLURM_JOB_ID {% if online_eval -%} export EVAL_DONE=$LAUNCHER_DIR/.sft_eval_done_$SLURM_JOB_ID {%- endif %} +{% if num_eval_nodes -%} +export EVAL_DONE=$LAUNCHER_DIR/.pretrain_eval_done_$SLURM_JOB_ID +{% endif -%} mkdir -p "$LAUNCHER_DIR" -rm -f "$TRAINER_DONE"{% if online_eval %} "$EVAL_DONE"{% endif %} +rm -f "$TRAINER_DONE"{% if online_eval or num_eval_nodes %} "$EVAL_DONE"{% endif %} export PRL_RUN_ID={{ prl_run_id }} export PRL_RUN_NAME={{ run_name }} @@ -155,6 +158,38 @@ srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' IB_HCA=$(ibv_devinfo 2>/dev/null | sed -n -e "/hca_id/p" -e "/link_layer:/p" | grep -B1 InfiniBand | grep hca_id | sed -e "s/^hca_id://g" | tr -d "[[:blank:]]" | paste -sd,) [ -n "$IB_HCA" ] && export NCCL_IB_HCA=$IB_HCA +{% if num_eval_nodes -%} +if [ "$SLURM_PROCID" -ge "$NUM_TRAIN_NODES" ]; then + # Eval worker nodes follow the trainer nodes. + EVAL_NODE_RANK=$((SLURM_PROCID - NUM_TRAIN_NODES)) + HOSTS=( $ALL_HOSTNAMES_STR ) +{%- for key, value in trainer_env_vars.items() %} + export {{ key }}="{{ value }}" +{%- endfor %} + mkdir -p $LOG_DIR/eval_worker + ( + uv run torchrun \ + --role=eval \ + --nnodes={{ num_eval_nodes }} \ + --nproc-per-node=$GPUS_PER_NODE \ + --node-rank=$EVAL_NODE_RANK \ + --rdzv-endpoint=${HOSTS[$NUM_TRAIN_NODES]}:29600 \ + --rdzv-id=eval_$SLURM_JOB_ID \ + --log-dir=$LOG_DIR/eval_worker/torchrun \ + --tee=3 \ + --redirects=3 \ + --local-ranks-filter=0 \ + -m prime_rl.trainer.pretrain.eval_worker \ + @ {{ eval_worker_config_path }} \ + 2>&1 | sed -u 's/^\[[a-zA-Z]*[0-9]*\]://' | tee -a $LOG_DIR/eval_worker/node_${EVAL_NODE_RANK}.log + COMPONENT_CODE=$? + if [ "$COMPONENT_CODE" -eq 0 ] && [ "$EVAL_NODE_RANK" -eq 0 ]; then + touch "$EVAL_DONE" || exit $? + fi + exit "$COMPONENT_CODE" + ) & +else +{% endif -%} {% if online_eval -%} if [ "$SLURM_PROCID" -lt "$NUM_INFER_NODES" ]; then INFER_NODE_RANK=$SLURM_PROCID @@ -254,9 +289,10 @@ else exit "$COMPONENT_CODE" ) & {% if online_eval %}fi{% endif %} +{%- if num_eval_nodes %}fi{% endif %} # Every node stays in the shared srun step until all terminal roles finish. -{% if online_eval -%} +{% if online_eval or num_eval_nodes -%} ( until [ -f "$TRAINER_DONE" ] && [ -f "$EVAL_DONE" ]; do sleep 10; done ) & {% else -%} ( until [ -f "$TRAINER_DONE" ]; do sleep 10; done ) & diff --git a/src/prime_rl/trainer/pretrain/eval_worker.py b/src/prime_rl/trainer/pretrain/eval_worker.py index 7446d65a8e3..7b088113d56 100644 --- a/src/prime_rl/trainer/pretrain/eval_worker.py +++ b/src/prime_rl/trainer/pretrain/eval_worker.py @@ -32,7 +32,13 @@ shared_wandb_env, snapshot_dir, ) -from prime_rl.trainer.pretrain.suite import prepare_set, score_requests, summarize, use_scoring_head +from prime_rl.trainer.pretrain.suite import ( + length_label, + prepare_multi_length_set, + score_requests, + summarize, + use_scoring_head, +) from prime_rl.trainer.utils import get_dcp_group, setup_torch_distributed from prime_rl.trainer.world import get_world from prime_rl.utils.config import cli @@ -92,10 +98,12 @@ def main_loop(config: EvalWorkerConfig) -> None: use_scoring_head(model) tokenizer = AutoTokenizer.from_pretrained(run.tokenizer.name, trust_remote_code=run.tokenizer.trust_remote_code) start = time.perf_counter() - eval_sets = [prepare_set(name, config.data_dir, tokenizer, config.seq_len) for name in config.sets] + eval_sets = [prepare_multi_length_set(name, config.data_dir, tokenizer, config.seq_lens) for name in config.sets] num_tokens = sum(len(ids) - begin for eval_set in eval_sets for ids, begin in eval_set.requests) - logger.info(f"Prepared {len(eval_sets)} sets ({num_tokens:,} scored tokens) in {time.perf_counter() - start:.0f}s") - + logger.info( + f"Prepared {len(eval_sets)} sets at context lengths {config.seq_lens} ({num_tokens:,} distinct scored tokens) " + f"in {time.perf_counter() - start:.0f}s" + ) last_step = -1 while True: newest = broadcast(max(finished_snapshot_steps(config.run_dir), default=-1) if world.is_master else 0) @@ -117,17 +125,23 @@ def main_loop(config: EvalWorkerConfig) -> None: scored = score_requests(model, eval_set.requests, config.token_budget, get_dcp_group()) if not world.is_master: continue - summary, arrays = summarize(eval_set, *scored) - metrics |= { - f"eval/{eval_set.name}/{key}": value for key, value in summary.items() if isinstance(value, float) - } - if config.save_arrays: - out = config.run_dir / "evals" / f"step_{newest}" - out.mkdir(parents=True, exist_ok=True) - np.savez(out / f"{eval_set.name}.npz", **arrays) - (out / f"{eval_set.name}.json").write_text( - json.dumps(summary | {"data_sha256": eval_set.sha256}, default=float) - ) + logprobs, hits = scored + for seq_len, version in eval_set.versions.items(): + index = eval_set.index[seq_len] + summary, arrays = summarize(version, [logprobs[i] for i in index], [hits[i] for i in index]) + label = length_label(seq_len) + metrics |= { + f"eval/{eval_set.name}/{key}@{label}": value + for key, value in summary.items() + if isinstance(value, float) + } + if config.save_arrays: + out = config.run_dir / "evals" / f"step_{newest}" + out.mkdir(parents=True, exist_ok=True) + np.savez(out / f"{eval_set.name}@{label}.npz", **arrays) + (out / f"{eval_set.name}@{label}.json").write_text( + json.dumps(summary | {"data_sha256": version.sha256}, default=float) + ) if world.is_master: round_time = time.perf_counter() - round_start metrics |= {"time/eval_round": round_time, "time/eval_load": load_time, "step": newest} diff --git a/src/prime_rl/trainer/pretrain/suite.py b/src/prime_rl/trainer/pretrain/suite.py index 4f03696ad01..5ef6c71eb0a 100644 --- a/src/prime_rl/trainer/pretrain/suite.py +++ b/src/prime_rl/trainer/pretrain/suite.py @@ -190,6 +190,36 @@ def prepare_set(name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_l return EvalSet(name, kind, table, requests, context, meta["parquet_sha256"]) +@dataclass +class MultiLengthSet: + """One eval set at several context lengths (``ppl`` window / ``ctxcont`` context limit). A request that is the + same at two lengths (e.g. a document shorter than both windows) is scored once.""" + + name: str + versions: dict[int, EvalSet] + requests: list[Request] + """The distinct requests over all lengths.""" + index: dict[int, list[int]] + """Per length: the position in ``requests`` of each of that version's requests.""" + + +def prepare_multi_length_set( + name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_lens: list[int] +) -> MultiLengthSet: + versions = {seq_len: prepare_set(name, data_dir, tokenizer, seq_len) for seq_len in seq_lens} + distinct: dict[tuple[tuple[int, ...], int], int] = {} + index = { + seq_len: [distinct.setdefault((tuple(ids), start), len(distinct)) for ids, start in version.requests] + for seq_len, version in versions.items() + } + return MultiLengthSet(name, versions, [(list(ids), start) for ids, start in distinct], index) + + +def length_label(seq_len: int) -> str: + """``2048`` -> ``2k``, the suffix of the metric names.""" + return f"{seq_len // 1024}k" if seq_len % 1024 == 0 else str(seq_len) + + def summarize(eval_set: EvalSet, logprobs: list[np.ndarray], hits: list[np.ndarray]) -> tuple[dict, dict]: """Summary metrics and per-document / per-item arrays of one scored set.""" if eval_set.kind == "ppl": diff --git a/tests/unit/train/pretrain/test_launcher.py b/tests/unit/train/pretrain/test_launcher.py index 9238d16d42e..db4e4750c7c 100644 --- a/tests/unit/train/pretrain/test_launcher.py +++ b/tests/unit/train/pretrain/test_launcher.py @@ -1,6 +1,6 @@ import re -from prime_rl.configs.pretrain import PretrainConfig +from prime_rl.configs.pretrain import EvalWorkerConfig, PretrainConfig from prime_rl.entrypoints.pretrain import pretrain from prime_rl.utils.config import cli @@ -31,3 +31,31 @@ def test_multi_node_dry_run_writes_trainer_sbatch(tmp_path, monkeypatch): trainer_config = cli(PretrainConfig, args=["@", config_path]) assert trainer_config.slurm is None assert trainer_config.model == config.model + + +def test_multi_node_dry_run_adds_the_eval_worker_node(tmp_path, monkeypatch): + monkeypatch.setenv("PRL_RUN_ID", "test") + monkeypatch.setenv("PRL_RUN_NAME", "test") + config = PretrainConfig.model_validate( + { + "output_dir": tmp_path, + "model": {"arch": {"model_type": "llama"}}, + "tokenizer": {"name": "PrimeIntellect/Qwen3-0.6B"}, + "data": {"type": "fake", "seq_len": 8, "batch_tokens": 16, "micro_batch_tokens": 8}, + "optim": {"lr": 1e-3}, + "max_steps": 1, + "eval_snapshots": {"interval": 1}, + "eval_worker": {"data_dir": tmp_path, "sets": ["holdout"], "seq_lens": [2048, 4096]}, + "deployment": {"type": "multi_node", "num_train_nodes": 2}, + "slurm": {}, + "dry_run": True, + } + ) + pretrain(config) + + script = (config.run_dir / "launcher" / "pretrain.sbatch").read_text() + assert "#SBATCH --nodes=3" in script + worker_config_path = re.search(r"-m prime_rl.trainer.pretrain.eval_worker \\\n\s+@ (\S+)", script).group(1) + worker_config = cli(EvalWorkerConfig, args=["@", worker_config_path]) + assert worker_config.run_dir == config.run_dir + assert worker_config.seq_lens == [2048, 4096] diff --git a/tests/unit/train/pretrain/test_suite.py b/tests/unit/train/pretrain/test_suite.py index 3fd63f10453..0b7afab772c 100644 --- a/tests/unit/train/pretrain/test_suite.py +++ b/tests/unit/train/pretrain/test_suite.py @@ -1,8 +1,11 @@ +import json import math import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq -from prime_rl.trainer.pretrain.suite import pack_requests, summarize_ctxcont +from prime_rl.trainer.pretrain.suite import pack_requests, prepare_multi_length_set, summarize_ctxcont def test_pack_requests_covers_every_request_within_the_budget(): @@ -24,3 +27,21 @@ def test_summarize_ctxcont_scores_the_gold_choice(): assert summary["acc"] == 0.5 # item 0: gold wins; item 1: the other choice wins assert math.isclose(summary["bpb"], (2.0 + 2.0) / (3 * math.log(2))) # gold nll / gold continuation bytes assert summary["greedy_em"] == 0.0 + + +class CharTokenizer: + bos_token_id, eos_token_id = 1, 2 + + def __call__(self, texts, add_special_tokens=False): + return {"input_ids": [[ord(char) for char in text] for text in texts]} + + +def test_multi_length_set_scores_shared_requests_once(tmp_path): + (tmp_path / "docs").mkdir() + pq.write_table(pa.table({"text": ["short", "a much longer document"]}), tmp_path / "docs" / "test.parquet") + (tmp_path / "docs" / "meta.json").write_text(json.dumps({"kind": "ppl", "parquet_sha256": "x"})) + multi = prepare_multi_length_set("docs", tmp_path, CharTokenizer(), [8, 32]) + # "short" (5 chars + BOS + EOS) is one window at both lengths; the long document is 3 windows at 8, 1 at 32. + assert len(multi.versions[8].requests) == 4 and len(multi.versions[32].requests) == 2 + assert len(multi.requests) == 5 + assert multi.index[8][0] == multi.index[32][0] From 2cffe13a0035e828139776a16db6711b726298e4 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 20:36:28 +0200 Subject: [PATCH 046/132] test(pretrain): unique basename for the pretrain metrics test --- .../train/pretrain/{test_metrics.py => test_pretrain_metrics.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/unit/train/pretrain/{test_metrics.py => test_pretrain_metrics.py} (100%) diff --git a/tests/unit/train/pretrain/test_metrics.py b/tests/unit/train/pretrain/test_pretrain_metrics.py similarity index 100% rename from tests/unit/train/pretrain/test_metrics.py rename to tests/unit/train/pretrain/test_pretrain_metrics.py From 06fc55df0288d92cd92357383c488dbc7ec0211c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 20:36:28 +0200 Subject: [PATCH 047/132] fix(pretrain): the eval worker only attaches to the trainer's platform run, and fails clearly without one The worker logs through the eval platform monitor, which attaches to the parent run for metrics and never reports its status (as SFT's online evals do). Waiting for the trainer's run id stops with an error once the trainer logs metrics without one. --- src/prime_rl/trainer/pretrain/eval_worker.py | 26 +++++++++++++++----- 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/eval_worker.py b/src/prime_rl/trainer/pretrain/eval_worker.py index 7b088113d56..f6573302d9e 100644 --- a/src/prime_rl/trainer/pretrain/eval_worker.py +++ b/src/prime_rl/trainer/pretrain/eval_worker.py @@ -13,6 +13,7 @@ import os import time from datetime import timedelta +from pathlib import Path import numpy as np import torch @@ -21,6 +22,7 @@ from transformers import AutoTokenizer from prime_rl import monitors +from prime_rl.configs.monitors import PrimeEvalMonitorConfig from prime_rl.configs.pretrain import EvalWorkerConfig, PretrainConfig from prime_rl.monitors.prime import read_platform_record from prime_rl.trainer.ckpt import AppState @@ -43,6 +45,7 @@ from prime_rl.trainer.world import get_world from prime_rl.utils.config import cli from prime_rl.utils.logger import setup_logger +from prime_rl.utils.pathing import get_file_monitor_dir from prime_rl.utils.process import set_proc_title from prime_rl.utils.utils import clean_exit @@ -61,6 +64,17 @@ def run_config(config: EvalWorkerConfig) -> PretrainConfig: return run +def platform_run_id(run_dir: Path, run: PretrainConfig, poll_seconds: int) -> str: + """The id of the trainer's platform run, which its monitor writes to the run dir when it starts, before the + trainer logs any metrics. Fails if the trainer logs metrics without it (e.g. platform runs disabled).""" + metrics_path = get_file_monitor_dir(run_dir) / run.monitors.file.path if run.monitors.file is not None else None + while (record := read_platform_record(run_dir)) is None or "id" not in record: + if metrics_path is not None and metrics_path.is_file() and metrics_path.stat().st_size > 0: + raise RuntimeError(f"The trainer logs metrics but left no platform run id in {run_dir}") + time.sleep(poll_seconds) + return record["id"] + + def broadcast(value: int) -> int: values = [value] dist.broadcast_object_list(values, src=0, group=get_dcp_group()) @@ -76,19 +90,19 @@ def main_loop(config: EvalWorkerConfig) -> None: torch.set_float32_matmul_precision(run.matmul_precision) os.environ.update(shared_wandb_env(config.run_dir, "eval")) + prime = None if run.monitors.prime is not None: - # Attach to the trainer's platform run, whose id the trainer's monitor leaves in the run dir. - while (record := read_platform_record(config.run_dir)) is None or "id" not in record: - time.sleep(config.poll_seconds) - os.environ["RUN_ID"] = record["id"] + # Metrics onto the trainer's platform run, which the trainer owns: the eval monitor only attaches to it + # and never reports its status. + os.environ["RUN_ID"] = platform_run_id(config.run_dir, run, config.poll_seconds) + prime = PrimeEvalMonitorConfig(name=run.monitors.prime.name) asyncio.run( monitors.setup( producer="eval", wandb=run.monitors.wandb, - prime=run.monitors.prime, + prime=prime, file=run.monitors.file, output_dir=config.run_dir, - run_config=run, overview_flavor="pretrain", ) ) From 3831683ae1bdbd7a775c3a0c51ec3c14ba91baa9 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 20:37:03 +0200 Subject: [PATCH 048/132] revert(configs): SFTConfig keeps its own deployment validators The multi-node merge had folded SFTConfig's validate_deployment, auto_setup_deployment and auto_setup_slurm_template into a shared setup_trainer_deployment that ran at a different point in validation (so an invalid multi-node SFT config could raise a different error first). SFTConfig's validators are back exactly as on main; PretrainConfig has its own auto_setup_deployment. The rendered SFT sbatch scripts (single node, multi-node, multi-node with online evals) are byte-identical to main's. --- .../src/prime_rl/configs/pretrain.py | 24 ++++++++++++-- .../src/prime_rl/configs/sft.py | 31 ++++++++++++++++-- .../src/prime_rl/configs/shared.py | 32 ++----------------- 3 files changed, 51 insertions(+), 36 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 3e2e686f2f6..8caa8241e82 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -16,7 +16,6 @@ SingleNodeDeploymentConfig, SlurmConfig, TrainerLogConfig, - setup_trainer_deployment, ) from prime_rl.configs.trainer import ( CheckpointConfig, @@ -26,7 +25,7 @@ OptimizerConfig, TokenizerConfig, ) -from prime_rl.utils.config import BaseConfig, default_output_dir +from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource TOKENIZER_FIELDS = ("vocab_size", "bos_token_id", "eos_token_id", "pad_token_id") @@ -305,7 +304,26 @@ def auto_setup_tokenizer(self): @model_validator(mode="after") def auto_setup_deployment(self): - setup_trainer_deployment(self) + """A multi-node deployment needs SLURM and sets ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the + SLURM job name defaults to the run name and the template to the bundled single- or multi-node template.""" + deployment = self.deployment + if deployment.type == "multi_node": + if self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + if deployment.nodes_per_fsdp_group is not None: + if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" + ) + self.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group + if self.slurm is not None: + if self.slurm.job_name is None: + self.slurm.job_name = self.run.name + if self.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + self.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" return self @model_validator(mode="after") diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 88866b04003..211230b64fe 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -24,7 +24,6 @@ SlurmConfig, TrainerLogConfig, WeightBroadcastConfig, - setup_trainer_deployment, ) from prime_rl.configs.trainer import ( AdamWConfig, @@ -37,7 +36,7 @@ TokenizerConfig, validate_scheduler, ) -from prime_rl.utils.config import BaseConfig, default_output_dir +from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource class BaseDataConfig(BaseConfig): @@ -366,6 +365,12 @@ def full_optimizer_offload_disables_grad_clipping(self): self.optim.max_norm = None return self + @model_validator(mode="after") + def validate_deployment(self): + if self.deployment.type == "multi_node" and self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + return self + @model_validator(mode="after") def auto_setup_online_eval(self): """Wire online evals and the trainer-to-inference weight transport.""" @@ -588,5 +593,25 @@ def auto_setup_tokenizer(self): @model_validator(mode="after") def auto_setup_deployment(self): - setup_trainer_deployment(self) + if self.deployment.type == "multi_node": + if self.deployment.nodes_per_fsdp_group is not None: + if self.deployment.num_train_nodes % self.deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_train_nodes ({self.deployment.num_train_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({self.deployment.nodes_per_fsdp_group})" + ) + self.model.dp_replicate = self.deployment.num_train_nodes // self.deployment.nodes_per_fsdp_group + return self + + @model_validator(mode="after") + def auto_setup_slurm_template(self): + if self.slurm is not None and self.slurm.job_name is None: + self.slurm.job_name = self.run.name + if self.slurm is not None and self.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + if self.deployment.type == "single_node": + self.slurm.template_path = templates_dir / "single_node_sft.sbatch.j2" + else: + self.slurm.template_path = templates_dir / "multi_node_sft.sbatch.j2" return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 52b2883207d..49938677daf 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -1,16 +1,12 @@ import os import re from pathlib import Path -from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias +from typing import Annotated, Literal, TypeAlias from pydantic import AfterValidator, AliasChoices, Field, model_validator from verifiers.v1.configs.client import BaseClientConfig -from prime_rl.utils.config import BaseConfig, find_package_resource - -if TYPE_CHECKING: - from prime_rl.configs.pretrain import PretrainConfig - from prime_rl.configs.sft import SFTConfig +from prime_rl.utils.config import BaseConfig # Launcher-managed env vars that a component's `env_vars` must not set: GPU partitioning # and the single shared W&B run. The launcher always sets these last, so allowing them in @@ -337,27 +333,3 @@ class MultiNodeDeploymentConfig(BaseDeploymentConfig): SFTDeploymentConfig: TypeAlias = Annotated[ SingleNodeDeploymentConfig | MultiNodeDeploymentConfig, Field(discriminator="type") ] - - -def setup_trainer_deployment(config: "SFTConfig | PretrainConfig") -> None: - """Shared by ``sft`` and ``pretrain``: a multi-node deployment needs SLURM and sets - ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the SLURM job name defaults to the - run name and the template to the bundled single- or multi-node trainer template.""" - deployment = config.deployment - if deployment.type == "multi_node": - if config.slurm is None: - raise ValueError("Must use SLURM for multi-node deployment.") - if deployment.nodes_per_fsdp_group is not None: - if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: - raise ValueError( - f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " - f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" - ) - config.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group - if config.slurm is not None: - if config.slurm.job_name is None: - config.slurm.job_name = config.run.name - if config.slurm.template_path is None: - templates_dir = find_package_resource("templates") - if templates_dir is not None: - config.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" From 80f9bbfffc4f9c428cf538cba4451cea81c3a752 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 22:57:56 +0200 Subject: [PATCH 049/132] fix(pretrain): the eval worker only reads platform state of the current launch --- src/prime_rl/trainer/pretrain/eval_worker.py | 30 +++++++++++++++----- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/eval_worker.py b/src/prime_rl/trainer/pretrain/eval_worker.py index f6573302d9e..9e97841a43f 100644 --- a/src/prime_rl/trainer/pretrain/eval_worker.py +++ b/src/prime_rl/trainer/pretrain/eval_worker.py @@ -45,15 +45,17 @@ from prime_rl.trainer.world import get_world from prime_rl.utils.config import cli from prime_rl.utils.logger import setup_logger -from prime_rl.utils.pathing import get_file_monitor_dir +from prime_rl.utils.pathing import get_file_monitor_dir, get_platform_run_path from prime_rl.utils.process import set_proc_title from prime_rl.utils.utils import clean_exit +RUN_CONFIG = "configs/latest/resolved/pretrain.json" + def run_config(config: EvalWorkerConfig) -> PretrainConfig: """The training run's resolved config (waiting for the launcher to write it), with the parallelism of the worker: plain FSDP over its GPUs.""" - path = config.run_dir / "configs/latest/resolved/pretrain.json" + path = config.run_dir / RUN_CONFIG while not path.exists(): time.sleep(config.poll_seconds) run = PretrainConfig.model_validate_json(path.read_text()) @@ -66,13 +68,27 @@ def run_config(config: EvalWorkerConfig) -> PretrainConfig: def platform_run_id(run_dir: Path, run: PretrainConfig, poll_seconds: int) -> str: """The id of the trainer's platform run, which its monitor writes to the run dir when it starts, before the - trainer logs any metrics. Fails if the trainer logs metrics without it (e.g. platform runs disabled).""" + trainer logs any metrics. Only state of the current launch counts (a resumed run dir holds an earlier + attempt's record and metrics). Fails if the trainer logs metrics of this launch without it (e.g. platform + runs disabled).""" + launched = (run_dir / RUN_CONFIG).stat().st_mtime + record_path = get_platform_run_path(run_dir) metrics_path = get_file_monitor_dir(run_dir) / run.monitors.file.path if run.monitors.file is not None else None - while (record := read_platform_record(run_dir)) is None or "id" not in record: - if metrics_path is not None and metrics_path.is_file() and metrics_path.stat().st_size > 0: - raise RuntimeError(f"The trainer logs metrics but left no platform run id in {run_dir}") + while True: + record = read_platform_record(run_dir) + if record is not None and "id" in record and record_path.stat().st_mtime >= launched: + return record["id"] + if metrics_path is not None and metrics_path.is_file() and last_row_time(metrics_path) >= launched: + raise RuntimeError(f"The trainer logs metrics but left no platform run id for this launch in {run_dir}") time.sleep(poll_seconds) - return record["id"] + + +def last_row_time(path: Path) -> float: + """The ``time`` of the last complete row of a ``metrics.jsonl``.""" + with path.open("rb") as f: + f.seek(max(0, f.seek(0, 2) - 65536)) + rows = f.read().splitlines() + return json.loads(rows[-1])["time"] if rows else 0.0 def broadcast(value: int) -> int: From 1608bbf470dfe85c7cc5edb6eb0fb92f549e0aeb Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 23:08:42 +0200 Subject: [PATCH 050/132] feat(pretrain): diagnostic step with attention/activation probe and update sizes Every diagnostics.interval steps (default 100): the optimizer update's size relative to the weights and the weight RMS, per parameter group and per norm role; and a forward pass on a fixed probe batch (the first probe_tokens tokens of each rank's first micro-batch, which is then trained on as usual) that records per layer the attention max logit and entropy per head (collapsed heads below 0.5 nats) and the residual-stream RMS and max. The attention statistics come from wrapping FlashAttention._compute_attention during the probe, as ring attention does. --- .../src/prime_rl/configs/pretrain.py | 17 +++ src/prime_rl/trainer/pretrain/diagnostics.py | 144 ++++++++++++++++++ src/prime_rl/trainer/pretrain/metrics.py | 25 +-- src/prime_rl/trainer/pretrain/train.py | 22 +++ tests/unit/train/pretrain/test_diagnostics.py | 22 +++ 5 files changed, 219 insertions(+), 11 deletions(-) create mode 100644 src/prime_rl/trainer/pretrain/diagnostics.py create mode 100644 tests/unit/train/pretrain/test_diagnostics.py diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index a752b1deb01..0211295bb8b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -231,6 +231,18 @@ class ScheduleConfig(BaseConfig): """Shape of the decay from 1 to ``end_ratio``.""" +class DiagnosticsConfig(BaseConfig): + """Diagnostic step: on every ``interval``-th step, the optimizer update's size per parameter group and a + forward pass on a fixed probe batch (the first ``probe_tokens`` tokens of each rank's first micro-batch) + that records attention max logits and entropy per head and the residual-stream size per layer.""" + + interval: int = Field(100, ge=1) + """Run the diagnostics every N steps.""" + + probe_tokens: int = Field(2048, ge=64) + """Tokens per rank in the probe batch. Fixed across sequence lengths so attention entropy stays comparable.""" + + class EvalSnapshotsConfig(BaseConfig): """Model-only snapshots for the asynchronous eval worker (`prime_rl.trainer.pretrain.eval_worker`), written in the background to ``/snapshots/step_N``. The worker deletes each one after scoring it.""" @@ -317,6 +329,9 @@ class PretrainConfig(BaseConfig): val: ValConfig | None = None """Held-out evals. If None, none run.""" + diagnostics: DiagnosticsConfig | None = DiagnosticsConfig() + """Diagnostic step (update sizes, attention and activation probe). If None, it never runs.""" + eval_snapshots: EvalSnapshotsConfig | None = None """Model-only snapshots for the asynchronous eval worker. If None, none are written.""" @@ -478,6 +493,8 @@ def validate_cp(self): raise ValueError("data.seq_len must be divisible by model.cp") if self.data.micro_batch_size != 1: raise ValueError("data.micro_batch_tokens must equal data.seq_len with context parallelism") + if self.diagnostics is not None: + raise ValueError("[diagnostics] does not support context parallelism yet; set diagnostics = null") return self @model_validator(mode="after") diff --git a/src/prime_rl/trainer/pretrain/diagnostics.py b/src/prime_rl/trainer/pretrain/diagnostics.py new file mode 100644 index 00000000000..bb9be1bbd22 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/diagnostics.py @@ -0,0 +1,144 @@ +import torch +from torch import Tensor, nn + +from prime_rl.trainer.model import forward +from prime_rl.trainer.models.layers.attn import FlashAttention +from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX +from prime_rl.trainer.pretrain.metrics import group_norms, param_group, reduce_stats +from prime_rl.utils.vlm import get_language_model + +MIN_VISIBLE_KEYS = 64 +"""Attention entropy is averaged over queries that see at least this many keys.""" + +COLLAPSED_ENTROPY = 0.5 +"""A head whose mean attention entropy (nats) is below this counts as collapsed.""" + + +def probe_batch(micro_batch: dict[str, Tensor], num_tokens: int) -> dict[str, Tensor]: + """The first ``num_tokens`` tokens of a training micro-batch, on the GPU: the fixed batch of the diagnostic + forward.""" + lengths, total = [], 0 + for length in micro_batch["seq_lens"].tolist(): + lengths.append(min(length, num_tokens - total)) + total += lengths[-1] + if total == num_tokens: + break + batch = {key: micro_batch[key][:, :total].to("cuda") for key in ("input_ids", "position_ids", "target_ids")} + batch["seq_lens"] = torch.tensor(lengths, device="cuda") + return batch + + +def attention_stats(q: Tensor, k: Tensor, cu_seqlens: Tensor, scale: float) -> tuple[Tensor, Tensor, Tensor]: + """Per head: max attention logit (q·k times the softmax scale, after RoPE and QK-norm, over causal pairs + within a document), summed attention entropy (nats) over the queries that see at least `MIN_VISIBLE_KEYS` + keys, and the number of those queries. ``q`` is `[tokens, heads, head_dim]`, ``k`` `[tokens, kv_heads, head_dim]`.""" + num_tokens, num_heads, _ = q.shape + positions = torch.arange(num_tokens, device=q.device) + document = torch.searchsorted(cu_seqlens[1:].long(), positions, right=True) + mask = (positions[None, :] <= positions[:, None]) & (document[None, :] == document[:, None]) + valid = mask.sum(dim=1) >= MIN_VISIBLE_KEYS + groups = num_heads // k.shape[1] + max_logits, entropy_sums = [], [] + for head in range(num_heads): + scores = (q[:, head].float() @ k[:, head // groups].float().T * scale).masked_fill(~mask, float("-inf")) + max_logits.append(scores.amax()) + entropy = torch.special.entr(scores.softmax(dim=-1)).sum(dim=-1) + entropy_sums.append(entropy[valid].sum()) + return torch.stack(max_logits), torch.stack(entropy_sums), valid.sum() + + +@torch.no_grad() +def run_probe(model: nn.Module, batch: dict[str, Tensor], group) -> dict[str, float]: + """Forward the probe batch with hooks: attention max logit and entropy per layer and head, residual-stream RMS + and max per layer, and the probe loss. Every rank probes its own batch; the statistics are reduced over + ``group``.""" + layers = get_language_model(model).layers + attentions = [module for module in layers.modules() if isinstance(module, FlashAttention)] + attention, activations = [], [] + + def capture(module: FlashAttention): + compute = module._compute_attention + + def compute_and_record(q, k, v, cu_seqlens, max_seqlen): + attention.append(attention_stats(q, k, cu_seqlens, module.scaling)) + return compute(q, k, v, cu_seqlens, max_seqlen) + + return compute_and_record + + def record_activations(module, args, output): + hidden = (output[0] if isinstance(output, tuple) else output).float() + activations.append( + (hidden.square().sum(), torch.tensor(hidden.numel(), device=hidden.device), hidden.abs().amax()) + ) + + for module in attentions: + module._compute_attention = capture(module) + hooks = [layer.register_forward_hook(record_activations) for layer in layers] + labels = batch["target_ids"] + out = forward(model, batch["input_ids"], batch["position_ids"], seq_lens=batch["seq_lens"], labels=labels) + for module in attentions: + del module._compute_attention + for hook in hooks: + hook.remove() + + max_logit, entropy_sum, queries = (torch.stack(values) for values in zip(*attention)) + square_sum, count, absmax = (torch.stack(values) for values in zip(*activations)) + sums = { + "entropy_sum": entropy_sum, + "queries": queries.double(), + "square_sum": square_sum, + "count": count.double(), + "loss": out["loss"].double(), + "tokens": (labels != IGNORE_INDEX).sum().double(), + } + maxes = {"max_logit": max_logit, "absmax": absmax} + stats = reduce_stats(sums, maxes, group) + + head_entropy = stats["entropy_sum"] / stats["queries"][:, None] # [layers, heads] + layer_max_logit = stats["max_logit"].amax(dim=1) + layer_entropy = head_entropy.mean(dim=1) + rms = (stats["square_sum"] / stats["count"]).sqrt() + metrics = { + "attn/max_logit/max": layer_max_logit.max().item(), + "attn/entropy/min": layer_entropy.min().item(), + "attn/entropy_min_head": head_entropy.min().item(), + "attn/collapsed_heads": int((head_entropy < COLLAPSED_ENTROPY).sum()), + "act/rms/max": rms.max().item(), + "act/absmax/max": stats["absmax"].max().item(), + "probe/loss": (stats["loss"] / stats["tokens"]).item(), + } + for layer in range(len(layer_max_logit)): + metrics[f"attn/max_logit/layer_{layer}"] = layer_max_logit[layer].item() + metrics[f"attn/entropy/layer_{layer}"] = layer_entropy[layer].item() + metrics[f"act/rms/layer_{layer}"] = rms[layer].item() + metrics[f"act/absmax/layer_{layer}"] = stats["absmax"][layer].item() + return metrics + + +def norm_role(name: str) -> str: + """Normalization layers by role (`input_layernorm`, `post_attention_layernorm`, `q_norm`, `k_norm`, `norm`), + everything else by `param_group`.""" + module = name.split(".")[-2] + return f"norm_gain/{module}" if "norm" in module else param_group(name) + + +class UpdateTracker: + """Size of one optimizer update relative to the weights, per parameter group: ||W_after - W_before|| / + ||W_before||, plus the weight RMS per group and of the norm gains by role.""" + + def before_step(self, model: nn.Module) -> None: + self.before = {name: param.detach().clone() for name, param in model.named_parameters()} + + def after_step(self, model: nn.Module) -> dict[str, float]: + params = dict(model.named_parameters()) + weight = group_norms(((name, self.before[name]) for name in params), norm_role) + update = group_norms(((name, params[name].detach() - self.before[name]) for name in params), norm_role) + numel = {} + for name, param in params.items(): + numel[norm_role(name)] = numel.get(norm_role(name), 0) + param.numel() + del self.before + metrics = {} + for group, norm in weight.items(): + metrics[f"param/rms/{group}"] = (norm / numel[group] ** 0.5).item() + metrics[f"update/rms_ratio/{group}"] = (update[group] / norm).item() + return metrics diff --git a/src/prime_rl/trainer/pretrain/metrics.py b/src/prime_rl/trainer/pretrain/metrics.py index 227c9292bc5..90e42fe7911 100644 --- a/src/prime_rl/trainer/pretrain/metrics.py +++ b/src/prime_rl/trainer/pretrain/metrics.py @@ -1,5 +1,6 @@ import math from collections import defaultdict, deque +from collections.abc import Iterable import torch import torch.distributed as dist @@ -106,21 +107,18 @@ def param_group(name: str) -> str: return "mlp" -def grad_norms_by_group(model: nn.Module) -> dict[str, Tensor]: - """Gradient L2 norm per parameter group (see `param_group`). The squared norms of the local shards are summed - over the sharded mesh dimensions only (replicated dimensions hold the same gradient), one collective per mesh - and placement.""" +def group_norms(named_tensors: Iterable[tuple[str, Tensor]], group_of=param_group) -> dict[str, Tensor]: + """L2 norm of (possibly sharded) tensors per group. The squared norms of the local shards are summed over the + sharded mesh dimensions only (replicated dimensions hold the same values), one collective per mesh and + placement.""" by_layout: dict = defaultdict(lambda: defaultdict(list)) - for name, param in model.named_parameters(): - grad = param.grad - if grad is None: - continue - layout = (grad.device_mesh, tuple(grad.placements)) if isinstance(grad, DTensor) else None - by_layout[layout][param_group(name)].append(grad.to_local() if isinstance(grad, DTensor) else grad) + for name, tensor in named_tensors: + layout = (tensor.device_mesh, tuple(tensor.placements)) if isinstance(tensor, DTensor) else None + by_layout[layout][group_of(name)].append(tensor.to_local() if isinstance(tensor, DTensor) else tensor) squares: dict[str, Tensor] = {} for layout, groups in by_layout.items(): local = torch.stack( - [torch.stack(torch._foreach_norm(grads)).double().square().sum() for grads in groups.values()] + [torch.stack(torch._foreach_norm(tensors)).double().square().sum() for tensors in groups.values()] ) if layout is not None: mesh, placements = layout @@ -130,3 +128,8 @@ def grad_norms_by_group(model: nn.Module) -> dict[str, Tensor]: for group, square in zip(groups, local): squares[group] = squares.get(group, 0) + square return {group: square.sqrt() for group, square in squares.items()} + + +def grad_norms_by_group(model: nn.Module) -> dict[str, Tensor]: + """Gradient L2 norm per parameter group (see `param_group`).""" + return group_norms((name, param.grad) for name, param in model.named_parameters() if param.grad is not None) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 7bbe33ff549..420776614a7 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -1,6 +1,7 @@ import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before import import asyncio +import itertools import os import math import time @@ -44,6 +45,7 @@ ) from prime_rl.trainer.pretrain.schedules import LRSchedule from prime_rl.trainer.data import get_dataset_progress, get_dataset_state +from prime_rl.trainer.pretrain.diagnostics import UpdateTracker, probe_batch, run_probe from prime_rl.trainer.pretrain.snapshots import SnapshotWriter, shared_wandb_env from prime_rl.trainer.pretrain.step import run_micro_batches from prime_rl.trainer.pretrain.metrics import DeadExperts, SpikeDetector, grad_norms_by_group @@ -210,6 +212,12 @@ def train(config: PretrainConfig): # Create the iterator only after a potential resume: iter() forks the workers with the # dataset's current state. dataiter = iter(dataloader) + probe = None + if config.diagnostics is not None: + # The diagnostic forward's fixed batch: the start of the first micro-batch, which is then trained on as usual. + first_micro_batch = next(dataiter) + probe = probe_batch(first_micro_batch, config.diagnostics.probe_tokens) + dataiter = itertools.chain([first_micro_batch], dataiter) cp_enabled = parallel_dims.cp_enabled cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() if cp_enabled else 0 @@ -293,6 +301,7 @@ def run_evals(step: int) -> None: loss_spikes, grad_norm_spikes = SpikeDetector(), SpikeDetector() last_bias_moves = None snapshots = SnapshotWriter(config.run_dir) if config.eval_snapshots is not None else None + update_tracker = UpdateTracker() logger.info(f"Starting training loop (max_steps={config.max_steps})") while True: torch.cuda.reset_peak_memory_stats() @@ -337,6 +346,9 @@ def run_evals(step: int) -> None: group_grad_norms = grad_norms_by_group(model) grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm) if config.optim.max_norm else None + is_diagnostic_step = config.diagnostics is not None and progress.step % config.diagnostics.interval == 0 + if is_diagnostic_step: + update_tracker.before_step(model) optimizer.step() optimizer.zero_grad() current_lr = optimizer.param_groups[0]["lr"] @@ -354,6 +366,15 @@ def run_evals(step: int) -> None: save_ckpt_time = time.perf_counter() - save_ckpt_start_time disk_metrics = get_ckpt_disk_metrics(config.run_dir) + diagnostic_metrics = {} + if is_diagnostic_step: + diagnostic_metrics = update_tracker.after_step(model) | run_probe(model, probe, dp_cp_group) + # The no-grad forward leaves the [lm_head, norm] FSDP group unsharded, and its routing must not count + # toward the next step's MoE statistics. + reshard_module(model) + if is_moe_model: + get_load_balance_stats(model) + if snapshots is not None and (is_last_step or progress.step % config.eval_snapshots.interval == 0): snapshots.save(model, progress.step) @@ -403,6 +424,7 @@ def run_evals(step: int) -> None: **disk_metrics, **result.metrics, **moe_metrics, + **diagnostic_metrics, "schedule/lr": lr_multiplier, "step": progress.step, } diff --git a/tests/unit/train/pretrain/test_diagnostics.py b/tests/unit/train/pretrain/test_diagnostics.py new file mode 100644 index 00000000000..c06ade505ce --- /dev/null +++ b/tests/unit/train/pretrain/test_diagnostics.py @@ -0,0 +1,22 @@ +import math + +import torch + +from prime_rl.trainer.pretrain.diagnostics import MIN_VISIBLE_KEYS, attention_stats + + +def test_attention_stats_respect_documents_and_causality(): + # Two documents; zero queries give uniform attention, so each query's entropy is log(visible keys). + lengths = (MIN_VISIBLE_KEYS + 6, MIN_VISIBLE_KEYS + 2) + num_tokens = sum(lengths) + cu_seqlens = torch.tensor([0, lengths[0], num_tokens], dtype=torch.int32) + q = torch.zeros(num_tokens, 2, 4) + k = torch.randn(num_tokens, 1, 4) + q[-1, 1] = 1.0 # one non-zero query in the second head, last token of the second document + max_logit, entropy_sum, queries = attention_stats(q, k, cu_seqlens, scale=0.5) + + visible = [i + 1 for length in lengths for i in range(length)] + assert int(queries) == sum(v >= MIN_VISIBLE_KEYS for v in visible) + assert torch.isclose(entropy_sum[0], torch.tensor(sum(math.log(v) for v in visible if v >= MIN_VISIBLE_KEYS))) + second_document_keys = k[lengths[0] :, 0] + assert torch.isclose(max_logit[1], ((second_document_keys @ q[-1, 1]).max() * 0.5).clamp(min=0)) # other queries: 0 From ef57f1b2c2a992b90164c2899abaf7909ef43357 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 23:29:44 +0200 Subject: [PATCH 051/132] feat(pretrain): the eval suite runs inside the trainer [evals] scores the eval suite every interval steps with forward passes without gradients on the training model (same sharding, dtype and fusions): packed document-masked rows dealt over the ranks, the LM head swapped for an fp32 log-prob head during the evals, per-token results gathered once on rank 0, logged as eval//@. The snapshot writer, the separate eval worker, its launcher and template integration and the W&B shared-mode plumbing are removed. --- docs/pretraining.md | 12 +- .../src/prime_rl/configs/pretrain.py | 59 ++---- src/prime_rl/dashboard/static/app.js | 2 +- src/prime_rl/entrypoints/pretrain.py | 54 +---- src/prime_rl/monitors/wandb/overview.py | 2 +- .../templates/multi_node_sft.sbatch.j2 | 42 +--- src/prime_rl/trainer/pretrain/eval_worker.py | 192 ------------------ src/prime_rl/trainer/pretrain/snapshots.py | 88 -------- src/prime_rl/trainer/pretrain/suite.py | 10 +- src/prime_rl/trainer/pretrain/train.py | 71 +++++-- src/prime_rl/utils/pathing.py | 5 - tests/unit/train/pretrain/test_launcher.py | 30 +-- 12 files changed, 93 insertions(+), 474 deletions(-) delete mode 100644 src/prime_rl/trainer/pretrain/eval_worker.py delete mode 100644 src/prime_rl/trainer/pretrain/snapshots.py diff --git a/docs/pretraining.md b/docs/pretraining.md index e71f5a88d47..d636e6c0eb4 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -127,22 +127,16 @@ interval = 500 web = { path = "data/web-val" } ``` -**Asynchronous evals on dedicated GPUs.** With `[eval_snapshots]`, the trainer writes a model-only snapshot (bf16, router in fp32) every `interval` steps and at the last step to `/snapshots/step_N`, in the background: it never waits for the eval side. With `[eval_worker]`, the launcher starts the eval worker on GPUs of its own (`num_gpus` of the node in a single-node run, `num_nodes` extra nodes in a multi-node run). The worker takes the newest finished snapshot, scores it on the eval suite (`ppl`, `ctxcont` and `pairs` sets in `//{test.parquet,meta.json}`) at every context length in `seq_lens`, logs `eval/{set}/{metric}@{length}` (e.g. `eval/holdout/bpb@4k`) at the snapshot's step, deletes that snapshot and every older one, and repeats until it has scored the last step. The eval frequency therefore follows from the cost of a round and the number of eval GPUs. Requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. +**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requires `model.cp = 1`. ```toml -[eval_snapshots] -interval = 50 - -[eval_worker] +[evals] +interval = 1000 data_dir = "/shared/evals-data" sets = ["fwedu-holdout", "lambada", "math"] seq_lens = [2048, 4096, 16384] -num_gpus = 1 # single-node runs -num_nodes = 1 # multi-node runs ``` -The trainer and the worker log into one W&B run (W&B shared mode, run id in `/wandb_run_id`) and, with `[monitors.prime]`, into one platform run. The worker can also be started by hand on any GPUs: `torchrun --nproc-per-node 8 -m prime_rl.trainer.pretrain.eval_worker --run-dir --data-dir --sets `. - ## Multi-node `uv run pretrain` takes the same `[deployment]` and `[slurm]` blocks as `uv run sft` (see [Scaling § SLURM](scaling.md#slurm)) and submits the same bundled templates. The launcher writes the model dir on the submitting node, then submits one job in which every node runs `torchrun` on the pretrain trainer: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 0211295bb8b..de67d454e50 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -10,7 +10,6 @@ from prime_rl.configs.shared import ( EnvVars, HeartbeatConfig, - LogConfig, ResumeConfig, RunConfig, SFTDeploymentConfig, @@ -243,16 +242,16 @@ class DiagnosticsConfig(BaseConfig): """Tokens per rank in the probe batch. Fixed across sequence lengths so attention entropy stays comparable.""" -class EvalSnapshotsConfig(BaseConfig): - """Model-only snapshots for the asynchronous eval worker (`prime_rl.trainer.pretrain.eval_worker`), written in - the background to ``/snapshots/step_N``. The worker deletes each one after scoring it.""" +class EvalsConfig(BaseConfig): + """The eval suite, scored inside the trainer: forward passes without gradients on the training model, every + ``interval`` steps. The sets are log-likelihood evals in the format of the pretrain-evals offline suite.""" interval: int = Field(ge=1) - """Write a snapshot every N steps, and at the final step.""" + """Score the suite every N steps (after that step's optimizer update), and at the last step.""" - -class EvalSuiteConfig(BaseConfig): - """What the asynchronous eval worker scores.""" + eval_on_start: bool = True + """Also score the weights the run starts from (the initial weights, logged as step 0, or the checkpoint it + resumes from, logged at that step).""" data_dir: Path """Eval sets, one directory per set with ``test.parquet`` and ``meta.json``.""" @@ -266,12 +265,9 @@ class EvalSuiteConfig(BaseConfig): same at several lengths are scored once.""" token_budget: int = Field(32768, ge=1) - """Tokens per packed row (one forward). At least the longest context length plus one.""" + """Tokens per packed row (one forward). Must exceed the longest context length.""" - poll_seconds: int = Field(30, ge=1) - """Seconds to wait before looking for a new snapshot again.""" - - save_arrays: bool = True + save_arrays: bool = False """Write per-document arrays (npz) and summaries to ``/evals/step_N/``.""" @model_validator(mode="after") @@ -281,27 +277,6 @@ def validate_token_budget(self): return self -class EvalWorkerConfig(EvalSuiteConfig): - """The asynchronous eval worker (``python -m prime_rl.trainer.pretrain.eval_worker``): scores the newest snapshot - of a training run on the eval suite, logs the results into the run's monitors at the snapshot's step, and deletes - scored snapshots. The pretrain launcher writes this config and starts the worker for ``[eval_worker]``.""" - - run_dir: Path - """The training run's directory (with ``configs/latest/resolved/pretrain.json`` and ``snapshots/``).""" - - log: LogConfig = LogConfig() - - -class PretrainEvalWorkerConfig(EvalSuiteConfig): - """The eval worker, started by the launcher next to the trainer on GPUs of its own.""" - - num_gpus: int = Field(1, ge=1) - """Single-node runs: GPUs of the node for the worker, next to the trainer's ``deployment.num_train_gpus``.""" - - num_nodes: int = Field(1, ge=1) - """Multi-node runs: nodes for the worker, next to the trainer's ``deployment.num_train_nodes``.""" - - class PretrainAdamWConfig(AdamWConfig): eps: float = Field(1e-8, gt=0) """Adam epsilon of the base shape (scaled per role under CompleteP).""" @@ -332,12 +307,8 @@ class PretrainConfig(BaseConfig): diagnostics: DiagnosticsConfig | None = DiagnosticsConfig() """Diagnostic step (update sizes, attention and activation probe). If None, it never runs.""" - eval_snapshots: EvalSnapshotsConfig | None = None - """Model-only snapshots for the asynchronous eval worker. If None, none are written.""" - - eval_worker: PretrainEvalWorkerConfig | None = None - """The asynchronous eval worker, started by the launcher on GPUs of its own. Needs ``[eval_snapshots]``. If None, - none is started (a worker can still be started by hand).""" + evals: EvalsConfig | None = None + """The eval suite, scored inside the trainer. If None, it never runs.""" optim: PretrainOptimizerConfig = PretrainMuonConfig() """Optimizer. ``optim.lr`` is required.""" @@ -480,12 +451,6 @@ def validate_offload(self): raise ValueError("Pretraining's per-role parameter groups need model.offload = 'none'.") return self - @model_validator(mode="after") - def validate_eval_worker(self): - if self.eval_worker is not None and self.eval_snapshots is None: - raise ValueError("[eval_worker] needs [eval_snapshots]") - return self - @model_validator(mode="after") def validate_cp(self): if self.model.cp > 1: @@ -495,6 +460,8 @@ def validate_cp(self): raise ValueError("data.micro_batch_tokens must equal data.seq_len with context parallelism") if self.diagnostics is not None: raise ValueError("[diagnostics] does not support context parallelism yet; set diagnostics = null") + if self.evals is not None: + raise ValueError("[evals] does not support context parallelism yet") return self @model_validator(mode="after") diff --git a/src/prime_rl/dashboard/static/app.js b/src/prime_rl/dashboard/static/app.js index c08f30e5d74..db341ffba4e 100644 --- a/src/prime_rl/dashboard/static/app.js +++ b/src/prime_rl/dashboard/static/app.js @@ -450,7 +450,7 @@ const SFT_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memor const PRETRAIN_TRAIN_PANELS = [{ metric: "loss/mean" }, { metric: "loss/spike" }, { regex: "val/.*/bpb" }, { regex: "eval/.*/bpb" }]; const PRETRAIN_STABILITY_PANELS = [{ metric: "optim/grad_norm" }, { metric: "optim/grad_norm_spike" }, { metric: "optim/lr" }, { metric: "loss/nan_count" }, { regex: "optim/grad_norm/.*" }]; const PRETRAIN_MOE_METRICS = ["expert_load/max_mean/mean", "expert_load/max_mean/max", "expert_load/micro_max_mean/mean", "expert_load/micro_max_mean/max", "expert_load/entropy/min", "expert_load/dead", "router/confidence", "router/bias/absmax", "loss/seq_balance"]; -const PRETRAIN_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memory/max", "time/step", "time/data/max", "time/eval_round"]; +const PRETRAIN_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memory/max", "time/step", "time/data/max", "time/eval_suite"]; // Multi-series inference panels (overview.py INFERENCE_PANELS): fleet aggregate // paired with the cross-engine tail that flags a single sick engine. diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py index 7f67363d9b9..3f55b9fe73b 100644 --- a/src/prime_rl/entrypoints/pretrain.py +++ b/src/prime_rl/entrypoints/pretrain.py @@ -3,8 +3,7 @@ import uuid from pathlib import Path -from prime_rl.configs.pretrain import EvalWorkerConfig, PretrainConfig -from prime_rl.configs.shared import LogConfig +from prime_rl.configs.pretrain import PretrainConfig from prime_rl.entrypoints.sft import submit_slurm_script, write_trainer_slurm_script from prime_rl.utils.config import cli, dump_resolved_config from prime_rl.utils.logger import get_logger, setup_logger @@ -12,7 +11,6 @@ format_config_message, format_log_message, get_launcher_dir, - get_wandb_run_id_path, prepare_attempt_dirs, validate_run_dir, write_launch_artifacts, @@ -21,14 +19,12 @@ DEFAULT_COMMON_ENV_VARS, DEFAULT_TRAINER_ENV_VARS, ProcessGroup, - partition_gpus, set_proc_title, torchrun_cmd, ) PRETRAIN_CONFIG = "pretrain.json" PRETRAIN_SBATCH = "pretrain.sbatch" -EVAL_WORKER_CONFIG = "eval_worker.json" def write_model_dir(config: PretrainConfig) -> None: @@ -59,25 +55,10 @@ def write_model_dir(config: PretrainConfig) -> None: tokenizer.save_pretrained(model_dir) -def write_eval_worker_config(config: PretrainConfig, config_dir: Path) -> Path: - """The eval worker's config: the ``[eval_worker]`` suite for this run's directory.""" - assert config.eval_worker is not None - suite = config.eval_worker.model_dump(exclude={"num_gpus", "num_nodes"}) - worker = EvalWorkerConfig(**suite, run_dir=config.run_dir, log=LogConfig(level=config.log.level)) - path = config_dir / EVAL_WORKER_CONFIG - path.write_text(json.dumps(dump_resolved_config(worker), indent=2)) - return path - - def pretrain_slurm(config: PretrainConfig, config_path: Path, log_dir: Path) -> None: logger = get_logger() script_path = get_launcher_dir(config.run_dir) / PRETRAIN_SBATCH - eval_worker_vars = {} - if config.eval_worker is not None and config.deployment.type == "multi_node": - eval_worker_vars = dict( - num_eval_nodes=config.eval_worker.num_nodes, eval_worker_config_path=config_path.parent / EVAL_WORKER_CONFIG - ) - write_trainer_slurm_script(config, "pretrain", config_path, log_dir, script_path, **eval_worker_vars) + write_trainer_slurm_script(config, "pretrain", config_path, log_dir, script_path) logger.info(f"Wrote SLURM script to {script_path}") num_nodes = config.deployment.num_train_nodes if config.deployment.type == "multi_node" else 1 log_message = format_log_message(log_dir=log_dir, trainer=True, num_train_nodes=num_nodes) @@ -110,9 +91,6 @@ def pretrain(config: PretrainConfig): if not config.dry_run: write_model_dir(config) - if config.eval_snapshots is not None and not get_wandb_run_id_path(config.run_dir).exists(): - # The trainer and the eval worker (a separate job) log into one W&B run; a resumed run keeps its id. - get_wandb_run_id_path(config.run_dir).write_text(os.environ["PRL_RUN_ID"]) config_path = config_dir / PRETRAIN_CONFIG exclude = {"dry_run", "clean"} @@ -121,10 +99,7 @@ def pretrain(config: PretrainConfig): # multi-node job read the config directly. exclude |= {"slurm", "deployment"} if config.deployment.type == "multi_node" else {"slurm"} config_path.write_text(json.dumps(dump_resolved_config(config, exclude=exclude), indent=2)) - components = [("Pretrain", config_path)] - if config.eval_worker is not None: - components.append(("Eval worker", write_eval_worker_config(config, config_dir))) - logger.info(f"Configs:\n{format_config_message(config_dir, 'pretrain', components)}") + logger.info(f"Configs:\n{format_config_message(config_dir, 'pretrain', [('Pretrain', config_path)])}") if config.slurm is not None: pretrain_slurm(config, config_path, log_dir) return @@ -133,25 +108,7 @@ def pretrain(config: PretrainConfig): return logger.info(format_log_message(log_dir=log_dir, trainer=True)) - trainer_env = {**os.environ, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, **config.env_vars} with ProcessGroup() as processes: - if config.eval_worker is not None: - # The worker's GPUs first, then the trainer's. - eval_gpu_ids, train_gpu_ids = partition_gpus(config.eval_worker.num_gpus, config.deployment.num_train_gpus) - trainer_env["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, train_gpu_ids)) - processes.start( - "eval-worker", - torchrun_cmd( - "prime_rl.trainer.pretrain.eval_worker", - config_dir / EVAL_WORKER_CONFIG, - config.eval_worker.num_gpus, - log_dir, - config.log.ranks_filter, - ), - env={**trainer_env, "CUDA_VISIBLE_DEVICES": ",".join(map(str, eval_gpu_ids))}, - log_path=log_dir / "eval_worker.log", - ) - logger.success(f"Started the eval worker on GPU(s) {' '.join(map(str, eval_gpu_ids))}") processes.start( "trainer", torchrun_cmd( @@ -161,14 +118,11 @@ def pretrain(config: PretrainConfig): log_dir, config.log.ranks_filter, ), - env=trainer_env, + env={**os.environ, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, **config.env_vars}, log_path=log_dir / "trainer.log", ) logger.success(f"Started trainer with {config.deployment.num_train_gpus} GPU(s)") processes.wait("trainer") - if config.eval_worker is not None: - # The worker scores the final snapshot after the trainer has written it. - processes.wait("eval-worker") logger.success("Pretraining finished!") diff --git a/src/prime_rl/monitors/wandb/overview.py b/src/prime_rl/monitors/wandb/overview.py index 6962828bb29..3b81cc9b9b7 100644 --- a/src/prime_rl/monitors/wandb/overview.py +++ b/src/prime_rl/monitors/wandb/overview.py @@ -113,7 +113,7 @@ "perf/peak_memory/max", "time/step", "time/data/max", - "time/eval_round", + "time/eval_suite", ] # Dense grid: more, smaller panels per row and enough rows that sections don't paginate. diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 index c494970ec95..6dcb0ca9e44 100644 --- a/src/prime_rl/templates/multi_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -1,7 +1,7 @@ #!/bin/bash #SBATCH --job-name={{ job_name }} -#SBATCH --nodes={{ num_train_nodes + num_infer_nodes + (num_eval_nodes or 0) }} +#SBATCH --nodes={{ num_train_nodes + num_infer_nodes }} #SBATCH --ntasks-per-node=1 #SBATCH --gres=gpu:{{ gpus_per_node }} {%- if partition %} @@ -59,11 +59,8 @@ export TRAINER_DONE=$LAUNCHER_DIR/.sft_trainer_done_$SLURM_JOB_ID {% if online_eval -%} export EVAL_DONE=$LAUNCHER_DIR/.sft_eval_done_$SLURM_JOB_ID {%- endif %} -{% if num_eval_nodes -%} -export EVAL_DONE=$LAUNCHER_DIR/.pretrain_eval_done_$SLURM_JOB_ID -{% endif -%} mkdir -p "$LAUNCHER_DIR" -rm -f "$TRAINER_DONE"{% if online_eval or num_eval_nodes %} "$EVAL_DONE"{% endif %} +rm -f "$TRAINER_DONE"{% if online_eval %} "$EVAL_DONE"{% endif %} export PRL_RUN_ID={{ prl_run_id }} export PRL_RUN_NAME={{ run_name }} @@ -158,38 +155,6 @@ srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH' IB_HCA=$(ibv_devinfo 2>/dev/null | sed -n -e "/hca_id/p" -e "/link_layer:/p" | grep -B1 InfiniBand | grep hca_id | sed -e "s/^hca_id://g" | tr -d "[[:blank:]]" | paste -sd,) [ -n "$IB_HCA" ] && export NCCL_IB_HCA=$IB_HCA -{% if num_eval_nodes -%} -if [ "$SLURM_PROCID" -ge "$NUM_TRAIN_NODES" ]; then - # Eval worker nodes follow the trainer nodes. - EVAL_NODE_RANK=$((SLURM_PROCID - NUM_TRAIN_NODES)) - HOSTS=( $ALL_HOSTNAMES_STR ) -{%- for key, value in trainer_env_vars.items() %} - export {{ key }}="{{ value }}" -{%- endfor %} - mkdir -p $LOG_DIR/eval_worker - ( - uv run torchrun \ - --role=eval \ - --nnodes={{ num_eval_nodes }} \ - --nproc-per-node=$GPUS_PER_NODE \ - --node-rank=$EVAL_NODE_RANK \ - --rdzv-endpoint=${HOSTS[$NUM_TRAIN_NODES]}:29600 \ - --rdzv-id=eval_$SLURM_JOB_ID \ - --log-dir=$LOG_DIR/eval_worker/torchrun \ - --tee=3 \ - --redirects=3 \ - --local-ranks-filter=0 \ - -m prime_rl.trainer.pretrain.eval_worker \ - @ {{ eval_worker_config_path }} \ - 2>&1 | sed -u 's/^\[[a-zA-Z]*[0-9]*\]://' | tee -a $LOG_DIR/eval_worker/node_${EVAL_NODE_RANK}.log - COMPONENT_CODE=$? - if [ "$COMPONENT_CODE" -eq 0 ] && [ "$EVAL_NODE_RANK" -eq 0 ]; then - touch "$EVAL_DONE" || exit $? - fi - exit "$COMPONENT_CODE" - ) & -else -{% endif -%} {% if online_eval -%} if [ "$SLURM_PROCID" -lt "$NUM_INFER_NODES" ]; then INFER_NODE_RANK=$SLURM_PROCID @@ -289,10 +254,9 @@ else exit "$COMPONENT_CODE" ) & {% if online_eval %}fi{% endif %} -{%- if num_eval_nodes %}fi{% endif %} # Every node stays in the shared srun step until all terminal roles finish. -{% if online_eval or num_eval_nodes -%} +{% if online_eval -%} ( until [ -f "$TRAINER_DONE" ] && [ -f "$EVAL_DONE" ]; do sleep 10; done ) & {% else -%} ( until [ -f "$TRAINER_DONE" ]; do sleep 10; done ) & diff --git a/src/prime_rl/trainer/pretrain/eval_worker.py b/src/prime_rl/trainer/pretrain/eval_worker.py deleted file mode 100644 index 9e97841a43f..00000000000 --- a/src/prime_rl/trainer/pretrain/eval_worker.py +++ /dev/null @@ -1,192 +0,0 @@ -"""Asynchronous eval worker for pretraining runs: on its own GPUs, it scores the newest model snapshot of a -running training run (`[eval_snapshots]` in the run's config) on the eval suite, logs the results into the run's -W&B run and `metrics.jsonl` at the snapshot's step, deletes the snapshot and every older one, and repeats until -it has scored the final step. Started with torchrun on the eval GPUs: - - torchrun --nproc-per-node 8 -m prime_rl.trainer.pretrain.eval_worker --run-dir --data-dir --sets - -The model is built once with the trainer's code (`setup_model` from the run's resolved config) and kept resident; -each round loads only the weights. Rows of every set are scored data-parallel over the worker's GPUs.""" - -import asyncio -import json -import os -import time -from datetime import timedelta -from pathlib import Path - -import numpy as np -import torch -import torch.distributed as dist -from torch.distributed.checkpoint.state_dict_loader import load as dcp_load -from transformers import AutoTokenizer - -from prime_rl import monitors -from prime_rl.configs.monitors import PrimeEvalMonitorConfig -from prime_rl.configs.pretrain import EvalWorkerConfig, PretrainConfig -from prime_rl.monitors.prime import read_platform_record -from prime_rl.trainer.ckpt import AppState -from prime_rl.trainer.model import setup_model -from prime_rl.trainer.parallel_dims import get_parallel_dims -from prime_rl.trainer.pretrain.snapshots import ( - delete_snapshots_up_to, - finished_snapshot_steps, - shared_wandb_env, - snapshot_dir, -) -from prime_rl.trainer.pretrain.suite import ( - length_label, - prepare_multi_length_set, - score_requests, - summarize, - use_scoring_head, -) -from prime_rl.trainer.utils import get_dcp_group, setup_torch_distributed -from prime_rl.trainer.world import get_world -from prime_rl.utils.config import cli -from prime_rl.utils.logger import setup_logger -from prime_rl.utils.pathing import get_file_monitor_dir, get_platform_run_path -from prime_rl.utils.process import set_proc_title -from prime_rl.utils.utils import clean_exit - -RUN_CONFIG = "configs/latest/resolved/pretrain.json" - - -def run_config(config: EvalWorkerConfig) -> PretrainConfig: - """The training run's resolved config (waiting for the launcher to write it), with the parallelism of the - worker: plain FSDP over its GPUs.""" - path = config.run_dir / RUN_CONFIG - while not path.exists(): - time.sleep(config.poll_seconds) - run = PretrainConfig.model_validate_json(path.read_text()) - run.model.compile = None # packed rows of varying length would recompile at every forward - run.model.ep = 1 - run.model.cp = 1 - run.model.dp_replicate = 1 - return run - - -def platform_run_id(run_dir: Path, run: PretrainConfig, poll_seconds: int) -> str: - """The id of the trainer's platform run, which its monitor writes to the run dir when it starts, before the - trainer logs any metrics. Only state of the current launch counts (a resumed run dir holds an earlier - attempt's record and metrics). Fails if the trainer logs metrics of this launch without it (e.g. platform - runs disabled).""" - launched = (run_dir / RUN_CONFIG).stat().st_mtime - record_path = get_platform_run_path(run_dir) - metrics_path = get_file_monitor_dir(run_dir) / run.monitors.file.path if run.monitors.file is not None else None - while True: - record = read_platform_record(run_dir) - if record is not None and "id" in record and record_path.stat().st_mtime >= launched: - return record["id"] - if metrics_path is not None and metrics_path.is_file() and last_row_time(metrics_path) >= launched: - raise RuntimeError(f"The trainer logs metrics but left no platform run id for this launch in {run_dir}") - time.sleep(poll_seconds) - - -def last_row_time(path: Path) -> float: - """The ``time`` of the last complete row of a ``metrics.jsonl``.""" - with path.open("rb") as f: - f.seek(max(0, f.seek(0, 2) - 65536)) - rows = f.read().splitlines() - return json.loads(rows[-1])["time"] if rows else 0.0 - - -def broadcast(value: int) -> int: - values = [value] - dist.broadcast_object_list(values, src=0, group=get_dcp_group()) - return values[0] - - -@clean_exit -def main_loop(config: EvalWorkerConfig) -> None: - world = get_world() - logger = setup_logger(config.log.level, json_logging=config.log.json_logging) - run = run_config(config) - setup_torch_distributed(timeout=timedelta(seconds=run.dist_timeout_seconds)) - torch.set_float32_matmul_precision(run.matmul_precision) - - os.environ.update(shared_wandb_env(config.run_dir, "eval")) - prime = None - if run.monitors.prime is not None: - # Metrics onto the trainer's platform run, which the trainer owns: the eval monitor only attaches to it - # and never reports its status. - os.environ["RUN_ID"] = platform_run_id(config.run_dir, run, config.poll_seconds) - prime = PrimeEvalMonitorConfig(name=run.monitors.prime.name) - asyncio.run( - monitors.setup( - producer="eval", - wandb=run.monitors.wandb, - prime=prime, - file=run.monitors.file, - output_dir=config.run_dir, - overview_flavor="pretrain", - ) - ) - - model = setup_model(run.model, get_parallel_dims(run.model), loading_from_checkpoint_later=True) - model.eval() - use_scoring_head(model) - tokenizer = AutoTokenizer.from_pretrained(run.tokenizer.name, trust_remote_code=run.tokenizer.trust_remote_code) - start = time.perf_counter() - eval_sets = [prepare_multi_length_set(name, config.data_dir, tokenizer, config.seq_lens) for name in config.sets] - num_tokens = sum(len(ids) - begin for eval_set in eval_sets for ids, begin in eval_set.requests) - logger.info( - f"Prepared {len(eval_sets)} sets at context lengths {config.seq_lens} ({num_tokens:,} distinct scored tokens) " - f"in {time.perf_counter() - start:.0f}s" - ) - last_step = -1 - while True: - newest = broadcast(max(finished_snapshot_steps(config.run_dir), default=-1) if world.is_master else 0) - if newest <= last_step: - if run.max_steps is not None and last_step >= run.max_steps: - break - time.sleep(config.poll_seconds) - continue - - round_start = time.perf_counter() - dcp_load( - {"app": AppState(model, [], None, None)}, - checkpoint_id=snapshot_dir(config.run_dir) / f"step_{newest}", - process_group=get_dcp_group(), - ) - load_time = time.perf_counter() - round_start - metrics = {} - for eval_set in eval_sets: - scored = score_requests(model, eval_set.requests, config.token_budget, get_dcp_group()) - if not world.is_master: - continue - logprobs, hits = scored - for seq_len, version in eval_set.versions.items(): - index = eval_set.index[seq_len] - summary, arrays = summarize(version, [logprobs[i] for i in index], [hits[i] for i in index]) - label = length_label(seq_len) - metrics |= { - f"eval/{eval_set.name}/{key}@{label}": value - for key, value in summary.items() - if isinstance(value, float) - } - if config.save_arrays: - out = config.run_dir / "evals" / f"step_{newest}" - out.mkdir(parents=True, exist_ok=True) - np.savez(out / f"{eval_set.name}@{label}.npz", **arrays) - (out / f"{eval_set.name}@{label}.json").write_text( - json.dumps(summary | {"data_sha256": version.sha256}, default=float) - ) - if world.is_master: - round_time = time.perf_counter() - round_start - metrics |= {"time/eval_round": round_time, "time/eval_load": load_time, "step": newest} - asyncio.run(monitors.log(metrics, step=newest)) - delete_snapshots_up_to(config.run_dir, newest) - logger.success(f"Evals | Step {newest} | {round_time:.0f}s | {len(eval_sets)} sets") - last_step = newest - - asyncio.run(monitors.finalize()) - - -def main(): - set_proc_title("PretrainEvalWorker") - main_loop(cli(EvalWorkerConfig)) - - -if __name__ == "__main__": - main() diff --git a/src/prime_rl/trainer/pretrain/snapshots.py b/src/prime_rl/trainer/pretrain/snapshots.py deleted file mode 100644 index e977067b460..00000000000 --- a/src/prime_rl/trainer/pretrain/snapshots.py +++ /dev/null @@ -1,88 +0,0 @@ -import shutil -from concurrent.futures import Future, ThreadPoolExecutor -from pathlib import Path - -import torch -from torch import nn -from torch.distributed.checkpoint.state_dict_saver import save as dcp_save - -from prime_rl.trainer.ckpt import AppState -from prime_rl.trainer.utils import get_dcp_group -from prime_rl.trainer.world import get_world -from prime_rl.utils.pathing import get_wandb_run_id_path - -FINISHED_MARKER = ".finished" - - -def shared_wandb_env(run_dir: Path, label: str) -> dict[str, str]: - """W&B shared-mode settings that put the trainer (``label = "trainer"``, creates the run) and the eval worker - (``"eval"``, finishes it) into one W&B run, whose id the launcher writes to ``/wandb_run_id``.""" - return { - "WANDB_SHARED_MODE": "1", - "WANDB_RUN_ID": get_wandb_run_id_path(run_dir).read_text().strip(), - "WANDB_SHARED_LABEL": label, - "WANDB_SHARED_PRIMARY": "trainer", - "WANDB_SHARED_FINISHER": "eval", - } - - -def snapshot_dir(run_dir: Path) -> Path: - return run_dir / "snapshots" - - -def finished_snapshot_steps(run_dir: Path) -> list[int]: - """Steps of the complete snapshots, ascending.""" - paths = snapshot_dir(run_dir).glob(f"step_*/{FINISHED_MARKER}") - return sorted(int(path.parent.name.removeprefix("step_")) for path in paths) - - -def delete_snapshots_up_to(run_dir: Path, step: int) -> None: - for path in snapshot_dir(run_dir).glob("step_*"): - if int(path.name.removeprefix("step_")) <= step: - shutil.rmtree(path, ignore_errors=True) - - -def snapshot_tensor(name: str, tensor: torch.Tensor) -> torch.Tensor: - """A copy of one model tensor for the snapshot, in bf16 (the dtype the forward computes in, so scoring is - unchanged and the snapshot is half the size), except the router (gate weights and selection bias), which - runs in fp32 and decides the expert selection.""" - if tensor.is_floating_point() and "router" not in name: - return tensor.to(torch.bfloat16, copy=True) - return tensor.clone() - - -class SnapshotWriter: - """Model-only DCP snapshots for the eval worker (`prime_rl.trainer.pretrain.eval_worker`), in the layout of a - model-only checkpoint (`{"app": AppState}`), in bf16 except the router (`snapshot_tensor`); loading casts back - to the model's dtypes. The trainer only waits for a copy of its weight shards on the GPU; - a background thread writes them. A snapshot is complete once its `.finished` marker exists; the eval worker - deletes it after scoring it.""" - - def __init__(self, run_dir: Path): - self.run_dir = run_dir - # A new thread starts on CUDA device 0; DCP's GPU-to-CPU copies must run on this rank's device. - self.executor = ThreadPoolExecutor( - max_workers=1, initializer=torch.cuda.set_device, initargs=(torch.cuda.current_device(),) - ) - self.pending: Future | None = None - # Created here, on every rank at the same point; the writer thread only uses it. - self.group = get_dcp_group() - - def save(self, model: nn.Module, step: int) -> None: - self.wait() - state = AppState(model, [], None, None).state_dict() - state["model"] = {name: snapshot_tensor(name, tensor) for name, tensor in state["model"].items()} - # DCP copies the tensors to the CPU on its own CUDA stream, which does not wait for the copies above. - torch.cuda.current_stream().synchronize() - self.pending = self.executor.submit(self.write, state, snapshot_dir(self.run_dir) / f"step_{step}") - - def write(self, state: dict, path: Path) -> None: - dcp_save({"app": state}, checkpoint_id=path, process_group=self.group) - if get_world().is_master: - (path / FINISHED_MARKER).touch() - - def wait(self) -> None: - """Block until the previous snapshot is on disk.""" - if self.pending is not None: - self.pending.result() - self.pending = None diff --git a/src/prime_rl/trainer/pretrain/suite.py b/src/prime_rl/trainer/pretrain/suite.py index 5ef6c71eb0a..f7c2ba9830c 100644 --- a/src/prime_rl/trainer/pretrain/suite.py +++ b/src/prime_rl/trainer/pretrain/suite.py @@ -15,6 +15,8 @@ import json import math import types +from collections.abc import Iterator +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Any @@ -70,8 +72,14 @@ def scoring_head(self, hidden_states, labels, temperature=None, sampling_mask=No return {"logprobs": logprobs, "entropy": hits} -def use_scoring_head(model: nn.Module) -> None: +@contextmanager +def scoring_head_active(model: nn.Module) -> Iterator[None]: + """The model's LM head returns per-token log-probs and greedy hits (`scoring_head`) inside the block.""" model.lm_head.forward = types.MethodType(scoring_head, model.lm_head) + try: + yield + finally: + del model.lm_head.forward def pack_requests(requests: list[Request], token_budget: int) -> list[list[int]]: diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 420776614a7..c36eb603df3 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -1,8 +1,8 @@ import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before import import asyncio +import json import itertools -import os import math import time from datetime import timedelta @@ -10,6 +10,7 @@ # ruff: noqa: I001 from prime_rl.utils.act_offloading import maybe_activation_offloading +import numpy as np import torch import torch.distributed as dist from torch import nn @@ -46,10 +47,17 @@ from prime_rl.trainer.pretrain.schedules import LRSchedule from prime_rl.trainer.data import get_dataset_progress, get_dataset_state from prime_rl.trainer.pretrain.diagnostics import UpdateTracker, probe_batch, run_probe -from prime_rl.trainer.pretrain.snapshots import SnapshotWriter, shared_wandb_env +from prime_rl.trainer.pretrain.suite import ( + length_label, + prepare_multi_length_set, + score_requests, + scoring_head_active, + summarize, +) from prime_rl.trainer.pretrain.step import run_micro_batches from prime_rl.trainer.pretrain.metrics import DeadExperts, SpikeDetector, grad_norms_by_group from prime_rl.trainer.utils import ( + get_dcp_group, GarbageCollection, clip_grad_norm_, get_ckpt_disk_metrics, @@ -94,8 +102,6 @@ def train(config: PretrainConfig): logger = setup_logger(config.log.level, json_logging=config.log.json_logging) logger.info(f"Starting pretrain trainer in {world} (output_dir={config.run_dir})") - if config.eval_snapshots is not None: - os.environ.update(shared_wandb_env(config.run_dir, "trainer")) asyncio.run( monitors.setup( producer="trainer", @@ -292,15 +298,60 @@ def run_evals(step: int) -> None: logger.success(f"Evals | Step {step} | " + " | ".join(f"{k} {v:.4f}" for k, v in sorted(metrics.items()))) asyncio.run(monitors.log({**metrics, "step": step}, step=step)) + eval_sets = [] + if config.evals is not None: + start_time = time.perf_counter() + eval_sets = [ + prepare_multi_length_set(name, config.evals.data_dir, tokenizer, config.evals.seq_lens) + for name in config.evals.sets + ] + logger.info(f"Prepared {len(eval_sets)} eval sets in {time.perf_counter() - start_time:.0f}s") + + @torch.no_grad() + def run_suite(step: int) -> None: + """The eval suite on the weights after ``step`` optimizer updates: the packed rows of every set are dealt + over the ranks, the per-token results gathered on rank 0, which summarizes and logs them.""" + start_time = time.perf_counter() + metrics = {} + with scoring_head_active(model): + for eval_set in eval_sets: + scored = score_requests(model, eval_set.requests, config.evals.token_budget, get_dcp_group()) + if scored is None: + continue + logprobs, hits = scored + for seq_len, version in eval_set.versions.items(): + index = eval_set.index[seq_len] + summary, arrays = summarize(version, [logprobs[i] for i in index], [hits[i] for i in index]) + label = length_label(seq_len) + metrics |= { + f"eval/{eval_set.name}/{key}@{label}": value + for key, value in summary.items() + if isinstance(value, float) + } + if config.evals.save_arrays: + out = config.run_dir / "evals" / f"step_{step}" + out.mkdir(parents=True, exist_ok=True) + np.savez(out / f"{eval_set.name}@{label}.npz", **arrays) + (out / f"{eval_set.name}@{label}.json").write_text(json.dumps(summary, default=float)) + # The no-grad forwards leave the [lm_head, norm] FSDP group unsharded, and their routing must not count + # toward the next step's MoE statistics. + reshard_module(model) + if is_moe_model: + get_load_balance_stats(model) + metrics["time/eval_suite"] = time.perf_counter() - start_time + logger.success(f"Eval suite | Step {step} | {metrics['time/eval_suite']:.0f}s | {len(eval_sets)} sets") + asyncio.run(monitors.log({**metrics, "step": step}, step=step)) + if checkpoint_step is None and config.val is not None and config.val.eval_on_start: run_evals(step=0) + if config.evals is not None and config.evals.eval_on_start: + run_suite(step=checkpoint_step or 0) gc_handler = GarbageCollection(config.gc.interval) if config.gc else None max_memory = torch.cuda.mem_get_info()[1] / 1024**3 dead_experts = DeadExperts() loss_spikes, grad_norm_spikes = SpikeDetector(), SpikeDetector() last_bias_moves = None - snapshots = SnapshotWriter(config.run_dir) if config.eval_snapshots is not None else None update_tracker = UpdateTracker() logger.info(f"Starting training loop (max_steps={config.max_steps})") while True: @@ -358,9 +409,6 @@ def run_evals(step: int) -> None: disk_metrics = {} if config.ckpt and config.ckpt.interval and progress.step % config.ckpt.interval == 0 and not is_last_step: save_ckpt_start_time = time.perf_counter() - if snapshots is not None: - # Both write through the same DCP process group; their collectives must not interleave. - snapshots.wait() ckpt_manager.save(progress.step, model, [optimizer], None, progress, dataloader=dataloader) ckpt_manager.maybe_clean() save_ckpt_time = time.perf_counter() - save_ckpt_start_time @@ -375,12 +423,11 @@ def run_evals(step: int) -> None: if is_moe_model: get_load_balance_stats(model) - if snapshots is not None and (is_last_step or progress.step % config.eval_snapshots.interval == 0): - snapshots.save(model, progress.step) - eval_start_time = time.perf_counter() if config.val is not None and (is_last_step or progress.step % config.val.interval == 0): run_evals(progress.step) + if config.evals is not None and (is_last_step or progress.step % config.evals.interval == 0): + run_suite(progress.step) eval_time = time.perf_counter() - eval_start_time num_tokens = result.num_tokens @@ -439,8 +486,6 @@ def run_evals(step: int) -> None: break progress.step += 1 - if snapshots is not None: - snapshots.wait() if config.ckpt is not None: logger.info(f"Saving final checkpoint at step {progress.step}") ckpt_manager.save(progress.step, model, [optimizer], None, progress, dataloader=dataloader) diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index 479eacd9aa6..7dc00a77993 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -284,11 +284,6 @@ def get_eval_dir(output_dir: Path) -> Path: return output_dir / "eval" -def get_wandb_run_id_path(output_dir: Path) -> Path: - """The W&B run id that a pretraining run's trainer and its eval worker (separate jobs) share.""" - return output_dir / "wandb_run_id" - - def get_broadcast_dir(output_dir: Path) -> Path: return output_dir / "broadcasts" diff --git a/tests/unit/train/pretrain/test_launcher.py b/tests/unit/train/pretrain/test_launcher.py index db4e4750c7c..9238d16d42e 100644 --- a/tests/unit/train/pretrain/test_launcher.py +++ b/tests/unit/train/pretrain/test_launcher.py @@ -1,6 +1,6 @@ import re -from prime_rl.configs.pretrain import EvalWorkerConfig, PretrainConfig +from prime_rl.configs.pretrain import PretrainConfig from prime_rl.entrypoints.pretrain import pretrain from prime_rl.utils.config import cli @@ -31,31 +31,3 @@ def test_multi_node_dry_run_writes_trainer_sbatch(tmp_path, monkeypatch): trainer_config = cli(PretrainConfig, args=["@", config_path]) assert trainer_config.slurm is None assert trainer_config.model == config.model - - -def test_multi_node_dry_run_adds_the_eval_worker_node(tmp_path, monkeypatch): - monkeypatch.setenv("PRL_RUN_ID", "test") - monkeypatch.setenv("PRL_RUN_NAME", "test") - config = PretrainConfig.model_validate( - { - "output_dir": tmp_path, - "model": {"arch": {"model_type": "llama"}}, - "tokenizer": {"name": "PrimeIntellect/Qwen3-0.6B"}, - "data": {"type": "fake", "seq_len": 8, "batch_tokens": 16, "micro_batch_tokens": 8}, - "optim": {"lr": 1e-3}, - "max_steps": 1, - "eval_snapshots": {"interval": 1}, - "eval_worker": {"data_dir": tmp_path, "sets": ["holdout"], "seq_lens": [2048, 4096]}, - "deployment": {"type": "multi_node", "num_train_nodes": 2}, - "slurm": {}, - "dry_run": True, - } - ) - pretrain(config) - - script = (config.run_dir / "launcher" / "pretrain.sbatch").read_text() - assert "#SBATCH --nodes=3" in script - worker_config_path = re.search(r"-m prime_rl.trainer.pretrain.eval_worker \\\n\s+@ (\S+)", script).group(1) - worker_config = cli(EvalWorkerConfig, args=["@", worker_config_path]) - assert worker_config.run_dir == config.run_dir - assert worker_config.seq_lens == [2048, 4096] From 72c5dd2051a5d5d95e274f7d22aeec1647965214 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 23:36:26 +0200 Subject: [PATCH 052/132] fix(pretrain): the diagnostic probe and the eval suite run compiled layers eagerly; probe skips attention stats without FlashAttention Compiled layers would not see the probe's attention hook and would recompile for every packed eval row; both now run under torch.compiler.set_stance("force_eager"). Without FlashAttention layers the probe logs a warning once and records only the activation statistics and the probe loss. --- src/prime_rl/trainer/pretrain/diagnostics.py | 54 +++++++++++++------- src/prime_rl/trainer/pretrain/train.py | 3 +- 2 files changed, 38 insertions(+), 19 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/diagnostics.py b/src/prime_rl/trainer/pretrain/diagnostics.py index bb9be1bbd22..fdfaa3f5023 100644 --- a/src/prime_rl/trainer/pretrain/diagnostics.py +++ b/src/prime_rl/trainer/pretrain/diagnostics.py @@ -1,3 +1,5 @@ +from functools import cache + import torch from torch import Tensor, nn @@ -5,6 +7,7 @@ from prime_rl.trainer.models.layers.attn import FlashAttention from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX from prime_rl.trainer.pretrain.metrics import group_norms, param_group, reduce_stats +from prime_rl.utils.logger import get_logger from prime_rl.utils.vlm import get_language_model MIN_VISIBLE_KEYS = 64 @@ -47,13 +50,22 @@ def attention_stats(q: Tensor, k: Tensor, cu_seqlens: Tensor, scale: float) -> t return torch.stack(max_logits), torch.stack(entropy_sums), valid.sum() +@cache +def warn_no_attention_stats() -> None: + get_logger().warning( + "The diagnostic probe records attention statistics only for FlashAttention layers; none found, skipping them" + ) + + @torch.no_grad() def run_probe(model: nn.Module, batch: dict[str, Tensor], group) -> dict[str, float]: - """Forward the probe batch with hooks: attention max logit and entropy per layer and head, residual-stream RMS - and max per layer, and the probe loss. Every rank probes its own batch; the statistics are reduced over - ``group``.""" + """Forward the probe batch with hooks: attention max logit and entropy per layer and head (FlashAttention + layers), residual-stream RMS and max per layer, and the probe loss. Compiled layers run eagerly for the probe, + so the hooks see every call. Every rank probes its own batch; the statistics are reduced over ``group``.""" layers = get_language_model(model).layers attentions = [module for module in layers.modules() if isinstance(module, FlashAttention)] + if not attentions: + warn_no_attention_stats() attention, activations = [], [] def capture(module: FlashAttention): @@ -75,43 +87,49 @@ def record_activations(module, args, output): module._compute_attention = capture(module) hooks = [layer.register_forward_hook(record_activations) for layer in layers] labels = batch["target_ids"] - out = forward(model, batch["input_ids"], batch["position_ids"], seq_lens=batch["seq_lens"], labels=labels) + with torch.compiler.set_stance("force_eager"): + out = forward(model, batch["input_ids"], batch["position_ids"], seq_lens=batch["seq_lens"], labels=labels) for module in attentions: del module._compute_attention for hook in hooks: hook.remove() - max_logit, entropy_sum, queries = (torch.stack(values) for values in zip(*attention)) square_sum, count, absmax = (torch.stack(values) for values in zip(*activations)) sums = { - "entropy_sum": entropy_sum, - "queries": queries.double(), "square_sum": square_sum, "count": count.double(), "loss": out["loss"].double(), "tokens": (labels != IGNORE_INDEX).sum().double(), } - maxes = {"max_logit": max_logit, "absmax": absmax} + maxes = {"absmax": absmax} + if attention: + max_logit, entropy_sum, queries = (torch.stack(values) for values in zip(*attention)) + sums |= {"entropy_sum": entropy_sum, "queries": queries.double()} + maxes["max_logit"] = max_logit stats = reduce_stats(sums, maxes, group) - head_entropy = stats["entropy_sum"] / stats["queries"][:, None] # [layers, heads] - layer_max_logit = stats["max_logit"].amax(dim=1) - layer_entropy = head_entropy.mean(dim=1) rms = (stats["square_sum"] / stats["count"]).sqrt() metrics = { - "attn/max_logit/max": layer_max_logit.max().item(), - "attn/entropy/min": layer_entropy.min().item(), - "attn/entropy_min_head": head_entropy.min().item(), - "attn/collapsed_heads": int((head_entropy < COLLAPSED_ENTROPY).sum()), "act/rms/max": rms.max().item(), "act/absmax/max": stats["absmax"].max().item(), "probe/loss": (stats["loss"] / stats["tokens"]).item(), } - for layer in range(len(layer_max_logit)): - metrics[f"attn/max_logit/layer_{layer}"] = layer_max_logit[layer].item() - metrics[f"attn/entropy/layer_{layer}"] = layer_entropy[layer].item() + for layer in range(len(rms)): metrics[f"act/rms/layer_{layer}"] = rms[layer].item() metrics[f"act/absmax/layer_{layer}"] = stats["absmax"][layer].item() + if attention: + head_entropy = stats["entropy_sum"] / stats["queries"][:, None] # [layers, heads] + layer_max_logit = stats["max_logit"].amax(dim=1) + layer_entropy = head_entropy.mean(dim=1) + metrics |= { + "attn/max_logit/max": layer_max_logit.max().item(), + "attn/entropy/min": layer_entropy.min().item(), + "attn/entropy_min_head": head_entropy.min().item(), + "attn/collapsed_heads": int((head_entropy < COLLAPSED_ENTROPY).sum()), + } + for layer in range(len(layer_max_logit)): + metrics[f"attn/max_logit/layer_{layer}"] = layer_max_logit[layer].item() + metrics[f"attn/entropy/layer_{layer}"] = layer_entropy[layer].item() return metrics diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index c36eb603df3..f3e629f0377 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -313,7 +313,8 @@ def run_suite(step: int) -> None: over the ranks, the per-token results gathered on rank 0, which summarizes and logs them.""" start_time = time.perf_counter() metrics = {} - with scoring_head_active(model): + # Eager, also for compiled layers: the packed rows vary in length and would recompile at every forward. + with scoring_head_active(model), torch.compiler.set_stance("force_eager"): for eval_set in eval_sets: scored = score_requests(model, eval_set.requests, config.evals.token_budget, get_dcp_group()) if scored is None: From e8ff0591354b5dd9983afe6e5408a9230b4ffb9d Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 14 Oct 2026 23:36:39 +0200 Subject: [PATCH 053/132] docs(pretrain): resume needs the same GPU layout The data-loader state only loads on the same layout and data.num_workers; resuming across layouts belongs to the production data loader. --- docs/pretraining.md | 2 +- src/prime_rl/trainer/pretrain/data.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index ba644b97b5c..e2d9a8ee616 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -74,7 +74,7 @@ start = 0.9 weights = { web = 0.5, code = 0.5 } ``` -Batch sizes are in tokens, so changing `seq_len` keeps the tokens per step and the number of steps. The resume state is the global row counter and the tokens read per source, the same on every rank. `data.type = "fake"` trains on random tokens, for smoke and throughput tests. +Batch sizes are in tokens, so changing `seq_len` keeps the tokens per step and the number of steps. Resume restores the data position on the same GPU layout and `data.num_workers`; resuming on a different layout is not supported yet. `data.type = "fake"` trains on random tokens, for smoke and throughput tests. ## Evals diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index 2ca6e4bb4dc..ec624b2ff9e 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -85,8 +85,8 @@ class TokenDataset(StatefulIterableDataset): Window ``w`` belongs to step ``w // batch_size``; its source is drawn from that step's phase weights with an RNG keyed by ``(seed, w)``, and it continues the source's stream where that source's last - window ended. Every data rank walks the whole window sequence but reads only its own windows, so the - state (the window counter and the tokens read per source) does not depend on the world size. Each + window ended. Every data rank walks the whole window sequence but reads only its own windows; the + state is the window counter and the tokens read per source. Each window reads one token past its end for the last target. With ``num_windows`` the dataset stops after that many windows.""" From 29c3bad31ce30c4bdaae93d1b89f908c0ef1ec86 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 09:53:52 +0200 Subject: [PATCH 054/132] fix(pretrain): global per-source token counts from the window sequence progress//num_tokens came from one data-loader worker's state, so it didn't count every rank's rows. Step s trains on the global windows [(s - 1) * batch_size, s * batch_size), and every data rank walks the same window sequence, so the loop counts each step's tokens per source from that sequence without communication: progress//step_tokens, num_tokens (cumulative) and epochs. A 2-GPU smoke run and its resume from step 20 on research-india give identical counts. --- src/prime_rl/trainer/pretrain/data.py | 9 ++++++ src/prime_rl/trainer/pretrain/train.py | 41 +++++++++++++++++++------- tests/unit/train/pretrain/test_data.py | 13 +++++++- 3 files changed, 51 insertions(+), 12 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index ec624b2ff9e..b5f9a660437 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -157,6 +157,15 @@ def __iter__(self): } +def source_tokens(dataset: TokenDataset, start: int, end: int) -> dict[str, int]: + """Tokens each source contributes to windows ``[start, end)``. Every data rank walks the same window + sequence, so these are global counts, computed without communication.""" + tokens = dict.fromkeys(dataset.names, 0) + for window in range(start, end): + tokens[dataset.source(window)] += dataset.seq_len + return tokens + + def collate_rows(rows: list[dict]) -> dict: """Concatenate windows into one packed sample of shape ``[1, len(rows) * seq_len]``. Positions restart at every document; attention never crosses a document or a window.""" diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index b6feedb89a9..bc88d4983c6 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -38,10 +38,11 @@ check_eos, setup_dataloader, setup_train_dataset, + source_tokens, setup_val_dataset, ) from prime_rl.trainer.pretrain.schedules import LRSchedule -from prime_rl.trainer.sft.data import get_dataset_progress, get_dataset_state +from prime_rl.trainer.sft.data import get_dataset_state from prime_rl.trainer.utils import ( GarbageCollection, begin_backward, @@ -150,13 +151,10 @@ def train(config: PretrainConfig): schedule_offset = checkpoint_step if restart_schedule else 0 lr_schedule = LRSchedule(config.schedule, optimizer, config.max_steps - schedule_offset, schedule_offset) - dataloader = setup_dataloader( - setup_train_dataset( - config.data, config.max_steps, len(tokenizer), tokenizer.eos_token_id, non_dp_size=config.model.cp - ), - config.data.micro_batch_size, - config.data.num_workers, + train_dataset = setup_train_dataset( + config.data, config.max_steps, len(tokenizer), tokenizer.eos_token_id, non_dp_size=config.model.cp ) + dataloader = setup_dataloader(train_dataset, config.data.micro_batch_size, config.data.num_workers) val_sources = {} if config.val is not None: @@ -187,6 +185,20 @@ def train(config: PretrainConfig): progress.step += 1 logger.info(f"Resuming from step {checkpoint_step} (dataset_state={get_dataset_state(dataloader)})") + # Step s trains on the global windows [(s - 1) * batch_size, s * batch_size). The tokens per source are counted + # from the window sequence itself, which every data rank walks, so they are global without communication. + next_window = 0 + source_totals = dict.fromkeys(train_dataset.names, 0) + if checkpoint_step is not None and not (config.ckpt or CheckpointConfig()).skip_dataloader: + next_window = checkpoint_step * config.data.batch_size + # The loaded state counts the windows before its own position, which is near next_window. + state = next(iter(get_dataset_state(dataloader).values())) + source_totals = dict(state["tokens"]) + start, end = sorted((state["step"], next_window)) + sign = 1 if next_window > state["step"] else -1 + for name, tokens in source_tokens(train_dataset, start, end).items(): + source_totals[name] += sign * tokens + # Create the iterator only after a potential resume: iter() forks the workers with the # dataset's current state. dataiter = iter(dataloader) @@ -345,8 +357,11 @@ def run_evals(step: int) -> None: dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) num_tokens = global_num_tokens.item() progress.total_tokens += num_tokens - dataset_progress = get_dataset_progress(dataloader) - progress.total_samples = dataset_progress["step"] + step_tokens = source_tokens(train_dataset, next_window, next_window + config.data.batch_size) + next_window += config.data.batch_size + for name, tokens in step_tokens.items(): + source_totals[name] += tokens + progress.total_samples = next_window step_time = time.perf_counter() - step_start_time - eval_time perf_counter = get_perf_counter(model, config.data.seq_len) perf_counter.count_tokens(num_tokens) @@ -368,8 +383,12 @@ def run_evals(step: int) -> None: "loss/perplexity": math.exp(min(batch_loss, 20)), "loss/nan_count": nan_loss_count, "optim/lr": current_lr, - **{f"progress/{name}/epochs": epochs for name, epochs in dataset_progress["epochs"].items()}, - **{f"progress/{name}/num_tokens": tokens for name, tokens in dataset_progress["tokens"].items()}, + **{f"progress/{name}/step_tokens": tokens for name, tokens in step_tokens.items()}, + **{f"progress/{name}/num_tokens": tokens for name, tokens in source_totals.items()}, + **{ + f"progress/{name}/epochs": tokens / train_dataset.sources[name].num_tokens + for name, tokens in source_totals.items() + }, "progress/num_samples": progress.total_samples, "progress/num_tokens": progress.total_tokens, "perf/throughput": throughput, diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index 4ae2bbb36e3..4a78fbe5132 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -4,7 +4,7 @@ import numpy as np from prime_rl.configs.pretrain import MixturePhaseConfig, TokenSourceConfig -from prime_rl.trainer.pretrain.data import TokenDataset, TokenShards, collate_rows +from prime_rl.trainer.pretrain.data import TokenDataset, TokenShards, collate_rows, source_tokens def write_shard(path: Path, documents: list[list[int]]) -> None: @@ -55,6 +55,17 @@ def dataset(): assert [window["tokens"].tolist() for window in islice(resumed, 5)] == full[3:8] +def test_source_tokens_count_every_window_of_the_mixture(tmp_path): + for name in "ab": + write_shard(tmp_path / name / "00.ds", [list(range(100))]) + sources = {name: TokenShards(TokenSourceConfig(path=tmp_path / name)) for name in "ab"} + dataset = TokenDataset(sources, [MixturePhaseConfig(weights={"a": 0.3, "b": 0.7})], seq_len=4, batch_size=8, max_steps=4) + list(islice(dataset, 32)) + # The training loop's global per-source counts equal the tokens the dataset itself read. + assert source_tokens(dataset, 0, 32) == dataset.tokens + assert sum(source_tokens(dataset, 8, 16).values()) == 8 * 4 + + def test_collate_rows_restarts_positions_at_documents(): rows = [ {"tokens": np.arange(5), "seq_lens": np.array([3, 1])}, From c83379844c5921fb479d4a5987919de9c9f2f643 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 09:54:18 +0200 Subject: [PATCH 055/132] docs(pretrain): CompleteP with QK-norm; the MoE example uses both The MoE QK-norm sweep (pretrain-experiments moe/2026-10-14-moe-completep-qknorm): without QK-norm, CompleteP's attention collapses one or two LR steps above its optimum; with QK-norm, CompleteP and SP tie and CompleteP's optimum (2^-7) stays fixed from d=512 to d=1024. The example turns on QK-norm, uses lr 2^-7 and CompleteP with its own shape as the base. --- docs/pretraining.md | 2 +- examples/pretrain/glm4_moe.toml | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index adaee8bc10e..bc099a2b9f4 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -59,7 +59,7 @@ With `optim.type = "muon"`, the hidden matrices and `o_proj`/`down_proj` are tra In MoE models, the expert matrices follow the hidden and `down_proj` rules per expert. The router maps the width to a fixed number of experts, so under Muon its learning rate falls as 1/m_N (with AdamW it follows the hidden rule). -The per-role groups need AdamW or Muon and `model.offload = "none"`. The standard parametrization is CompleteP with the base shape equal to the model shape. Keep `head_dim` and the ratios `intermediate_size / hidden_size` (and `moe_intermediate_size / hidden_size`) and `num_attention_heads / num_key_value_heads` fixed across shapes. +Use QK-norm (`model.arch.use_qk_norm = true`) with CompleteP. Without it, CompleteP's attention collapses one or two LR steps above the optimum (the norm gains get no weight decay), and the cliff moves down with width. With QK-norm, CompleteP and the standard parametrization reach the same loss, and CompleteP's optimal LR stays fixed across widths (MoE QK-norm sweep). The per-role groups need AdamW or Muon and `model.offload = "none"`. The standard parametrization is CompleteP with the base shape equal to the model shape. Keep `head_dim` and the ratios `intermediate_size / hidden_size` (and `moe_intermediate_size / hidden_size`) and `num_attention_heads / num_key_value_heads` fixed across shapes. ## Learning-rate schedule diff --git a/examples/pretrain/glm4_moe.toml b/examples/pretrain/glm4_moe.toml index 45f4b8af198..7bcd9c5eba9 100644 --- a/examples/pretrain/glm4_moe.toml +++ b/examples/pretrain/glm4_moe.toml @@ -1,6 +1,7 @@ # GLM-4.5-style MoE (64 routed experts, 6 active, 2 shared) with d = 512 and 8 layers, trained from scratch on # the Dolma 3 mixture tokenized with the DeepSeek-V4 tokenizer. -# 4096 steps x 256K tokens = 1.1B tokens. The learning rate is the optimum of the MoE base-shape sweep. +# 4096 steps x 256K tokens = 1.1B tokens. QK-norm and the learning rate follow the MoE QK-norm sweep (optimum 2^-7 +# with QK-norm at d = 512). [parametrization] with this base shape keeps the LR when the width grows. # # uv run pretrain @ examples/pretrain/glm4_moe.toml @@ -26,7 +27,7 @@ first_k_dense_replace = 1 norm_topk_prob = true routed_scaling_factor = 2.5 attention_bias = false -use_qk_norm = false +use_qk_norm = true partial_rotary_factor = 1.0 rope_parameters = { rope_type = "default", rope_theta = 500000.0, partial_rotary_factor = 1.0 } tie_word_embeddings = false @@ -53,9 +54,13 @@ weights = { web = 0.60, pdfs = 0.20, code = 0.12, math = 0.05, reference = 0.03 [optim] type = "muon" -lr = 0.00390625 +lr = 0.0078125 max_norm = 1.0 +[parametrization] +base_hidden_size = 512 +base_num_layers = 8 + [val] interval = 512 From 1c203939e20f4df02e9bc97c32544d07427ecd6c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 10:01:18 +0200 Subject: [PATCH 056/132] test(pretrain): keep test_data.py as in #4070 --- tests/unit/train/pretrain/test_data.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index 52db6180bd1..4a78fbe5132 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -59,9 +59,7 @@ def test_source_tokens_count_every_window_of_the_mixture(tmp_path): for name in "ab": write_shard(tmp_path / name / "00.ds", [list(range(100))]) sources = {name: TokenShards(TokenSourceConfig(path=tmp_path / name)) for name in "ab"} - dataset = TokenDataset( - sources, [MixturePhaseConfig(weights={"a": 0.3, "b": 0.7})], seq_len=4, batch_size=8, max_steps=4 - ) + dataset = TokenDataset(sources, [MixturePhaseConfig(weights={"a": 0.3, "b": 0.7})], seq_len=4, batch_size=8, max_steps=4) list(islice(dataset, 32)) # The training loop's global per-source counts equal the tokens the dataset itself read. assert source_tokens(dataset, 0, 32) == dataset.tokens From 23ae6cbbb234959dae30a6b3356c03b473887824 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 10:36:17 +0200 Subject: [PATCH 057/132] style(pretrain): ruff format test_data.py (ruff 0.15.14, as in CI) --- tests/unit/train/pretrain/test_data.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index 4a78fbe5132..52db6180bd1 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -59,7 +59,9 @@ def test_source_tokens_count_every_window_of_the_mixture(tmp_path): for name in "ab": write_shard(tmp_path / name / "00.ds", [list(range(100))]) sources = {name: TokenShards(TokenSourceConfig(path=tmp_path / name)) for name in "ab"} - dataset = TokenDataset(sources, [MixturePhaseConfig(weights={"a": 0.3, "b": 0.7})], seq_len=4, batch_size=8, max_steps=4) + dataset = TokenDataset( + sources, [MixturePhaseConfig(weights={"a": 0.3, "b": 0.7})], seq_len=4, batch_size=8, max_steps=4 + ) list(islice(dataset, 32)) # The training loop's global per-source counts equal the tokens the dataset itself read. assert source_tokens(dataset, 0, 32) == dataset.tokens From 4bb2a3d66e7a4c92b228243015256f4779e0aa03 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 10:45:01 +0200 Subject: [PATCH 058/132] feat(pretrain): expert-load max/mean and top-expert share per MoE layer --- src/prime_rl/trainer/pretrain/metrics.py | 15 +++++++++++++-- src/prime_rl/trainer/pretrain/step.py | 4 ++-- .../unit/train/pretrain/test_pretrain_metrics.py | 5 ++++- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/metrics.py b/src/prime_rl/trainer/pretrain/metrics.py index 90e42fe7911..813bf29e32b 100644 --- a/src/prime_rl/trainer/pretrain/metrics.py +++ b/src/prime_rl/trainer/pretrain/metrics.py @@ -8,6 +8,8 @@ from torch.distributed.tensor import DTensor, Partial, Replicate from prime_rl.trainer.model import iter_moe_blocks +from prime_rl.trainer.models.layers.moe import MoE +from prime_rl.utils.vlm import get_language_model def reduce_stats(sums: dict[str, Tensor], maxes: dict[str, Tensor], group: dist.ProcessGroup) -> dict[str, Tensor]: @@ -39,11 +41,17 @@ def take_routing_counts(model: nn.Module) -> tuple[Tensor, Tensor]: return counts, confidence -def expert_load_metrics(counts: Tensor) -> dict[str, float]: +def moe_layer_indices(model: nn.Module) -> list[int]: + """The decoder-layer index of every MoE layer, in the order of `iter_moe_blocks`.""" + return [index for index, layer in enumerate(get_language_model(model).layers) if isinstance(getattr(layer, "mlp", None), MoE)] + + +def expert_load_metrics(counts: Tensor, layer_indices: list[int]) -> dict[str, float]: """Balance of the `[num_moe_layers, num_experts]` expert counts of a whole step, per layer, summarized as the mean and the worst layer. `max_mean` / `min_mean`: busiest / least-loaded expert over the mean load; `cv`: std / mean; `cold_frac`: experts under 0.1x the mean load; `entropy`: entropy of the experts' shares of the - assignments over log(num_experts) (1 = uniform).""" + assignments over log(num_experts) (1 = uniform). Also per decoder layer (``layer_indices``): `max_mean` and + `top_share`, the busiest expert's share of the layer's assignments.""" mean = counts.mean(dim=1) share = counts / counts.sum(dim=1, keepdim=True) per_layer = { @@ -57,6 +65,9 @@ def expert_load_metrics(counts: Tensor) -> dict[str, float]: for name, (values, worst) in per_layer.items(): metrics[f"expert_load/{name}/mean"] = values.mean().item() metrics[f"expert_load/{name}/{worst}"] = (values.amax() if worst == "max" else values.amin()).item() + for layer, max_mean, top_share in zip(layer_indices, per_layer["max_mean"][0], share.amax(dim=1)): + metrics[f"expert_load/layer_{layer}/max_mean"] = max_mean.item() + metrics[f"expert_load/layer_{layer}/top_share"] = top_share.item() return metrics diff --git a/src/prime_rl/trainer/pretrain/step.py b/src/prime_rl/trainer/pretrain/step.py index 60bd169ffd0..0d18a33547b 100644 --- a/src/prime_rl/trainer/pretrain/step.py +++ b/src/prime_rl/trainer/pretrain/step.py @@ -11,7 +11,7 @@ from prime_rl.trainer.model import is_tt_moe_model, iter_moe_blocks from prime_rl.trainer.parallel_dims import ParallelDims -from prime_rl.trainer.pretrain.metrics import expert_load_metrics, reduce_stats, take_routing_counts +from prime_rl.trainer.pretrain.metrics import expert_load_metrics, moe_layer_indices, reduce_stats, take_routing_counts from prime_rl.trainer.utils import begin_backward, finish_backward, prepare_gradient_offload, scale_gradients_ @@ -128,7 +128,7 @@ def run_micro_batches( if is_moe_model: counts = stats["expert_counts"] routed_tokens = counts.sum(dim=1) / next(iter_moe_blocks(model)).router.top_k - metrics |= expert_load_metrics(counts) | { + metrics |= expert_load_metrics(counts, moe_layer_indices(model)) | { "expert_load/micro_max_mean/mean": stats["micro_max_mean"].item() / (grad_accum_steps * num_ranks), "expert_load/micro_max_mean/max": stats["micro_max_mean_max"].item(), "router/confidence": (stats["routing_confidence"] / routed_tokens).mean().item(), diff --git a/tests/unit/train/pretrain/test_pretrain_metrics.py b/tests/unit/train/pretrain/test_pretrain_metrics.py index 55e8de6b3ba..b88bb30a35e 100644 --- a/tests/unit/train/pretrain/test_pretrain_metrics.py +++ b/tests/unit/train/pretrain/test_pretrain_metrics.py @@ -6,12 +6,15 @@ def test_expert_load_metrics_per_layer_summaries(): counts = torch.tensor([[1.0, 1.0, 1.0, 1.0], [4.0, 0.0, 0.0, 0.0]], dtype=torch.float64) - metrics = expert_load_metrics(counts) + metrics = expert_load_metrics(counts, layer_indices=[1, 2]) assert metrics["expert_load/max_mean/mean"] == 2.5 # layers: 1 and 4 assert metrics["expert_load/max_mean/max"] == 4.0 assert metrics["expert_load/min_mean/min"] == 0.0 assert metrics["expert_load/cold_frac/max"] == 0.75 assert metrics["expert_load/entropy/mean"] == 0.5 # uniform layer 1, one-hot layer 0 + assert metrics["expert_load/layer_2/max_mean"] == 4.0 + assert metrics["expert_load/layer_1/top_share"] == 0.25 + assert metrics["expert_load/layer_2/top_share"] == 1.0 def test_grad_norms_by_group_add_up_to_the_total_norm(): From cac6bf4bf7c2f68c9c1ffffd11e2fb08f6626b66 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 10:45:48 +0200 Subject: [PATCH 059/132] style(pretrain): ruff format --- src/prime_rl/trainer/pretrain/metrics.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/pretrain/metrics.py b/src/prime_rl/trainer/pretrain/metrics.py index 813bf29e32b..90fdebd82e0 100644 --- a/src/prime_rl/trainer/pretrain/metrics.py +++ b/src/prime_rl/trainer/pretrain/metrics.py @@ -43,7 +43,11 @@ def take_routing_counts(model: nn.Module) -> tuple[Tensor, Tensor]: def moe_layer_indices(model: nn.Module) -> list[int]: """The decoder-layer index of every MoE layer, in the order of `iter_moe_blocks`.""" - return [index for index, layer in enumerate(get_language_model(model).layers) if isinstance(getattr(layer, "mlp", None), MoE)] + return [ + index + for index, layer in enumerate(get_language_model(model).layers) + if isinstance(getattr(layer, "mlp", None), MoE) + ] def expert_load_metrics(counts: Tensor, layer_indices: list[int]) -> dict[str, float]: From 8dc1866fc8af6e4b848727f3fc3309342f255aa4 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 11:50:41 +0200 Subject: [PATCH 060/132] refactor(pretrain): decide once whether the resume restores the data position --- src/prime_rl/trainer/pretrain/train.py | 6 +++--- tests/unit/train/pretrain/test_schedules.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index bc88d4983c6..94bbff0a7ba 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -171,15 +171,15 @@ def train(config: PretrainConfig): logger.info(f"Val source {name!r}: {val_windows[name] * config.data.seq_len:,} tokens") progress = Progress() + restore_dataloader = checkpoint_step is not None and not (config.ckpt or CheckpointConfig()).skip_dataloader if checkpoint_step is not None: - skip = config.ckpt or CheckpointConfig() ckpt_manager.load( checkpoint_step, model, [optimizer], None, progress, - dataloader=dataloader if not skip.skip_dataloader else None, + dataloader=dataloader if restore_dataloader else None, path=config.resume.dir / "trainer" if config.resume.dir is not None else None, ) progress.step += 1 @@ -189,7 +189,7 @@ def train(config: PretrainConfig): # from the window sequence itself, which every data rank walks, so they are global without communication. next_window = 0 source_totals = dict.fromkeys(train_dataset.names, 0) - if checkpoint_step is not None and not (config.ckpt or CheckpointConfig()).skip_dataloader: + if restore_dataloader: next_window = checkpoint_step * config.data.batch_size # The loaded state counts the windows before its own position, which is near next_window. state = next(iter(get_dataset_state(dataloader).values())) diff --git a/tests/unit/train/pretrain/test_schedules.py b/tests/unit/train/pretrain/test_schedules.py index dc867d73d11..e02f526fa74 100644 --- a/tests/unit/train/pretrain/test_schedules.py +++ b/tests/unit/train/pretrain/test_schedules.py @@ -39,7 +39,7 @@ def test_schedule_phases_must_fit_max_steps(): ) with pytest.raises(ValidationError, match="exceed max_steps"): PretrainConfig(**{**config.model_dump(), "schedule": {"warmup": 50, "decay": 51}}) - # Fractions summing to 1 fit for any max_steps (rounding both up used to exceed it). + # Fractions summing to 1 fit for any max_steps. PretrainConfig(**{**config.model_dump(), "max_steps": 21095}) From a5d6e2a42a3291d9c9709081fc630898b4fe11b3 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 11:54:59 +0200 Subject: [PATCH 061/132] refactor(pretrain): own micro-batch step; SFT trainer files as on main After the monitoring PR the pretrain loop runs its own step, so the shared trainer/step.py only served SFT. The SFT trainer, its data module and its test are back to their base versions; pretrain's step lives in trainer/pretrain/step.py without the SFT-only options and zero-token guards. --- src/prime_rl/trainer/data.py | 73 ----------- src/prime_rl/trainer/pretrain/data.py | 2 +- src/prime_rl/trainer/{ => pretrain}/step.py | 28 ++-- src/prime_rl/trainer/pretrain/train.py | 4 +- src/prime_rl/trainer/sft/data.py | 70 +++++++++- src/prime_rl/trainer/sft/train.py | 135 +++++++++++++++----- tests/unit/train/sft/test_dataloader.py | 3 +- 7 files changed, 186 insertions(+), 129 deletions(-) delete mode 100644 src/prime_rl/trainer/data.py rename src/prime_rl/trainer/{ => pretrain}/step.py (82%) diff --git a/src/prime_rl/trainer/data.py b/src/prime_rl/trainer/data.py deleted file mode 100644 index dd0e725b084..00000000000 --- a/src/prime_rl/trainer/data.py +++ /dev/null @@ -1,73 +0,0 @@ -from collections import defaultdict - -from torch.distributed.checkpoint.stateful import Stateful -from torch.utils.data import IterableDataset, get_worker_info -from torchdata.stateful_dataloader import StatefulDataLoader - -from prime_rl.trainer.world import get_world - - -class StatefulIterableDataset(Stateful, IterableDataset): - """SFT dataset are iterable (infinite) and stateful (can be checkpointed).""" - - def __init__(self, non_dp_size: int = 1): - self.step, self.epoch = 0, 0 - self.num_samples = defaultdict(int) - self.num_tokens = defaultdict(int) - self.fast_forward = False - self.non_dp_size = non_dp_size - self._setup_world_info() - - def state_dict(self) -> dict: - return {"step": self.step, "epoch": self.epoch} - - def load_state_dict(self, state_dict: dict): - assert "step" in state_dict and "epoch" in state_dict - self.fast_forward = True - self.step = state_dict["step"] - self.epoch = state_dict["epoch"] - - def _setup_world_info(self): - worker_info = get_worker_info() - if worker_info is not None: - worker_id = worker_info.id - num_workers = worker_info.num_workers - else: - worker_id, num_workers = 0, 1 - world = get_world() - assert world.world_size % self.non_dp_size == 0, "world_size must be divisible by non_dp_size" - self.data_rank = world.rank // self.non_dp_size * num_workers + worker_id - self.data_world_size = world.world_size // self.non_dp_size * num_workers - - -def get_dataset_state(dataloader: StatefulDataLoader) -> dict: - """Dataset position per worker for the startup log, parsed from ``StatefulDataLoader.state_dict()``. - - The loader is the only source that is correct in every case: after a resume's - ``load_state_dict`` the restored position exists solely in the loader's stashed - state (it reaches the dataset copies inside workers when the iterator forks them; - the main-process dataset object stays at position zero). The keys are torchdata's - private worker-snapshot layout.""" - snapshots = dataloader.state_dict()["_snapshot"]["_worker_snapshots"] - return {wid: snap["dataset_state"]["dataset"] for wid, snap in sorted(snapshots.items())} - - -def get_dataset_progress(dataloader: StatefulDataLoader) -> dict: - """Dataset position and aggregate counters from dataloader workers.""" - snapshot = dataloader.state_dict()["_snapshot"] - worker_snapshots = snapshot["_worker_snapshots"] - positions = [worker_snapshot["dataset_state"]["dataset"] for worker_snapshot in worker_snapshots.values()] - furthest = max(positions, key=lambda position: position["step"]) - num_samples = defaultdict(int) - num_tokens = defaultdict(int) - for worker_snapshot in worker_snapshots.values(): - progress = worker_snapshot["dataset_state"].get("progress", {}) - for name, count in progress.get("num_samples", {}).items(): - num_samples[name] += count - for name, count in progress.get("num_tokens", {}).items(): - num_tokens[name] += count - return { - **furthest, - "num_samples": dict(num_samples), - "num_tokens": dict(num_tokens), - } diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index f0e4d069d85..b5f9a660437 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -6,7 +6,7 @@ from torchdata.stateful_dataloader import StatefulDataLoader from prime_rl.configs.pretrain import MixturePhaseConfig, PretrainDataConfig, TokenSourceConfig -from prime_rl.trainer.data import StatefulIterableDataset +from prime_rl.trainer.sft.data import StatefulIterableDataset from prime_rl.utils.logger import get_logger diff --git a/src/prime_rl/trainer/step.py b/src/prime_rl/trainer/pretrain/step.py similarity index 82% rename from src/prime_rl/trainer/step.py rename to src/prime_rl/trainer/pretrain/step.py index 7ea93626ab6..469a4c6c271 100644 --- a/src/prime_rl/trainer/step.py +++ b/src/prime_rl/trainer/pretrain/step.py @@ -1,5 +1,4 @@ from collections.abc import Callable, Iterator -from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass import torch @@ -28,7 +27,7 @@ class StepResult: ``routing_confidence/mean``). Empty for dense models.""" tokens_per_expert: Tensor | None - """Global ``[num_moe_layers, num_experts]`` token counts of the step. None for dense models.""" + """This rank's ``[num_moe_layers, num_experts]`` token counts of the step. None for dense models.""" def run_micro_batches( @@ -38,9 +37,6 @@ def run_micro_batches( grad_accum_steps: int, parallel_dims: ParallelDims, gradient_manager=None, - *, - overlap_optimizer: bool = True, - record_function: Callable[[str], AbstractContextManager] = nullcontext, ) -> StepResult: """Forward and backward passes over the ``grad_accum_steps`` micro-batches of one step, leaving the gradients of the summed token loss divided by the global token count (the mean loss per token, as on one @@ -62,12 +58,8 @@ def run_micro_batches( ) dist.all_reduce(global_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) global_token_count = global_token_count.item() // parallel_dims.cp - grad_scale = ( - parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count - if global_token_count > 0 - else 1.0 - ) - prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=overlap_optimizer) + grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count + prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=True) loss_sum = torch.tensor(0.0, device="cuda") nan_count = torch.tensor(0, device="cuda") @@ -86,10 +78,9 @@ def run_micro_batches( else: loss_sum += micro_loss_sum.detach() scaled_loss = micro_loss_sum / grad_accum_steps - with record_function("backward"): - begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) - scaled_loss.backward() - finish_backward(gradient_manager) + begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) + scaled_loss.backward() + finish_backward(gradient_manager) if is_moe_model: micro_moe_stats, micro_tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) @@ -102,10 +93,7 @@ def run_micro_batches( if gradient_manager is None: dist.all_reduce(local_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) global_token_count = local_token_count.item() - if global_token_count > 0: - scale_gradients_( - None, model, parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count - ) + scale_gradients_(None, model, parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count) dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) dist.all_reduce(nan_count, op=dist.ReduceOp.SUM) @@ -113,7 +101,7 @@ def run_micro_batches( num_tokens = torch.tensor(local_num_tokens, dtype=torch.int64, device="cuda") dist.all_reduce(num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) return StepResult( - loss=(loss_sum / global_token_count).item() if global_token_count > 0 else 0.0, + loss=(loss_sum / global_token_count).item(), nan_count=int(nan_count.item()), num_tokens=int(num_tokens.item()), moe_stats={name: value.item() for name, value in moe_stats.items()} if is_moe_model else {}, diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 363217cf236..9522ba39d93 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -41,8 +41,8 @@ setup_val_dataset, ) from prime_rl.trainer.pretrain.schedules import LRSchedule -from prime_rl.trainer.data import get_dataset_state -from prime_rl.trainer.step import run_micro_batches +from prime_rl.trainer.pretrain.step import run_micro_batches +from prime_rl.trainer.sft.data import get_dataset_state from prime_rl.trainer.utils import ( GarbageCollection, clip_grad_norm_, diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index d8bb479faa8..6fd0a5dc7da 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -13,11 +13,13 @@ from renderers import AutoRendererConfig, RendererConfig, merge_chat_template_kwargs from renderers.base import MultiModalData, PlaceholderRange, Renderer, build_training_sample, create_renderer from torch import Tensor +from torch.distributed.checkpoint.stateful import Stateful +from torch.utils.data import IterableDataset, get_worker_info from torchdata.stateful_dataloader import StatefulDataLoader from transformers.tokenization_utils import PreTrainedTokenizer from prime_rl.configs.sft import DataConfig, LossMaskConfig, SFTColumnsConfig, SFTDataConfig -from prime_rl.trainer.data import StatefulIterableDataset +from prime_rl.trainer.world import get_world from prime_rl.utils.chat_template import deserialize_tool_calls, normalize_messages from prime_rl.utils.logger import get_logger from prime_rl.utils.utils import format_time @@ -43,6 +45,39 @@ class Batch(TypedDict): mm_token_type_ids: Int[Tensor, "batch seq"] | None +class StatefulIterableDataset(Stateful, IterableDataset): + """SFT dataset are iterable (infinite) and stateful (can be checkpointed).""" + + def __init__(self, non_dp_size: int = 1): + self.step, self.epoch = 0, 0 + self.num_samples = defaultdict(int) + self.num_tokens = defaultdict(int) + self.fast_forward = False + self.non_dp_size = non_dp_size + self._setup_world_info() + + def state_dict(self) -> dict: + return {"step": self.step, "epoch": self.epoch} + + def load_state_dict(self, state_dict: dict): + assert "step" in state_dict and "epoch" in state_dict + self.fast_forward = True + self.step = state_dict["step"] + self.epoch = state_dict["epoch"] + + def _setup_world_info(self): + worker_info = get_worker_info() + if worker_info is not None: + worker_id = worker_info.id + num_workers = worker_info.num_workers + else: + worker_id, num_workers = 0, 1 + world = get_world() + assert world.world_size % self.non_dp_size == 0, "world_size must be divisible by non_dp_size" + self.data_rank = world.rank // self.non_dp_size * num_workers + worker_id + self.data_world_size = world.world_size // self.non_dp_size * num_workers + + class FakeDataset(StatefulIterableDataset): """A dataset of fake tokens""" @@ -734,3 +769,36 @@ def setup_dataloader(dataset: StatefulIterableDataset, config: DataConfig, cp: i num_workers=config.num_workers, pin_memory=True, ) + + +def get_dataset_state(dataloader: StatefulDataLoader) -> dict: + """Dataset position per worker for the startup log, parsed from ``StatefulDataLoader.state_dict()``. + + The loader is the only source that is correct in every case: after a resume's + ``load_state_dict`` the restored position exists solely in the loader's stashed + state (it reaches the dataset copies inside workers when the iterator forks them; + the main-process dataset object stays at position zero). The keys are torchdata's + private worker-snapshot layout.""" + snapshots = dataloader.state_dict()["_snapshot"]["_worker_snapshots"] + return {wid: snap["dataset_state"]["dataset"] for wid, snap in sorted(snapshots.items())} + + +def get_dataset_progress(dataloader: StatefulDataLoader) -> dict: + """Dataset position and aggregate counters from dataloader workers.""" + snapshot = dataloader.state_dict()["_snapshot"] + worker_snapshots = snapshot["_worker_snapshots"] + positions = [worker_snapshot["dataset_state"]["dataset"] for worker_snapshot in worker_snapshots.values()] + furthest = max(positions, key=lambda position: position["step"]) + num_samples = defaultdict(int) + num_tokens = defaultdict(int) + for worker_snapshot in worker_snapshots.values(): + progress = worker_snapshot["dataset_state"].get("progress", {}) + for name, count in progress.get("num_samples", {}).items(): + num_samples[name] += count + for name, count in progress.get("num_tokens", {}).items(): + num_tokens[name] += count + return { + **furthest, + "num_samples": dict(num_samples), + "num_tokens": dict(num_tokens), + } diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 234074d354a..98727b861b9 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -28,6 +28,7 @@ forward, get_full_offload_dtype_policy, get_expert_load_stats, + get_global_moe_stats, get_load_balance_stats, is_tt_moe_model, reshard_module, @@ -37,19 +38,23 @@ ) from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter -from prime_rl.trainer.data import get_dataset_progress, get_dataset_state from prime_rl.trainer.sft.data import ( + get_dataset_progress, + get_dataset_state, load_sft_dataset, setup_dataloader, setup_dataset, ) -from prime_rl.trainer.step import run_micro_batches from prime_rl.trainer.utils import ( GarbageCollection, MemoryProfiler, + begin_backward, clip_grad_norm_, + finish_backward, get_ckpt_disk_metrics, + prepare_gradient_offload, print_sample, + scale_gradients_, setup_full_cpu_optimizer_offload, setup_torch_distributed, ) @@ -230,8 +235,8 @@ def train(config: SFTConfig): cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() if cp_enabled else 0 cp_group = parallel_dims.world_mesh["cp"].get_group() if cp_enabled else None dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() + ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None cp_size = parallel_dims.cp - is_moe_model = is_tt_moe_model(model) def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: """Forward pass returning (loss_sum, token_count) over unmasked tokens.""" @@ -293,14 +298,6 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: del out return loss_sum, token_count - def compute_train_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: - if config.log.log_data: - print_sample( - micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer - ) - with maybe_record_function("forward"): - return compute_loss(micro_batch) - maybe_record_function = nullcontext def run_eval_loop(data_iter): @@ -400,7 +397,7 @@ def is_online_eval_step(step: int) -> bool: if config.trace_path: logger.info(f"Tracing to {config.trace_path}") prof = profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True).__enter__() - maybe_record_function = record_function + maybe_record_function = record_function # noqa: F841 – captured by run_forward_loop closure max_peak_memory = 0.0 while True: # Reset peak memory stats @@ -416,22 +413,88 @@ def is_online_eval_step(step: int) -> bool: step_start_time = time.perf_counter() forward_backward_start_time = time.perf_counter() + step_loss_sum = torch.tensor(0.0, device="cuda") + nan_loss_count = torch.tensor(0, device="cuda") + is_moe_model = is_tt_moe_model(model) + moe_stats = ( + { + "max_vio/mean": torch.tensor(0.0), + "max_vio/max": torch.tensor(0.0), + "routing_confidence/mean": torch.tensor(0.0), + } + if is_moe_model + else {} + ) run_validation_this_step = config.val is not None and ( (is_first_step and config.val.eval_on_start) or (not is_first_step and progress.step % config.val.interval == 0) ) - result = run_micro_batches( - model, - dataiter, - compute_train_loss, - grad_accum_steps, - parallel_dims, - gradient_manager, - overlap_optimizer=not run_validation_this_step, - record_function=maybe_record_function, - ) + if gradient_manager is None: + micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) + step_local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") + else: + micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] + local_token_count = sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches) + global_step_token_count = torch.tensor(local_token_count, dtype=torch.int64, device="cuda") + dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count_val = global_step_token_count.item() // cp_size + grad_scale = ( + parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val + if global_token_count_val > 0 + else 1.0 + ) + prepare_gradient_offload( + gradient_manager, + grad_scale, + overlap_optimizer=not run_validation_this_step, + ) + + step_tokens_per_expert = 0 + step_local_num_tokens = 0 + for micro_step, micro_batch in enumerate(micro_batches): + step_local_num_tokens += micro_batch["input_ids"].shape[1] + if config.log.log_data: + print_sample( + micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer + ) + + with maybe_record_function("forward"): + local_loss_sum, batch_token_count = compute_loss(micro_batch) + + if gradient_manager is None: + step_local_token_count += batch_token_count + + if torch.isnan(local_loss_sum.detach()): + nan_loss_count += 1 + logger.warning("Local loss is nan, excluding this micro step from backward") + scaled_loss = torch.nan_to_num(local_loss_sum, nan=0.0) / grad_accum_steps + else: + step_loss_sum += local_loss_sum.detach() + scaled_loss = local_loss_sum / grad_accum_steps + + with maybe_record_function("backward"): + begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) + scaled_loss.backward() + finish_backward(gradient_manager) + + if is_moe_model: + micro_moe_stats, tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) + for name, value in micro_moe_stats.items(): + moe_stats[f"{name}/mean"] += value / grad_accum_steps + if name == "max_vio": + moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) + step_tokens_per_expert += tokens_per_expert + forward_backward_time = time.perf_counter() - forward_backward_start_time - expert_load_stats = get_expert_load_stats(result.tokens_per_expert, dp_cp_group) if is_moe_model else {} + expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) if is_moe_model else {} + + if gradient_manager is None: + global_step_token_count = step_local_token_count.clone() + dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count_val = global_step_token_count.item() + if global_token_count_val > 0: + grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val + scale_gradients_(None, model, grad_scale) # Run validation after forward-backward (so torch.compile sees training graph first) but before # optimizer step (so eval_on_start evaluates untrained weights) @@ -442,6 +505,15 @@ def is_online_eval_step(step: int) -> bool: # unsharded parameters. Reshard so clipping sees every gradient. reshard_module(model) + # Compute the global mean loss for logging. + dist.all_reduce(step_loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) + dist.all_reduce(nan_loss_count, op=dist.ReduceOp.SUM) + if global_token_count_val > 0: + batch_loss = (step_loss_sum / global_token_count_val).item() + else: + batch_loss = 0.0 + nan_loss_count = nan_loss_count.item() + grad_norm: torch.Tensor | None = None if config.optim.max_norm is not None: logger.debug(f"Clipping gradients with max norm {config.optim.max_norm}") @@ -478,7 +550,10 @@ def is_online_eval_step(step: int) -> bool: if memory_profiler is not None: memory_profiler.step() - num_tokens = result.num_tokens + # The dp mesh excludes cp, whose ranks hold the same rows. + global_num_tokens = torch.tensor(step_local_num_tokens, dtype=torch.int64, device="cuda") + dist.all_reduce(global_num_tokens, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) + num_tokens = global_num_tokens.item() progress.total_tokens += num_tokens dataset_progress = get_dataset_progress(dataloader) progress.total_samples = dataset_progress["step"] @@ -493,13 +568,13 @@ def is_online_eval_step(step: int) -> bool: max_peak_memory = max(max_peak_memory, peak_memory) # Log step metrics - step_message = f"Step {progress.step} | {format_time(step_time):>7} | Loss {result.loss:.4f}" + step_message = f"Step {progress.step} | {format_time(step_time):>7} | Loss {batch_loss:.4f}" if grad_norm is not None: step_message += f" | Grad. Norm {grad_norm:.4f}" step_message += f" | LR {current_lr:.2e} | Throughput {throughput:.0f} tokens/s | MFU {mfu:.1f}% | Peak Mem. {peak_memory:.1f}/{max_memory:.1f} GiB ({peak_memory / max_memory * 100:.1f}%)" if is_moe_model: for name, label in (("max_vio/mean", "Max Vio"), ("routing_confidence/mean", "Routing Conf.")): - value = result.moe_stats[name] + value = moe_stats[name].item() if value > 0: step_message += f" | {label} {value:.4f}" logger.success(step_message) @@ -549,9 +624,9 @@ def is_online_eval_step(step: int) -> bool: asyncio.run(monitors.log(optim_metrics, step=progress.step)) loss_log_metrics = { - "loss/mean": result.loss, - "loss/perplexity": math.exp(min(result.loss, 20)), - "loss/nan_count": result.nan_count, + "loss/mean": batch_loss, + "loss/perplexity": math.exp(min(batch_loss, 20)), + "loss/nan_count": nan_loss_count, "step": progress.step, } # Log tensor stats @@ -572,7 +647,7 @@ def is_online_eval_step(step: int) -> bool: disk_metrics["step"] = progress.step asyncio.run(monitors.log(disk_metrics, step=progress.step)) - moe_log_metrics = result.moe_stats | expert_load_stats + moe_log_metrics = {name: value.item() for name, value in moe_stats.items()} | expert_load_stats if moe_log_metrics: asyncio.run(monitors.log({**moe_log_metrics, "step": progress.step}, step=progress.step)) diff --git a/tests/unit/train/sft/test_dataloader.py b/tests/unit/train/sft/test_dataloader.py index 9615cf05104..7f2c7d47839 100644 --- a/tests/unit/train/sft/test_dataloader.py +++ b/tests/unit/train/sft/test_dataloader.py @@ -6,8 +6,7 @@ from datasets import Dataset from prime_rl.configs.sft import FakeDataConfig, SFTDataConfig -from prime_rl.trainer.data import get_dataset_progress, get_dataset_state -from prime_rl.trainer.sft.data import FakeDataset, SFTDataset, setup_dataloader +from prime_rl.trainer.sft.data import FakeDataset, SFTDataset, get_dataset_progress, get_dataset_state, setup_dataloader from prime_rl.trainer.world import reset_world From 0d003841d13e4a3f64d698e54c0be64a3b082ab9 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 11:59:54 +0200 Subject: [PATCH 062/132] refactor(pretrain): drop the CPU-offload paths pretraining cannot take model.offload only allows 'none' (a Literal instead of a validator), so the pretrain loop and its step lose the gradient-offload branches and the shared setup_optimizer loses a guard only pretrain could reach. Also: normal_shard call back on one line, docs list the active parameter count. --- docs/pretraining.md | 2 +- .../src/prime_rl/configs/pretrain.py | 11 ++----- src/prime_rl/trainer/init.py | 7 +--- src/prime_rl/trainer/optim/__init__.py | 2 -- src/prime_rl/trainer/pretrain/step.py | 33 +++++-------------- src/prime_rl/trainer/pretrain/train.py | 20 +++-------- 6 files changed, 16 insertions(+), 59 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index bc099a2b9f4..327e4c374c0 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -22,7 +22,7 @@ num_experts_per_tok = 6 name = "tokenizers/deepseek-v4-flash-base" ``` -The launcher writes the config and the tokenizer to `/configs/model` and trains from there, so the trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all see an ordinary model directory. The trainer logs the non-embedding, embedding (input embedding and LM head) and total parameter counts at startup. +The launcher writes the config and the tokenizer to `/configs/model` and trains from there, so the trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all see an ordinary model directory. The trainer logs the active non-embedding (routed experts counted at `top_k / num_experts`), non-embedding, embedding (input embedding and LM head) and total parameter counts at startup, and as `params/*` at step 0. ## Initialization diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index d25d10dca81..998d0c78db6 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -24,7 +24,6 @@ ModelConfig, MoERuntimeConfig, MuonConfig, - OffloadMode, TokenizerConfig, ) from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource @@ -85,8 +84,8 @@ class PretrainModelConfig(ModelConfig): moe: PretrainMoEConfig = PretrainMoEConfig() - offload: OffloadMode = "none" - """CPU offloading. Pretraining's per-role parameter groups need ``none``.""" + offload: Literal["none"] = "none" + """CPU offloading is not supported: pretraining's per-role parameter groups need ``none``.""" @model_validator(mode="after") def validate_arch(self): @@ -392,12 +391,6 @@ def deepep_disables_grad_clipping(self): self.optim.max_norm = None return self - @model_validator(mode="after") - def validate_offload(self): - if self.model.offload != "none": - raise ValueError("Pretraining's per-role parameter groups need model.offload = 'none'.") - return self - @model_validator(mode="after") def validate_monitors(self): if self.monitors.prime is not None: diff --git a/src/prime_rl/trainer/init.py b/src/prime_rl/trainer/init.py index dbd7d83689b..0f0fc6a5e87 100644 --- a/src/prime_rl/trainer/init.py +++ b/src/prime_rl/trainer/init.py @@ -99,12 +99,7 @@ def init_weights(model: nn.Module, config: InitConfig, m_n: float = 1.0, m_l: fl else: local_shape, global_offset = tuple(param.shape), (0,) * param.ndim values = normal_shard( - name, - tuple(param.shape), - tuple(local_shape), - tuple(global_offset), - config.seed, - local.device, + name, tuple(param.shape), tuple(local_shape), tuple(global_offset), config.seed, local.device ) local.copy_(values.mul_(role_std(role, config.std, m_n, m_l))) # MoE selection biases are buffers: they start at 0, so routing starts from the router scores alone. diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index 3d94d596ee2..99a5e10c22a 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/src/prime_rl/trainer/optim/__init__.py @@ -60,8 +60,6 @@ def setup_optimizer( ``algorithm`` key). It is not supported with CPU offloading.""" cpu_offload = offload == "optimizer" full_offload = offload == "full" - if param_groups is not None and offload != "none": - raise ValueError("Custom parameter groups are not supported with CPU offloading") if full_offload and config.type not in ("adamw", "sign_sgd"): raise ValueError("Full optimizer offload only supports AdamW and SignSGD") if full_offload and config.max_norm is not None: diff --git a/src/prime_rl/trainer/pretrain/step.py b/src/prime_rl/trainer/pretrain/step.py index 1eb8a5e817d..67d9555e5ff 100644 --- a/src/prime_rl/trainer/pretrain/step.py +++ b/src/prime_rl/trainer/pretrain/step.py @@ -7,7 +7,7 @@ from prime_rl.trainer.model import get_global_moe_stats, is_tt_moe_model from prime_rl.trainer.parallel_dims import ParallelDims -from prime_rl.trainer.utils import begin_backward, finish_backward, prepare_gradient_offload, scale_gradients_ +from prime_rl.trainer.utils import scale_gradients_ from prime_rl.utils.logger import get_logger @@ -36,7 +36,6 @@ def run_micro_batches( compute_loss: Callable[[dict], tuple[Tensor, Tensor, Tensor]], grad_accum_steps: int, parallel_dims: ParallelDims, - gradient_manager=None, ) -> StepResult: """Forward and backward passes over the ``grad_accum_steps`` micro-batches of one step, leaving the gradients of the summed token loss divided by the global token count (the mean loss per token, as on one @@ -47,30 +46,17 @@ def run_micro_batches( ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None is_moe_model = is_tt_moe_model(model) - if gradient_manager is None: - micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) - local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") - else: - # The offloaded gradients are scaled as they stream out, so the token count is needed up front. - micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] - global_token_count = torch.tensor( - sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches), device="cuda" - ) - dist.all_reduce(global_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) - global_token_count = global_token_count.item() // parallel_dims.cp - grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count - prepare_gradient_offload(gradient_manager, grad_scale, overlap_optimizer=True) - loss_sum = torch.tensor(0.0, device="cuda") + token_count = torch.tensor(0, dtype=torch.int64, device="cuda") nan_count = torch.tensor(0, device="cuda") moe_stats = {name: torch.tensor(0.0) for name in ("max_vio/mean", "max_vio/max", "routing_confidence/mean")} tokens_per_expert = None local_num_tokens = 0 - for micro_step, micro_batch in enumerate(micro_batches): + for _ in range(grad_accum_steps): + micro_batch = next(dataiter) local_num_tokens += micro_batch["input_ids"].shape[1] micro_loss_sum, micro_token_count, aux_loss_sum = compute_loss(micro_batch) - if gradient_manager is None: - local_token_count += micro_token_count + token_count += micro_token_count if torch.isnan(micro_loss_sum.detach()): nan_count += 1 logger.warning("Local loss is nan, excluding this micro step from backward") @@ -78,9 +64,7 @@ def run_micro_batches( else: loss_sum += micro_loss_sum.detach() scaled_loss = (micro_loss_sum + aux_loss_sum) / grad_accum_steps - begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) scaled_loss.backward() - finish_backward(gradient_manager) if is_moe_model: micro_moe_stats, micro_tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) @@ -90,10 +74,9 @@ def run_micro_batches( moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) tokens_per_expert = micro_tokens_per_expert + (0 if tokens_per_expert is None else tokens_per_expert) - if gradient_manager is None: - dist.all_reduce(local_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) - global_token_count = local_token_count.item() - scale_gradients_(None, model, parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count) + dist.all_reduce(token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) + global_token_count = token_count.item() + scale_gradients_(None, model, parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count) dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) dist.all_reduce(nan_count, op=dist.ReduceOp.SUM) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 19965e52f1f..232f7d0edf0 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -19,7 +19,6 @@ from prime_rl.trainer.model import ( forward, get_expert_load_stats, - get_full_offload_dtype_policy, get_load_balance_stats, is_tt_moe_model, iter_moe_blocks, @@ -50,7 +49,6 @@ GarbageCollection, clip_grad_norm_, get_ckpt_disk_metrics, - setup_full_cpu_optimizer_offload, setup_torch_distributed, ) from prime_rl.trainer.world import get_world @@ -105,11 +103,7 @@ def train(config: PretrainConfig): heart = Heartbeat(config.heartbeat) if config.heartbeat is not None and world.rank == 0 else None - setup_torch_distributed( - timeout=timedelta(seconds=config.dist_timeout_seconds), enable_gloo=config.model.offload == "full" - ) - if config.model.offload == "full": - setup_full_cpu_optimizer_offload() + setup_torch_distributed(timeout=timedelta(seconds=config.dist_timeout_seconds)) torch.set_float32_matmul_precision(config.matmul_precision) resolve_ep(config.model) @@ -150,15 +144,11 @@ def train(config: PretrainConfig): tokenizer = setup_tokenizer(config.tokenizer) - optimizer, gradient_manager = setup_optimizer( + optimizer, _ = setup_optimizer( config.optim, list(model.named_parameters()), parallel_dims, - offload=config.model.offload, model=model, - full_offload_dtype_policy=( - get_full_offload_dtype_policy(model, config.model) if config.model.offload == "full" else None - ), param_groups=param_groups(model, config.optim, m_n, m_l, parallel_dims.ep_enabled), ) @@ -302,7 +292,7 @@ def run_evals(step: int) -> None: step_start_time = time.perf_counter() lr_multiplier = lr_schedule.apply(progress.step - 1) - result = run_micro_batches(model, dataiter, compute_loss, grad_accum_steps, parallel_dims, gradient_manager) + result = run_micro_batches(model, dataiter, compute_loss, grad_accum_steps, parallel_dims) forward_backward_time = time.perf_counter() - step_start_time expert_load_stats = {} if is_moe_model: @@ -318,7 +308,7 @@ def run_evals(step: int) -> None: if biases: expert_load_stats["expert_load/selection_bias_absmax"] = torch.stack(biases).abs().max().item() - grad_norm = clip_grad_norm_(gradient_manager, model, config.optim.max_norm) if config.optim.max_norm else None + grad_norm = clip_grad_norm_(None, model, config.optim.max_norm) if config.optim.max_norm else None optimizer.step() optimizer.zero_grad() current_lr = optimizer.param_groups[0]["lr"] @@ -400,8 +390,6 @@ def run_evals(step: int) -> None: logger.info(f"Saving final checkpoint at step {progress.step}") ckpt_manager.save(progress.step, model, [optimizer], None, progress, dataloader=dataloader) ckpt_manager.maybe_clean() - if gradient_manager is not None: - gradient_manager.close() logger.success("Pretrain trainer finished") asyncio.run(monitors.finalize()) From 4da039556317a5e306459de3769e77bcd22f8828 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 12:00:04 +0200 Subject: [PATCH 063/132] docs(pretrain): the MoE example relies on the only offload mode --- examples/pretrain/glm4_moe.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/pretrain/glm4_moe.toml b/examples/pretrain/glm4_moe.toml index 7bcd9c5eba9..f99c6849455 100644 --- a/examples/pretrain/glm4_moe.toml +++ b/examples/pretrain/glm4_moe.toml @@ -8,7 +8,6 @@ max_steps = 4096 [model] -offload = "none" ac = "None" [model.arch] From f43ded9ca4b4199d52dc380b491158c5b0ea9b3c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 12:01:38 +0200 Subject: [PATCH 064/132] fix(pretrain): peak memory measured once, after the optimizer step perf/peak_memory/{max,mean} was read before optimizer.step() and missed the optimizer's buffers, while the step log line read the counter again later. Both now use one reading after the optimizer step, reduced over the ranks. --- src/prime_rl/trainer/pretrain/step.py | 10 +++------- src/prime_rl/trainer/pretrain/train.py | 14 ++++++++++---- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/step.py b/src/prime_rl/trainer/pretrain/step.py index 1a56f64b16f..a9c306281b4 100644 --- a/src/prime_rl/trainer/pretrain/step.py +++ b/src/prime_rl/trainer/pretrain/step.py @@ -30,7 +30,7 @@ class StepResult: """Global ``[num_moe_layers, num_experts]`` token counts of the step (fp32, on the GPU). None for dense models.""" metrics: dict[str, float] - """Statistics of the step over all ranks: data wait and peak memory (max and mean over ranks), and for MoE + """Statistics of the step over all ranks: data wait (max and mean over ranks), and for MoE models the expert balance of the global counts (`expert_load_metrics`), the busiest expert of one rank's micro-batch over the mean (``expert_load/micro_max_mean``, mean and max over ranks and micro-batches) and the routing confidence.""" @@ -85,10 +85,8 @@ def run_micro_batches( sums["micro_max_mean"] = sums.get("micro_max_mean", zero) + micro_max_mean.mean() maxes["micro_max_mean_max"] = torch.maximum(maxes.get("micro_max_mean_max", zero), micro_max_mean.amax()) - # Peak memory since the caller last reset the counter (at the start of the step: the forward and backward passes). - peak_memory = torch.cuda.max_memory_reserved() / 1024**3 - sums |= {"data_time": zero + data_time, "peak_memory": zero + peak_memory} - maxes |= {"data_time_max": zero + data_time, "peak_memory_max": zero + peak_memory} + sums["data_time"] = zero + data_time + maxes["data_time_max"] = zero + data_time extra_sums = extra_sums or {} stats = reduce_stats(sums | {f"extra/{name}": value for name, value in extra_sums.items()}, maxes, dp_cp_group) @@ -99,8 +97,6 @@ def run_micro_batches( metrics = { "time/data/max": stats["data_time_max"].item(), "time/data/mean": stats["data_time"].item() / num_ranks, - "perf/peak_memory/max": stats["peak_memory_max"].item(), - "perf/peak_memory/mean": stats["peak_memory"].item() / num_ranks, } tokens_per_expert = None if is_moe_model: diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 83ec8279cd9..12c5064dd02 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -54,7 +54,7 @@ summarize, ) from prime_rl.trainer.pretrain.step import run_micro_batches -from prime_rl.trainer.pretrain.metrics import DeadExperts, SpikeDetector, grad_norms_by_group +from prime_rl.trainer.pretrain.metrics import DeadExperts, SpikeDetector, grad_norms_by_group, reduce_stats from prime_rl.trainer.sft.data import get_dataset_state from prime_rl.trainer.utils import ( get_dcp_group, @@ -405,6 +405,13 @@ def run_suite(step: int) -> None: optimizer.step() optimizer.zero_grad() current_lr = optimizer.param_groups[0]["lr"] + # Peak memory of the step's forward, backward and optimizer passes, over all ranks. + local_peak_memory = torch.tensor(torch.cuda.max_memory_reserved() / 1024**3, device="cuda") + peak_memory = reduce_stats({"sum": local_peak_memory}, {"max": local_peak_memory}, dp_cp_group) + peak_memory_metrics = { + "perf/peak_memory/max": peak_memory["max"].item(), + "perf/peak_memory/mean": peak_memory["sum"].item() / dist.get_world_size(dp_cp_group), + } # The last step's checkpoint is written once after the loop. save_ckpt_time = 0.0 @@ -444,14 +451,12 @@ def run_suite(step: int) -> None: perf_counter.count_tokens(num_tokens) throughput = perf_counter.get_step_tokens_per_second(num_tokens, step_time) mfu = perf_counter.get_step_mfu(num_tokens, step_time) - peak_memory = torch.cuda.max_memory_reserved() / 1024**3 - step_message = f"Step {progress.step} | {step_time:.2f}s | Loss {result.loss:.4f}" if grad_norm is not None: step_message += f" | Grad. Norm {grad_norm:.4f}" step_message += ( f" | LR {current_lr:.2e} | Throughput {throughput:.0f} tokens/s | MFU {mfu:.1f}%" - f" | Peak Mem. {peak_memory:.1f}/{max_memory:.1f} GiB" + f" | Peak Mem. {peak_memory_metrics['perf/peak_memory/max']:.1f}/{max_memory:.1f} GiB" ) logger.success(step_message) @@ -478,6 +483,7 @@ def run_suite(step: int) -> None: "time/save_ckpt": save_ckpt_time, "time/evals": eval_time, **disk_metrics, + **peak_memory_metrics, **result.metrics, **moe_metrics, **diagnostic_metrics, From 7f5335f341ac0c3d0091603b2bafdf2f3c694993 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 12:03:22 +0200 Subject: [PATCH 065/132] refactor(pretrain): one expert-load implementation and one no-grad cleanup expert_load_metrics builds on compute_expert_load_stats for cv, max_mean and cold_frac and adds only its own metrics (the test keeps only those). The three copies of 'reshard + drop the eval routing counts' become finish_no_grad_forwards, which drops the counts with take_routing_counts. run_evals is run_val, as it runs [val]. Docs and a dashboard comment name the current metrics and the eval suite. --- docs/pretraining.md | 2 +- src/prime_rl/dashboard/static/app.js | 2 +- src/prime_rl/trainer/pretrain/metrics.py | 27 +++++------- src/prime_rl/trainer/pretrain/train.py | 42 +++++++++---------- .../train/pretrain/test_pretrain_metrics.py | 3 -- 5 files changed, 34 insertions(+), 42 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 9faed93ac2b..3ffe92116d8 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -85,7 +85,7 @@ decay_shape = "1-sqrt" # or "linear" (default), "cosine" | `bias_freeze` | Fraction of `max_steps` at the end during which the biases stop updating. Default 0.1. | | `seq_aux_coeff` | Weight (alpha) of the sequence-wise balance loss, computed per packed row of `data.seq_len` tokens and added to the training loss. Default 1e-3; 0 turns it off. | -The metrics are `expert_load/*` (global balance of the step), `expert_load/seq_balance` (1 is balanced) and `expert_load/selection_bias_absmax`. +The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` (1 is balanced) and `router/bias/*`. ## Data diff --git a/src/prime_rl/dashboard/static/app.js b/src/prime_rl/dashboard/static/app.js index db341ffba4e..f09d1cc3689 100644 --- a/src/prime_rl/dashboard/static/app.js +++ b/src/prime_rl/dashboard/static/app.js @@ -446,7 +446,7 @@ const SFT_TRAIN_METRICS = ["loss/mean", "loss/perplexity", "val/loss", "val/perp // max_vio only exists on MoE models: a regex panel stays away on dense ones const SFT_STABILITY_PANELS = [{ metric: "optim/grad_norm" }, { metric: "optim/lr" }, { metric: "loss/nan_count" }, { regex: "max_vio/(mean|max)" }]; const SFT_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memory", "time/forward_backward"]; -// Pretrain flavor (overview.py PRETRAIN_*): held-out bpb from the trainer (val/) and the eval worker (eval/) +// Pretrain flavor (overview.py PRETRAIN_*): held-out bpb from [val] (val/) and the eval suite (eval/) const PRETRAIN_TRAIN_PANELS = [{ metric: "loss/mean" }, { metric: "loss/spike" }, { regex: "val/.*/bpb" }, { regex: "eval/.*/bpb" }]; const PRETRAIN_STABILITY_PANELS = [{ metric: "optim/grad_norm" }, { metric: "optim/grad_norm_spike" }, { metric: "optim/lr" }, { metric: "loss/nan_count" }, { regex: "optim/grad_norm/.*" }]; const PRETRAIN_MOE_METRICS = ["expert_load/max_mean/mean", "expert_load/max_mean/max", "expert_load/micro_max_mean/mean", "expert_load/micro_max_mean/max", "expert_load/entropy/min", "expert_load/dead", "router/confidence", "router/bias/absmax", "loss/seq_balance"]; diff --git a/src/prime_rl/trainer/pretrain/metrics.py b/src/prime_rl/trainer/pretrain/metrics.py index 90fdebd82e0..5ef34f9564f 100644 --- a/src/prime_rl/trainer/pretrain/metrics.py +++ b/src/prime_rl/trainer/pretrain/metrics.py @@ -7,7 +7,7 @@ from torch import Tensor, nn from torch.distributed.tensor import DTensor, Partial, Replicate -from prime_rl.trainer.model import iter_moe_blocks +from prime_rl.trainer.model import compute_expert_load_stats, iter_moe_blocks from prime_rl.trainer.models.layers.moe import MoE from prime_rl.utils.vlm import get_language_model @@ -52,24 +52,19 @@ def moe_layer_indices(model: nn.Module) -> list[int]: def expert_load_metrics(counts: Tensor, layer_indices: list[int]) -> dict[str, float]: """Balance of the `[num_moe_layers, num_experts]` expert counts of a whole step, per layer, summarized as the - mean and the worst layer. `max_mean` / `min_mean`: busiest / least-loaded expert over the mean load; `cv`: - std / mean; `cold_frac`: experts under 0.1x the mean load; `entropy`: entropy of the experts' shares of the - assignments over log(num_experts) (1 = uniform). Also per decoder layer (``layer_indices``): `max_mean` and - `top_share`, the busiest expert's share of the layer's assignments.""" + mean and the worst layer: `cv`, `max_mean` and `cold_frac` as in `compute_expert_load_stats`, plus `min_mean` + (least-loaded expert over the mean load) and `entropy` (entropy of the experts' shares of the assignments over + log(num_experts), 1 = uniform). Also per decoder layer (``layer_indices``): `max_mean` and `top_share`, the + busiest expert's share of the layer's assignments.""" + metrics = {name: value.item() for name, value in compute_expert_load_stats(counts).items()} mean = counts.mean(dim=1) share = counts / counts.sum(dim=1, keepdim=True) - per_layer = { - "max_mean": (counts.amax(dim=1) / mean, "max"), - "min_mean": (counts.amin(dim=1) / mean, "min"), - "cv": (counts.std(dim=1, correction=0) / mean, "max"), - "cold_frac": ((counts < 0.1 * mean[:, None]).double().mean(dim=1), "max"), - "entropy": (-(share * share.clamp_min(1e-30).log()).sum(dim=1) / math.log(counts.shape[1]), "min"), - } - metrics = {} - for name, (values, worst) in per_layer.items(): + min_mean = counts.amin(dim=1) / mean + entropy = -(share * share.clamp_min(1e-30).log()).sum(dim=1) / math.log(counts.shape[1]) + for name, values in (("min_mean", min_mean), ("entropy", entropy)): metrics[f"expert_load/{name}/mean"] = values.mean().item() - metrics[f"expert_load/{name}/{worst}"] = (values.amax() if worst == "max" else values.amin()).item() - for layer, max_mean, top_share in zip(layer_indices, per_layer["max_mean"][0], share.amax(dim=1)): + metrics[f"expert_load/{name}/min"] = values.amin().item() + for layer, max_mean, top_share in zip(layer_indices, counts.amax(dim=1) / mean, share.amax(dim=1)): metrics[f"expert_load/layer_{layer}/max_mean"] = max_mean.item() metrics[f"expert_load/layer_{layer}/top_share"] = top_share.item() return metrics diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 12c5064dd02..656c9bb8378 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -21,7 +21,6 @@ from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager from prime_rl.trainer.model import ( forward, - get_load_balance_stats, is_tt_moe_model, iter_moe_blocks, reshard_module, @@ -54,7 +53,13 @@ summarize, ) from prime_rl.trainer.pretrain.step import run_micro_batches -from prime_rl.trainer.pretrain.metrics import DeadExperts, SpikeDetector, grad_norms_by_group, reduce_stats +from prime_rl.trainer.pretrain.metrics import ( + DeadExperts, + SpikeDetector, + grad_norms_by_group, + reduce_stats, + take_routing_counts, +) from prime_rl.trainer.sft.data import get_dataset_state from prime_rl.trainer.utils import ( get_dcp_group, @@ -274,8 +279,15 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.T ) return out["loss"], loss_mask.sum(dtype=torch.int64), aux_loss_sum + def finish_no_grad_forwards() -> None: + """No-grad forwards leave the [lm_head, norm] FSDP group unsharded, and their routing must not count toward + the next step's MoE statistics.""" + reshard_module(model) + if is_moe_model: + take_routing_counts(model) + @torch.no_grad() - def run_evals(step: int) -> None: + def run_val(step: int) -> None: """Held-out loss, perplexity and bits per byte on the weights after ``step`` optimizer updates.""" metrics = {} for name, source in val_sources.items(): @@ -292,12 +304,8 @@ def run_evals(step: int) -> None: metrics[f"val/{name}/loss"] = loss metrics[f"val/{name}/ppl"] = math.exp(min(loss, 20)) metrics[f"val/{name}/bpb"] = loss_sum / (math.log(2) * num_bytes) - # The no-grad eval forwards leave the [lm_head, norm] FSDP group unsharded, and their - # routing must not count toward the next step's MoE statistics. - reshard_module(model) - if is_moe_model: - get_load_balance_stats(model) - logger.success(f"Evals | Step {step} | " + " | ".join(f"{k} {v:.4f}" for k, v in sorted(metrics.items()))) + finish_no_grad_forwards() + logger.success(f"Val | Step {step} | " + " | ".join(f"{k} {v:.4f}" for k, v in sorted(metrics.items()))) asyncio.run(monitors.log({**metrics, "step": step}, step=step)) eval_sets = [] @@ -336,17 +344,13 @@ def run_suite(step: int) -> None: out.mkdir(parents=True, exist_ok=True) np.savez(out / f"{eval_set.name}@{label}.npz", **arrays) (out / f"{eval_set.name}@{label}.json").write_text(json.dumps(summary, default=float)) - # The no-grad forwards leave the [lm_head, norm] FSDP group unsharded, and their routing must not count - # toward the next step's MoE statistics. - reshard_module(model) - if is_moe_model: - get_load_balance_stats(model) + finish_no_grad_forwards() metrics["time/eval_suite"] = time.perf_counter() - start_time logger.success(f"Eval suite | Step {step} | {metrics['time/eval_suite']:.0f}s | {len(eval_sets)} sets") asyncio.run(monitors.log({**metrics, "step": step}, step=step)) if checkpoint_step is None and config.val is not None and config.val.eval_on_start: - run_evals(step=0) + run_val(step=0) if config.evals is not None and config.evals.eval_on_start: run_suite(step=checkpoint_step or 0) @@ -426,15 +430,11 @@ def run_suite(step: int) -> None: diagnostic_metrics = {} if is_diagnostic_step: diagnostic_metrics = update_tracker.after_step(model) | run_probe(model, probe, dp_cp_group) - # The no-grad forward leaves the [lm_head, norm] FSDP group unsharded, and its routing must not count - # toward the next step's MoE statistics. - reshard_module(model) - if is_moe_model: - get_load_balance_stats(model) + finish_no_grad_forwards() eval_start_time = time.perf_counter() if config.val is not None and (is_last_step or progress.step % config.val.interval == 0): - run_evals(progress.step) + run_val(progress.step) if config.evals is not None and (is_last_step or progress.step % config.evals.interval == 0): run_suite(progress.step) eval_time = time.perf_counter() - eval_start_time diff --git a/tests/unit/train/pretrain/test_pretrain_metrics.py b/tests/unit/train/pretrain/test_pretrain_metrics.py index b88bb30a35e..eaf12314414 100644 --- a/tests/unit/train/pretrain/test_pretrain_metrics.py +++ b/tests/unit/train/pretrain/test_pretrain_metrics.py @@ -7,10 +7,7 @@ def test_expert_load_metrics_per_layer_summaries(): counts = torch.tensor([[1.0, 1.0, 1.0, 1.0], [4.0, 0.0, 0.0, 0.0]], dtype=torch.float64) metrics = expert_load_metrics(counts, layer_indices=[1, 2]) - assert metrics["expert_load/max_mean/mean"] == 2.5 # layers: 1 and 4 - assert metrics["expert_load/max_mean/max"] == 4.0 assert metrics["expert_load/min_mean/min"] == 0.0 - assert metrics["expert_load/cold_frac/max"] == 0.75 assert metrics["expert_load/entropy/mean"] == 0.5 # uniform layer 1, one-hot layer 0 assert metrics["expert_load/layer_2/max_mean"] == 4.0 assert metrics["expert_load/layer_1/top_share"] == 0.25 From 8faa17df69249d15ee47c76fce014590a2f45f2b Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 15:47:20 +0200 Subject: [PATCH 066/132] fix(pretrain): bounded eval-suite logits; free the diagnostic weight copy before checkpoint saves The eval suite scored 16384 tokens per chunk with full-vocab fp32 logits (and an fp32 copy of the LM head per chunk): ~17 GB at a 129k vocab. Chunks are now 2048 tokens and the fp32 weight is made once. The diagnostic update sizes are computed right after the optimizer step, so the weight copy is freed before a checkpoint save. --- src/prime_rl/trainer/pretrain/suite.py | 8 +++++--- src/prime_rl/trainer/pretrain/train.py | 4 ++-- tests/unit/train/pretrain/test_suite.py | 12 +++++++++++- 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/suite.py b/src/prime_rl/trainer/pretrain/suite.py index f7c2ba9830c..0ad70ba94f4 100644 --- a/src/prime_rl/trainer/pretrain/suite.py +++ b/src/prime_rl/trainer/pretrain/suite.py @@ -48,11 +48,13 @@ class EvalSet: sha256: str -def head_logprobs(hidden, weight, targets, chunk: int = 16384): - """Log-prob and greedy hit of ``targets`` under ``hidden @ weight.T``, in fp32 chunks over tokens.""" +def head_logprobs(hidden, weight, targets, chunk: int = 2048): + """Log-prob and greedy hit of ``targets`` under ``hidden @ weight.T``, in fp32 chunks over tokens, so the + fp32 logits stay at ``chunk x vocab``.""" + weight = weight.float() logprobs, hits = [], [] for start in range(0, len(targets), chunk): - logits = hidden[start : start + chunk].float() @ weight.float().T + logits = hidden[start : start + chunk].float() @ weight.T target = targets[start : start + chunk] logprobs.append(logits.gather(-1, target[:, None])[:, 0] - torch.logsumexp(logits, -1)) hits.append(logits.argmax(-1) == target) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 656c9bb8378..dd3ab1d33d6 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -409,6 +409,7 @@ def run_suite(step: int) -> None: optimizer.step() optimizer.zero_grad() current_lr = optimizer.param_groups[0]["lr"] + diagnostic_metrics = update_tracker.after_step(model) if is_diagnostic_step else {} # Peak memory of the step's forward, backward and optimizer passes, over all ranks. local_peak_memory = torch.tensor(torch.cuda.max_memory_reserved() / 1024**3, device="cuda") peak_memory = reduce_stats({"sum": local_peak_memory}, {"max": local_peak_memory}, dp_cp_group) @@ -427,9 +428,8 @@ def run_suite(step: int) -> None: save_ckpt_time = time.perf_counter() - save_ckpt_start_time disk_metrics = get_ckpt_disk_metrics(config.run_dir) - diagnostic_metrics = {} if is_diagnostic_step: - diagnostic_metrics = update_tracker.after_step(model) | run_probe(model, probe, dp_cp_group) + diagnostic_metrics |= run_probe(model, probe, dp_cp_group) finish_no_grad_forwards() eval_start_time = time.perf_counter() diff --git a/tests/unit/train/pretrain/test_suite.py b/tests/unit/train/pretrain/test_suite.py index 0b7afab772c..ea8d59f9ea7 100644 --- a/tests/unit/train/pretrain/test_suite.py +++ b/tests/unit/train/pretrain/test_suite.py @@ -4,8 +4,9 @@ import numpy as np import pyarrow as pa import pyarrow.parquet as pq +import torch -from prime_rl.trainer.pretrain.suite import pack_requests, prepare_multi_length_set, summarize_ctxcont +from prime_rl.trainer.pretrain.suite import head_logprobs, pack_requests, prepare_multi_length_set, summarize_ctxcont def test_pack_requests_covers_every_request_within_the_budget(): @@ -45,3 +46,12 @@ def test_multi_length_set_scores_shared_requests_once(tmp_path): assert len(multi.versions[8].requests) == 4 and len(multi.versions[32].requests) == 2 assert len(multi.requests) == 5 assert multi.index[8][0] == multi.index[32][0] + + +def test_head_logprobs_chunks_match_the_full_softmax(): + hidden, weight = torch.randn(7, 4, dtype=torch.bfloat16), torch.randn(11, 4, dtype=torch.bfloat16) + targets = torch.randint(0, 11, (7,)) + logprobs, hits = head_logprobs(hidden, weight, targets, chunk=3) + logits = hidden.float() @ weight.float().T + torch.testing.assert_close(logprobs, logits.log_softmax(-1).gather(-1, targets[:, None])[:, 0]) + assert torch.equal(hits, logits.argmax(-1) == targets) From 11f0ee1c807544692c8a8bd5051eab1622456e15 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 16:48:15 +0200 Subject: [PATCH 067/132] perf(pretrain): score all eval sets in one packing pass The suite packed and dealt each set's requests separately, so every set ended with a partly empty round of padding forwards on many ranks (~46% of the forwards at 48 GPUs). score_sets packs the requests of all sets together, gathers once and splits the results per set. A GPU test checks the results equal scoring each set alone. --- src/prime_rl/trainer/pretrain/suite.py | 17 ++++++++++ src/prime_rl/trainer/pretrain/train.py | 11 +++---- tests/unit/train/pretrain/test_suite.py | 41 ++++++++++++++++++++++++- 3 files changed, 61 insertions(+), 8 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/suite.py b/src/prime_rl/trainer/pretrain/suite.py index 0ad70ba94f4..fe00d716457 100644 --- a/src/prime_rl/trainer/pretrain/suite.py +++ b/src/prime_rl/trainer/pretrain/suite.py @@ -145,6 +145,23 @@ def score_requests( return logprobs, hits +def score_sets( + model: nn.Module, eval_sets: list["MultiLengthSet"], token_budget: int, group: dist.ProcessGroup +) -> list[tuple[list[np.ndarray], list[np.ndarray]]] | None: + """`score_requests` over the requests of all sets at once, so the rows of every set are packed and dealt + together (fewer padding forwards on many ranks), split back per set on rank 0 (None on the others).""" + scored = score_requests(model, [request for s in eval_sets for request in s.requests], token_budget, group) + if scored is None: + return None + logprobs, hits = scored + per_set, offset = [], 0 + for eval_set in eval_sets: + end = offset + len(eval_set.requests) + per_set.append((logprobs[offset:end], hits[offset:end])) + offset = end + return per_set + + def continuation_ids(tokenizer: PreTrainedTokenizer, context: str, continuation: str) -> Request: """BOS + jointly tokenized context + continuation, and the index of the first continuation token (the first token that ends after the context; a token straddling the boundary counts as continuation).""" diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index dd3ab1d33d6..95723e52798 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -48,7 +48,7 @@ from prime_rl.trainer.pretrain.suite import ( length_label, prepare_multi_length_set, - score_requests, + score_sets, scoring_head_active, summarize, ) @@ -319,17 +319,14 @@ def run_val(step: int) -> None: @torch.no_grad() def run_suite(step: int) -> None: - """The eval suite on the weights after ``step`` optimizer updates: the packed rows of every set are dealt + """The eval suite on the weights after ``step`` optimizer updates: the packed rows of all sets are dealt over the ranks, the per-token results gathered on rank 0, which summarizes and logs them.""" start_time = time.perf_counter() metrics = {} # Eager, also for compiled layers: the packed rows vary in length and would recompile at every forward. with scoring_head_active(model), torch.compiler.set_stance("force_eager"): - for eval_set in eval_sets: - scored = score_requests(model, eval_set.requests, config.evals.token_budget, get_dcp_group()) - if scored is None: - continue - logprobs, hits = scored + scored = score_sets(model, eval_sets, config.evals.token_budget, get_dcp_group()) + for eval_set, (logprobs, hits) in zip(eval_sets, scored or []): for seq_len, version in eval_set.versions.items(): index = eval_set.index[seq_len] summary, arrays = summarize(version, [logprobs[i] for i in index], [hits[i] for i in index]) diff --git a/tests/unit/train/pretrain/test_suite.py b/tests/unit/train/pretrain/test_suite.py index ea8d59f9ea7..7d0c3b5b2dc 100644 --- a/tests/unit/train/pretrain/test_suite.py +++ b/tests/unit/train/pretrain/test_suite.py @@ -4,9 +4,20 @@ import numpy as np import pyarrow as pa import pyarrow.parquet as pq +import pytest import torch +import torch.distributed as dist +from torch import nn -from prime_rl.trainer.pretrain.suite import head_logprobs, pack_requests, prepare_multi_length_set, summarize_ctxcont +from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX +from prime_rl.trainer.pretrain.suite import ( + head_logprobs, + pack_requests, + prepare_multi_length_set, + score_requests, + score_sets, + summarize_ctxcont, +) def test_pack_requests_covers_every_request_within_the_budget(): @@ -55,3 +66,31 @@ def test_head_logprobs_chunks_match_the_full_softmax(): logits = hidden.float() @ weight.float().T torch.testing.assert_close(logprobs, logits.log_softmax(-1).gather(-1, targets[:, None])[:, 0]) assert torch.equal(hits, logits.argmax(-1) == targets) + + +class TokenPositionModel(nn.Module): + """Scores each label by its value and position, so a request's scores do not depend on how it is packed.""" + + def forward(self, input_ids, position_ids, labels, seq_lens): + keep = labels != IGNORE_INDEX + return {"logprobs": (labels + position_ids / 1000)[keep].double(), "entropy": (labels % 2 == 0)[keep]} + + +@pytest.mark.gpu +def test_score_sets_matches_scoring_each_set_alone(tmp_path): + dist.init_process_group("gloo", init_method=f"file://{tmp_path / 'store'}", rank=0, world_size=1) + try: + for name, texts in {"a": ["short", "a much longer document"], "b": ["xyz", "another text here"]}.items(): + (tmp_path / name).mkdir() + pq.write_table(pa.table({"text": texts}), tmp_path / name / "test.parquet") + (tmp_path / name / "meta.json").write_text(json.dumps({"kind": "ppl", "parquet_sha256": "x"})) + eval_sets = [prepare_multi_length_set(name, tmp_path, CharTokenizer(), [8, 32]) for name in "ab"] + together = score_sets(TokenPositionModel(), eval_sets, token_budget=40, group=dist.group.WORLD) + for eval_set, (logprobs, hits) in zip(eval_sets, together): + alone_logprobs, alone_hits = score_requests( + TokenPositionModel(), eval_set.requests, token_budget=40, group=dist.group.WORLD + ) + assert all(np.array_equal(x, y) for x, y in zip(logprobs, alone_logprobs)) + assert all(np.array_equal(x, y) for x, y in zip(hits, alone_hits)) + finally: + dist.destroy_process_group() From 8a47d4b737746b1816ec1de75af475314faba74d Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 17:02:49 +0200 Subject: [PATCH 068/132] fix(pretrain): eval-suite metrics as floats; pairs requests cut to the window run_suite kept only summary values that are Python floats; set_metrics now converts every summary value with float(). pairs sentences were not cut to seq_len, so one could exceed token_budget; they are cut to the window like ppl and ctxcont. --- src/prime_rl/trainer/pretrain/suite.py | 9 ++++++++- src/prime_rl/trainer/pretrain/train.py | 7 ++----- tests/unit/train/pretrain/test_suite.py | 17 +++++++++++++++++ 3 files changed, 27 insertions(+), 6 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/suite.py b/src/prime_rl/trainer/pretrain/suite.py index fe00d716457..55a97c5eee8 100644 --- a/src/prime_rl/trainer/pretrain/suite.py +++ b/src/prime_rl/trainer/pretrain/suite.py @@ -210,7 +210,8 @@ def prepare_set(name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_l elif kind == "pairs": good, bad = table.column("good").to_pylist(), table.column("bad").to_pylist() encoded = tokenizer(good + bad, add_special_tokens=False)["input_ids"] - requests = [([tokenizer.bos_token_id, *ids], 1) for ids in encoded] + # Cut to the window like the other kinds, so every request fits in a row of ``token_budget`` tokens. + requests = [([tokenizer.bos_token_id, *ids][: seq_len + 1], 1) for ids in encoded] context = {"num_good": len(good)} else: raise ValueError(f"Unknown eval set kind {kind!r} ({name})") @@ -242,6 +243,12 @@ def prepare_multi_length_set( return MultiLengthSet(name, versions, [(list(ids), start) for ids, start in distinct], index) +def set_metrics(name: str, seq_len: int, summary: dict) -> dict[str, float]: + """A set's summary as metrics ``eval//@``, as Python floats (the summaries hold NumPy + scalars).""" + return {f"eval/{name}/{key}@{length_label(seq_len)}": float(value) for key, value in summary.items()} + + def length_label(seq_len: int) -> str: """``2048`` -> ``2k``, the suffix of the metric names.""" return f"{seq_len // 1024}k" if seq_len % 1024 == 0 else str(seq_len) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 95723e52798..45fd1f8999f 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -49,6 +49,7 @@ length_label, prepare_multi_length_set, score_sets, + set_metrics, scoring_head_active, summarize, ) @@ -331,11 +332,7 @@ def run_suite(step: int) -> None: index = eval_set.index[seq_len] summary, arrays = summarize(version, [logprobs[i] for i in index], [hits[i] for i in index]) label = length_label(seq_len) - metrics |= { - f"eval/{eval_set.name}/{key}@{label}": value - for key, value in summary.items() - if isinstance(value, float) - } + metrics |= set_metrics(eval_set.name, seq_len, summary) if config.evals.save_arrays: out = config.run_dir / "evals" / f"step_{step}" out.mkdir(parents=True, exist_ok=True) diff --git a/tests/unit/train/pretrain/test_suite.py b/tests/unit/train/pretrain/test_suite.py index 7d0c3b5b2dc..15b7f87004c 100644 --- a/tests/unit/train/pretrain/test_suite.py +++ b/tests/unit/train/pretrain/test_suite.py @@ -14,8 +14,11 @@ head_logprobs, pack_requests, prepare_multi_length_set, + prepare_set, score_requests, score_sets, + set_metrics, + summarize, summarize_ctxcont, ) @@ -39,6 +42,8 @@ def test_summarize_ctxcont_scores_the_gold_choice(): assert summary["acc"] == 0.5 # item 0: gold wins; item 1: the other choice wins assert math.isclose(summary["bpb"], (2.0 + 2.0) / (3 * math.log(2))) # gold nll / gold continuation bytes assert summary["greedy_em"] == 0.0 + metrics = set_metrics("mc", 2048, summary) + assert metrics["eval/mc/acc@2k"] == 0.5 and all(type(value) is float for value in metrics.values()) class CharTokenizer: @@ -94,3 +99,15 @@ def test_score_sets_matches_scoring_each_set_alone(tmp_path): assert all(np.array_equal(x, y) for x, y in zip(hits, alone_hits)) finally: dist.destroy_process_group() + + +def test_pairs_requests_fit_the_window(tmp_path): + (tmp_path / "pairs").mkdir() + table = pa.table({"good": ["a long good sentence", "ok"], "bad": ["a long bad sentence!", "no"]}) + pq.write_table(table, tmp_path / "pairs" / "test.parquet") + (tmp_path / "pairs" / "meta.json").write_text(json.dumps({"kind": "pairs", "parquet_sha256": "x"})) + eval_set = prepare_set("pairs", tmp_path, CharTokenizer(), seq_len=8) + assert all(len(ids) - 1 <= 8 for ids, _ in eval_set.requests) + logprobs = [np.full(len(ids) - start, -1.0) for ids, start in eval_set.requests] + summary, _ = summarize(eval_set, logprobs, [np.ones_like(lp, dtype=bool) for lp in logprobs]) + assert set_metrics("pairs", 8, summary)["eval/pairs/acc@8"] == 0.0 # equal scores: good never beats bad From b1bec3d7ec6566288da530d45407f157a4b95193 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 17:18:36 +0200 Subject: [PATCH 069/132] docs(pretrain): monitoring section; step-to-row numbering; config docstrings The docs had no section on the step metrics and [diagnostics] after the stack was merged, and numbered the rows of step s from 0 while steps count from 1. The [val] and [monitors] docstrings name what they cover now. --- docs/pretraining.md | 8 +++++++- .../prime-rl-configs/src/prime_rl/configs/pretrain.py | 4 ++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 3ffe92116d8..c185bf10e32 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -115,6 +115,12 @@ weights = { web = 0.5, code = 0.5 } Batch sizes are in tokens, so changing `seq_len` keeps the tokens per step and the number of steps. Resume restores the data position on the same GPU layout and `data.num_workers`; resuming on a different layout is not supported yet. `data.type = "fake"` trains on random tokens, for smoke and throughput tests. +## Monitoring + +Every step logs the loss, learning rate and schedule multiplier, gradient norms (total and per parameter group, `optim/grad_norm/`), loss and grad-norm spike flags, per-source token counts and epochs (`progress//*`), throughput, MFU, data wait and peak memory (max and mean over ranks). MoE models add the expert balance of the step's global counts (`expert_load/*`, also per layer), `router/*` and `loss/seq_balance`. The step statistics are reduced once per step. + +`[diagnostics]` (default: every 100 steps; `diagnostics = null` turns it off) adds the size of the optimizer update relative to the weights per parameter group (`update/rms_ratio/*`, `param/rms/*`) and a forward pass on a fixed probe batch: attention max logit and entropy per layer (`attn/*`, FlashAttention layers) and the residual-stream size per layer (`act/*`). It does not support context parallelism yet. + ## Evals `[val]` computes the held-out loss (nats per token), perplexity and bits per UTF-8 byte (comparable across tokenizers) of every source in `val.sources`, through the same loss code as training, so it also works with context parallelism. Each eval reads the same leading windows of every source: as many as fill whole micro-batches on every data-parallel rank. Evals run after the optimizer update, so the metrics at step N describe the same weights as the step-N checkpoint. With `eval_on_start` (default), they also run before the first step (logged as step 0), and they always run at the last step. @@ -155,4 +161,4 @@ partition = "all" time = "04:00:00" ``` -Row `w` of the token stream goes to data stream `w % num_streams`, with `num_streams = data_parallel_ranks * data.num_workers`, and step `s` trains on rows `s * rows_per_step` to `(s + 1) * rows_per_step - 1`. So runs with different numbers of GPUs train on the same rows at every step, only split differently into micro-batches. The loss and gradient are normalized by the token count of the whole global batch, as on one GPU with gradient accumulation. +Row `w` of the token stream goes to data stream `w % num_streams`, with `num_streams = data_parallel_ranks * data.num_workers`, and step `s` (counted from 1) trains on rows `(s - 1) * rows_per_step` to `s * rows_per_step - 1`. So runs with different numbers of GPUs train on the same rows at every step, only split differently into micro-batches. The loss and gradient are normalized by the token count of the whole global batch, as on one GPU with gradient accumulation. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 8d90964ad7d..67d9beb7d2e 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -301,7 +301,7 @@ class PretrainConfig(BaseConfig): data: PretrainDataConfig val: ValConfig | None = None - """Held-out evals. If None, none run.""" + """Held-out loss on token shards (``val/*``). If None, it never runs.""" diagnostics: DiagnosticsConfig | None = DiagnosticsConfig() """Diagnostic step (update sizes, attention and activation probe). If None, it never runs.""" @@ -324,7 +324,7 @@ class PretrainConfig(BaseConfig): log: TrainerLogConfig = TrainerLogConfig() monitors: TrainMonitorsConfig = TrainMonitorsConfig() - """Metric monitors (``monitors.wandb``, ``monitors.file``).""" + """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" run: RunConfig = Field(default_factory=RunConfig) """Run metadata. ``run.name`` names the run directory under ``output_dir``.""" From 0bc6af520b6cb303f1aa79f96f13861cc6608783 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 19:26:14 +0200 Subject: [PATCH 070/132] feat(pretrain): shuffle each epoch of a token source Every epoch reads every token once, in blocks of 2^20 tokens in a seeded random order (data.seed, source, epoch) from a seeded random start offset, so epochs differ in order and window boundaries and part of an epoch samples the whole source, whatever the order of the shard files. Within an epoch, the windows of each group of 4,096 consecutive windows are read in a seeded random order, so the pieces of one long document spread over many steps. The resume state is unchanged (tokens read per source); held-out sources keep the file-name order. Experiment: pretrain-experiments data/2026-10-15-multi-epoch-shuffle. --- docs/pretraining.md | 2 +- .../src/prime_rl/configs/pretrain.py | 2 +- src/prime_rl/trainer/pretrain/data.py | 78 ++++++++++++++++--- tests/unit/train/pretrain/test_data.py | 29 +++++++ 4 files changed, 99 insertions(+), 12 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index c185bf10e32..1d144fb080a 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -89,7 +89,7 @@ The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` ## Data -`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of a source are read in name order as one stream; documents continue across rows, and attention restarts at every document and every row. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. +`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. Each epoch reads every token of a source once: in blocks of 2^20 tokens in a random order and from a random start offset, both drawn from `data.seed`, the source and the epoch. Within an epoch, the rows of each group of 4,096 consecutive rows are read in a seeded random order, so the pieces of one long document spread over many steps. So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole source, whatever the order of the files. Documents continue across rows, and attention restarts at every document, every row and every block boundary. Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. Each row comes from one source, drawn with the weights of the current mixture phase. Phases start at fractions of `max_steps`. At startup the trainer computes how many passes over each source the run makes and fails above `data.max_epochs`. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 67d9beb7d2e..9fafb5cb757 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -132,7 +132,7 @@ class BaseBatchConfig(BaseConfig): """Dataloader worker processes per GPU.""" seed: int = 0 - """Seed for choosing each row's source (and the fake tokens).""" + """Seed for choosing each row's source, the per-epoch order of every source (and the fake tokens).""" @property def batch_size(self) -> int: diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index b5f9a660437..0d60d68cc94 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -1,4 +1,5 @@ import bisect +import zlib from pathlib import Path import numpy as np @@ -9,15 +10,28 @@ from prime_rl.trainer.sft.data import StatefulIterableDataset from prime_rl.utils.logger import get_logger +# Unit of the per-epoch shuffle of a token source. A block boundary inside a document splits it, about once +# per 256 windows of 4,096 tokens. +BLOCK_TOKENS = 1 << 20 +# Unit of the window shuffle of a token source: consecutive windows, so the pieces of one long document spread +# over many steps. +WINDOW_GROUP = 4096 + class TokenShards: - """Datatrove tokenized shards in one folder, read as one token stream in file-name order. + """Datatrove tokenized shards in one folder, the ``num_tokens`` first tokens of the files in name order. ``.ds`` holds the token ids (little-endian, 2 or 4 bytes as recorded on the first line of ``.ds.metadata``) and ``.ds.index`` the uint64 end offset of every document. File - boundaries are document boundaries. The stream wraps around after ``num_tokens``.""" + boundaries are document boundaries. + + Stream position ``p`` belongs to epoch ``p // num_tokens``, and every epoch reads every token once. With a + ``seed``, each epoch reads the tokens in blocks of ``BLOCK_TOKENS`` in a seeded random order, starting at a + seeded random offset, so epochs differ in order and window boundaries, and part of an epoch is a sample of + the whole source, and the windows are read shuffled within groups (``window_start``). Without one + (held-out sets), the stream is the files in name order.""" - def __init__(self, config: TokenSourceConfig): + def __init__(self, config: TokenSourceConfig, seed: tuple[int, str] | None = None): self.files = sorted(config.path.glob("*.ds")) if not self.files: raise FileNotFoundError(f"No datatrove .ds files in {config.path}") @@ -28,7 +42,10 @@ def __init__(self, config: TokenSourceConfig): if config.max_tokens is not None and config.max_tokens > total: raise ValueError(f"{config.path} has {total:,} tokens, fewer than max_tokens={config.max_tokens:,}") self.num_tokens = total if config.max_tokens is None else config.max_tokens + self.seed = None if seed is None else [seed[0], zlib.crc32(seed[1].encode())] self.memmaps: dict[int, tuple[np.memmap, np.memmap]] = {} + self.layouts: dict[int, tuple[np.ndarray, np.ndarray, int]] = {} + self.window_order: tuple[tuple[int, int] | None, np.ndarray | None] = (None, None) def open(self, file_index: int) -> tuple[np.memmap, np.memmap]: # Opened lazily, so every dataloader worker maps the files itself. @@ -39,17 +56,39 @@ def open(self, file_index: int) -> tuple[np.memmap, np.memmap]: self.memmaps[file_index] = tokens, doc_ends return self.memmaps[file_index] + def layout(self, epoch: int) -> tuple[np.ndarray, np.ndarray, int]: + """An epoch's block order, the stream offsets where its blocks start, and its start offset.""" + if epoch not in self.layouts: + rng = np.random.default_rng([*self.seed, epoch]) + order = rng.permutation(-(-self.num_tokens // BLOCK_TOKENS)) + lengths = np.minimum(BLOCK_TOKENS, self.num_tokens - order * BLOCK_TOKENS) + self.layouts[epoch] = order, np.cumsum([0, *lengths]), int(rng.integers(self.num_tokens)) + return self.layouts[epoch] + + def locate(self, position: int, length: int) -> tuple[int, int, int]: + """File index, offset in the file and length of the contiguous file run at stream ``position``, + at most ``length`` tokens.""" + epoch, position = divmod(position, self.num_tokens) + length = min(length, self.num_tokens - position) + if self.seed is not None: + order, block_starts, offset = self.layout(epoch) + position = (position + offset) % self.num_tokens + block = int(np.searchsorted(block_starts, position, side="right")) - 1 + length = min(length, int(block_starts[block + 1]) - position) + position = int(order[block]) * BLOCK_TOKENS + position - int(block_starts[block]) + file_index = int(np.searchsorted(self.starts, position, side="right")) - 1 + length = min(length, int(self.starts[file_index + 1]) - position) + return file_index, position - int(self.starts[file_index]), length + def read(self, start: int, length: int) -> tuple[np.ndarray, list[int]]: - """Tokens ``[start, start + length)`` of the stream and the offsets in them where a document starts.""" + """Tokens ``[start, start + length)`` of the stream and the offsets in them where a document starts. + A block cut through a document starts a new one.""" pieces, doc_starts = [], [] offset = 0 while offset < length: if offset > 0: doc_starts.append(offset) - position = (start + offset) % self.num_tokens - file_index = int(np.searchsorted(self.starts, position, side="right")) - 1 - local = position - int(self.starts[file_index]) - size = min(length - offset, int(self.starts[file_index + 1]) - position, self.num_tokens - position) + file_index, local, size = self.locate(start + offset, length - offset) tokens, doc_ends = self.open(file_index) pieces.append(tokens[local : local + size]) first, last = np.searchsorted(doc_ends, [local, local + size], side="right") @@ -57,6 +96,22 @@ def read(self, start: int, length: int) -> tuple[np.ndarray, list[int]]: offset += size return np.concatenate(pieces).astype(np.int64), doc_starts + def window_start(self, position: int, seq_len: int) -> int: + """Stream position of the window to read when ``position`` tokens have been read. With a seed, each epoch's + windows are read in a seeded random order within groups of ``WINDOW_GROUP`` consecutive windows.""" + if self.seed is None: + return position + epoch = position // self.num_tokens + group = (position - epoch * self.num_tokens) // (WINDOW_GROUP * seq_len) + group_start = epoch * self.num_tokens + group * WINDOW_GROUP * seq_len + group_end = min((epoch + 1) * self.num_tokens, group_start + WINDOW_GROUP * seq_len) + # The windows that start in [group_start, group_end). + first, last = -(-group_start // seq_len), -(-group_end // seq_len) + if self.window_order[0] != (epoch, group): + order = np.random.default_rng([*self.seed, epoch, group, 1]).permutation(last - first) + self.window_order = (epoch, group), order + return (first + int(self.window_order[1][position // seq_len - first])) * seq_len + def last_token_of_first_document(self) -> int: tokens, doc_ends = self.open(0) return int(tokens[int(doc_ends[0]) - 1]) @@ -79,6 +134,9 @@ def __init__(self, vocab_size: int, seed: int): def read(self, start: int, length: int) -> tuple[np.ndarray, list[int]]: return np.random.default_rng([self.seed, start]).integers(0, self.vocab_size, length), [] + def window_start(self, position: int, seq_len: int) -> int: + return position + class TokenDataset(StatefulIterableDataset): """``seq_len`` windows of token sources, each window from one source of a weighted mixture. @@ -145,7 +203,7 @@ def __iter__(self): while self.num_windows is None or self.step < self.num_windows: window = self.step name = self.source(window) - start = self.tokens[name] + start = self.sources[name].window_start(self.tokens[name], self.seq_len) self.tokens[name] += self.seq_len self.step += 1 if window % self.data_world_size != self.data_rank: @@ -195,7 +253,7 @@ def setup_train_dataset( sources = {"fake": RandomTokens(vocab_size, config.seed)} phases = [MixturePhaseConfig(weights={"fake": 1.0})] else: - sources = {name: TokenShards(source) for name, source in config.sources.items()} + sources = {name: TokenShards(source, seed=(config.seed, name)) for name, source in config.sources.items()} check_eos(sources, eos_token_id) phases = config.phases dataset = TokenDataset( diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index 52db6180bd1..127240f5ec4 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -4,6 +4,7 @@ import numpy as np from prime_rl.configs.pretrain import MixturePhaseConfig, TokenSourceConfig +from prime_rl.trainer.pretrain import data from prime_rl.trainer.pretrain.data import TokenDataset, TokenShards, collate_rows, source_tokens @@ -28,6 +29,34 @@ def test_token_shards_read_across_files_and_wrap(tmp_path): assert doc_starts == [2] +def test_shuffled_epochs_read_every_token_once_in_a_new_order(tmp_path, monkeypatch): + monkeypatch.setattr(data, "BLOCK_TOKENS", 4) + write_shard(tmp_path / "a" / "00.ds", [list(range(0, 30)), list(range(30, 50))]) + write_shard(tmp_path / "a" / "01.ds", [list(range(50, 70))]) + shards = TokenShards(TokenSourceConfig(path=tmp_path / "a"), seed=(0, "a")) + epochs = [shards.read(epoch * 70, 70)[0].tolist() for epoch in range(3)] + assert all(sorted(epoch) == list(range(70)) for epoch in epochs) + assert len({tuple(epoch) for epoch in epochs}) == 3 + # Reading in windows gives the same stream, so resuming at any position continues it. + windows = [shards.read(start, 7)[0].tolist() for start in range(70, 140, 7)] + assert sum(windows, []) == epochs[1] + + +def test_window_shuffle_reads_every_window_of_an_epoch_once(tmp_path, monkeypatch): + monkeypatch.setattr(data, "WINDOW_GROUP", 5) + write_shard(tmp_path / "a" / "00.ds", [list(range(70))]) + shards = TokenShards(TokenSourceConfig(path=tmp_path / "a"), seed=(0, "a")) + starts = [shards.window_start(position, 7) for position in range(0, 140, 7)] + groups = [starts[i : i + 5] for i in range(0, 20, 5)] + assert [sorted(group) for group in groups] == [list(range(i * 7, i * 7 + 35, 7)) for i in range(0, 20, 5)] + assert starts != list(range(0, 140, 7)) + # A group never reaches into the next epoch. + monkeypatch.setattr(data, "WINDOW_GROUP", 100) + shards = TokenShards(TokenSourceConfig(path=tmp_path / "a"), seed=(0, "a")) + starts = [shards.window_start(position, 7) for position in range(0, 140, 7)] + assert sorted(starts[:10]) == list(range(0, 70, 7)) and sorted(starts[10:]) == list(range(70, 140, 7)) + + def test_token_dataset_resumes_and_follows_phases(tmp_path): write_shard(tmp_path / "a" / "00.ds", [list(range(100, 200))]) write_shard(tmp_path / "b" / "00.ds", [list(range(200, 300))]) From fdc61825b0fc8df8ccec26b2442d7f29ecd3d99f Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 19:26:59 +0200 Subject: [PATCH 071/132] refactor(pretrain): specify the model like SFT/RL; move the init to the shared model setup model.name points at an HF model id or model dir with a config.json, as in SFT and RL; [model.arch] and the launcher-written model dir are removed. Checked-in configs for the example (examples/pretrain/glm4-moe-d512) and the debug run. The weight init moves to prime_rl.trainer.models.init and runs in setup_model when given an init config; SFT and RL do not pass one. --- configs/debug/fake/pretrain-model/config.json | 28 +++++++++ configs/debug/fake/pretrain.toml | 9 +-- docs/pretraining.md | 18 +++--- examples/pretrain/glm4-moe-d512/config.json | 37 ++++++++++++ examples/pretrain/glm4_moe.toml | 24 +------- .../src/prime_rl/configs/pretrain.py | 58 ++++--------------- .../src/prime_rl/configs/trainer.py | 23 ++++++++ src/prime_rl/entrypoints/pretrain.py | 32 +--------- src/prime_rl/monitors/prime.py | 2 +- src/prime_rl/trainer/model.py | 14 ++++- src/prime_rl/trainer/{ => models}/init.py | 18 ++++-- .../trainer/pretrain/parametrization.py | 12 +--- src/prime_rl/trainer/pretrain/train.py | 14 +++-- tests/unit/train/pretrain/test_init.py | 5 +- tests/unit/train/pretrain/test_launcher.py | 2 +- tests/unit/train/pretrain/test_model_dir.py | 34 ----------- tests/unit/train/pretrain/test_schedules.py | 2 +- 17 files changed, 154 insertions(+), 178 deletions(-) create mode 100644 configs/debug/fake/pretrain-model/config.json create mode 100644 examples/pretrain/glm4-moe-d512/config.json rename src/prime_rl/trainer/{ => models}/init.py (86%) delete mode 100644 tests/unit/train/pretrain/test_model_dir.py diff --git a/configs/debug/fake/pretrain-model/config.json b/configs/debug/fake/pretrain-model/config.json new file mode 100644 index 00000000000..e991ed03b69 --- /dev/null +++ b/configs/debug/fake/pretrain-model/config.json @@ -0,0 +1,28 @@ +{ + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": null, + "eos_token_id": 151645, + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 256, + "initializer_range": 0.02, + "intermediate_size": 1024, + "max_position_embeddings": 2048, + "mlp_bias": false, + "model_type": "llama", + "num_attention_heads": 4, + "num_hidden_layers": 2, + "num_key_value_heads": 2, + "pad_token_id": 151643, + "pretraining_tp": 1, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "rope_theta": 10000.0, + "rope_type": "default" + }, + "tie_word_embeddings": false, + "transformers_version": "5.6.2", + "use_cache": true, + "vocab_size": 151680 +} diff --git a/configs/debug/fake/pretrain.toml b/configs/debug/fake/pretrain.toml index db0302142b5..0b9de1df8db 100644 --- a/configs/debug/fake/pretrain.toml +++ b/configs/debug/fake/pretrain.toml @@ -1,16 +1,9 @@ max_steps = 5 [model] +name = "configs/debug/fake/pretrain-model" compile = "None" -[model.arch] -model_type = "llama" -hidden_size = 256 -num_hidden_layers = 2 -num_attention_heads = 4 -num_key_value_heads = 2 -intermediate_size = 1024 - [tokenizer] name = "PrimeIntellect/Qwen3-0.6B" diff --git a/docs/pretraining.md b/docs/pretraining.md index c185bf10e32..25c99db1237 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -8,25 +8,21 @@ uv run pretrain @ examples/pretrain/glm4_moe.toml ## Model -`[model.arch]` is the whole architecture: a `model_type` from `prime_rl.trainer.models` and the fields of its config class. Unknown fields are rejected. The vocabulary size (`len(tokenizer)` rounded up to a multiple of 128) and the BOS, EOS and pad ids come from the tokenizer, so they cannot disagree with it. +The model is set up as in SFT and RL: `model.name` is an HF model id or a local model directory with a `config.json`. Its weights are not loaded. The model starts from a random init (below) or from a checkpoint. A new architecture is a directory with only a `config.json` for a model type in `prime_rl.trainer.models`; [`examples/pretrain/glm4-moe-d512`](../examples/pretrain/glm4-moe-d512/config.json) is the one in the example. `vocab_size` and the BOS, EOS and pad ids in the config must fit the tokenizer. The checked-in configs round `len(tokenizer)` up to a multiple of 128 and use the tokenizer's special ids, with EOS as the pad id if the tokenizer has none. ```toml -[model.arch] -model_type = "glm4_moe" -hidden_size = 512 -num_hidden_layers = 8 -n_routed_experts = 64 -num_experts_per_tok = 6 +[model] +name = "examples/pretrain/glm4-moe-d512" [tokenizer] -name = "tokenizers/deepseek-v4-flash-base" +name = "tokenizers/deepseek-v4-flash-base" # defaults to model.name ``` -The launcher writes the config and the tokenizer to `/configs/model` and trains from there, so the trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all see an ordinary model directory. The trainer logs the active non-embedding (routed experts counted at `top_k / num_experts`), non-embedding, embedding (input embedding and LM head) and total parameter counts at startup, and as `params/*` at step 0. +The trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all read the same model directory and tokenizer. The trainer logs the active non-embedding (routed experts counted at `top_k / num_experts`), non-embedding, embedding (input embedding and LM head) and total parameter counts at startup, and as `params/*` at step 0. ## Initialization -Weights start from a random init, unless the run resumes from a checkpoint. +Weights start from a random init, unless the run resumes from a checkpoint. The init lives with the models (`prime_rl.trainer.models.init`) and runs in the shared `setup_model` when it gets an `init` config. SFT and RL do not pass one, so their behavior is unchanged. | Field | Meaning | |---|---| @@ -59,7 +55,7 @@ With `optim.type = "muon"`, the hidden matrices and `o_proj`/`down_proj` are tra In MoE models, the expert matrices follow the hidden and `down_proj` rules per expert. The router maps the width to a fixed number of experts, so under Muon its learning rate falls as 1/m_N (with AdamW it follows the hidden rule). -Use QK-norm (`model.arch.use_qk_norm = true`) with CompleteP. Without it, CompleteP's attention collapses one or two LR steps above the optimum (the norm gains get no weight decay), and the cliff moves down with width. With QK-norm, CompleteP and the standard parametrization reach the same loss, and CompleteP's optimal LR stays fixed across widths (MoE QK-norm sweep). The per-role groups need AdamW or Muon and `model.offload = "none"`. The standard parametrization is CompleteP with the base shape equal to the model shape. Keep `head_dim` and the ratios `intermediate_size / hidden_size` (and `moe_intermediate_size / hidden_size`) and `num_attention_heads / num_key_value_heads` fixed across shapes. +Use QK-norm (`"use_qk_norm": true` in the model config) with CompleteP. Without it, CompleteP's attention collapses one or two LR steps above the optimum (the norm gains get no weight decay), and the cliff moves down with width. With QK-norm, CompleteP and the standard parametrization reach the same loss, and CompleteP's optimal LR stays fixed across widths (MoE QK-norm sweep). The per-role groups need AdamW or Muon and `model.offload = "none"`. The standard parametrization is CompleteP with the base shape equal to the model shape. Keep `head_dim` and the ratios `intermediate_size / hidden_size` (and `moe_intermediate_size / hidden_size`) and `num_attention_heads / num_key_value_heads` fixed across shapes. ## Learning-rate schedule diff --git a/examples/pretrain/glm4-moe-d512/config.json b/examples/pretrain/glm4-moe-d512/config.json new file mode 100644 index 00000000000..f8e1ad8cb96 --- /dev/null +++ b/examples/pretrain/glm4-moe-d512/config.json @@ -0,0 +1,37 @@ +{ + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 0, + "eos_token_id": 1, + "first_k_dense_replace": 1, + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 512, + "initializer_range": 0.02, + "intermediate_size": 2048, + "max_position_embeddings": 131072, + "model_type": "glm4_moe", + "moe_intermediate_size": 256, + "n_routed_experts": 64, + "n_shared_experts": 2, + "norm_topk_prob": true, + "num_attention_heads": 8, + "num_experts_per_tok": 6, + "num_hidden_layers": 8, + "num_key_value_heads": 2, + "pad_token_id": 1, + "partial_rotary_factor": 1.0, + "rms_norm_eps": 1e-05, + "rope_parameters": { + "partial_rotary_factor": 1.0, + "rope_theta": 500000.0, + "rope_type": "default" + }, + "rope_theta": 10000.0, + "routed_scaling_factor": 2.5, + "tie_word_embeddings": false, + "transformers_version": "5.6.2", + "use_cache": true, + "use_qk_norm": true, + "vocab_size": 129280 +} diff --git a/examples/pretrain/glm4_moe.toml b/examples/pretrain/glm4_moe.toml index f99c6849455..12154c82924 100644 --- a/examples/pretrain/glm4_moe.toml +++ b/examples/pretrain/glm4_moe.toml @@ -1,4 +1,4 @@ -# GLM-4.5-style MoE (64 routed experts, 6 active, 2 shared) with d = 512 and 8 layers, trained from scratch on +# GLM-4.5-style MoE (64 routed experts, 6 active, 2 shared) with d = 512 and 8 layers (glm4-moe-d512/config.json), trained from scratch on # the Dolma 3 mixture tokenized with the DeepSeek-V4 tokenizer. # 4096 steps x 256K tokens = 1.1B tokens. QK-norm and the learning rate follow the MoE QK-norm sweep (optimum 2^-7 # with QK-norm at d = 512). [parametrization] with this base shape keeps the LR when the width grows. @@ -8,29 +8,9 @@ max_steps = 4096 [model] +name = "examples/pretrain/glm4-moe-d512" # the architecture: config.json only, no weights ac = "None" -[model.arch] -model_type = "glm4_moe" -hidden_size = 512 -num_hidden_layers = 8 -num_attention_heads = 8 -num_key_value_heads = 2 -head_dim = 64 -intermediate_size = 2048 # the dense first layer -moe_intermediate_size = 256 -n_routed_experts = 64 -num_experts_per_tok = 6 -n_shared_experts = 2 -first_k_dense_replace = 1 -norm_topk_prob = true -routed_scaling_factor = 2.5 -attention_bias = false -use_qk_norm = true -partial_rotary_factor = 1.0 -rope_parameters = { rope_type = "default", rope_theta = 500000.0, partial_rotary_factor = 1.0 } -tie_word_embeddings = false - [tokenizer] name = "tokenizers/deepseek-v4-flash-base" # a local copy of deepseek-ai/DeepSeek-V4-Flash-Base without config.json diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 67d9beb7d2e..0d40f850918 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -2,7 +2,7 @@ import uuid import warnings from pathlib import Path -from typing import Annotated, Any, Literal, TypeAlias +from typing import Annotated, Literal, TypeAlias from pydantic import Field, model_validator @@ -21,38 +21,15 @@ AdamWConfig, CheckpointConfig, GCConfig, + InitConfig, ModelConfig, MoERuntimeConfig, MuonConfig, + ParametrizationConfig, TokenizerConfig, ) from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource -TOKENIZER_FIELDS = ("vocab_size", "bos_token_id", "eos_token_id", "pad_token_id") - - -class InitConfig(BaseConfig): - """Random weight initialization: every weight matrix and the embedding from N(0, std), scaled per role under - CompleteP (see ``[parametrization]``). Norm weights start at 1, biases at 0.""" - - std: float = Field(0.02, gt=0) - """Base standard deviation.""" - - seed: int = 0 - """Seed for the initialization. The values depend only on the seed, not on the number of GPUs.""" - - -class ParametrizationConfig(BaseConfig): - """CompleteP (muP extended to depth, Dey et al. 2025, alpha = 1) relative to a base shape: init std, learning - rate, Adam eps and weight decay scale per parameter role with ``m_N = hidden_size / base_hidden_size`` and - ``m_L = num_layers / base_num_layers``, so hyperparameters tuned on the base shape carry over.""" - - base_hidden_size: int = Field(ge=1) - """Hidden size of the base shape.""" - - base_num_layers: int = Field(ge=1) - """Number of layers of the base shape.""" - class MoEBalanceConfig(BaseConfig): """MoE load balancing as in DeepSeek-V3: an aux-loss-free selection bias plus a small sequence-wise balance @@ -76,9 +53,6 @@ class PretrainMoEConfig(MoERuntimeConfig): class PretrainModelConfig(ModelConfig): - arch: dict[str, Any] - """The architecture: ``model_type`` (a model in ``prime_rl.trainer.models``, e.g. ``glm4_moe``) and the fields of its config class. ``vocab_size`` and the special token ids come from the tokenizer. The launcher writes the config and the tokenizer to ``/configs/model``, which ``name`` points at.""" - init: InitConfig = InitConfig() """Random weight initialization. Used unless the run resumes from a checkpoint.""" @@ -87,14 +61,6 @@ class PretrainModelConfig(ModelConfig): offload: Literal["none"] = "none" """CPU offloading is not supported: pretraining's per-role parameter groups need ``none``.""" - @model_validator(mode="after") - def validate_arch(self): - if "model_type" not in self.arch: - raise ValueError("model.arch needs a model_type, e.g. 'glm4_moe'") - if derived := [key for key in TOKENIZER_FIELDS if key in self.arch]: - raise ValueError(f"model.arch must not set {derived}: they come from the tokenizer") - return self - class TokenSourceConfig(BaseConfig): path: Path @@ -290,13 +256,15 @@ class PretrainMuonConfig(MuonConfig): class PretrainConfig(BaseConfig): - model: PretrainModelConfig + model: PretrainModelConfig = PretrainModelConfig() + """The model. ``model.name`` is an HF model id or a local model directory with a ``config.json`` (as in SFT and + RL); its weights are not loaded: the model is initialized from scratch (``model.init``) or resumed.""" env_vars: EnvVars = {} """Extra environment variables for the trainer processes.""" - tokenizer: TokenizerConfig - """The tokenizer (``tokenizer.name`` is required). It sets the vocabulary and special tokens of the model.""" + tokenizer: TokenizerConfig = TokenizerConfig() + """The tokenizer. Defaults to the one in ``model.name``.""" data: PretrainDataConfig @@ -388,23 +356,17 @@ def validate_lr(self): @model_validator(mode="after") def auto_setup_run_identity(self): if self.run.name is None: - self.run.name = f"{self.model.arch['model_type']}--{uuid.uuid4().hex[:8]}" + self.run.name = f"{Path(self.model.name).name}--{uuid.uuid4().hex[:8]}" if self.run.dir is None: self.run.dir = self.run.name if self.monitors.wandb is not None and self.monitors.wandb.name is None: self.monitors.wandb.name = self.run.name return self - @model_validator(mode="after") - def auto_setup_model_dir(self): - """The model is trained from ``/configs/model``, written by the launcher from ``model.arch``.""" - self.model.name = (self.run_dir / "configs" / "model").as_posix() - return self - @model_validator(mode="after") def auto_setup_tokenizer(self): if self.tokenizer.name is None: - raise ValueError("tokenizer.name is required") + self.tokenizer.name = self.model.name if self.tokenizer.trust_remote_code is None: self.tokenizer.trust_remote_code = self.model.trust_remote_code return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index f7035c4e949..3c443fff14a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -230,6 +230,29 @@ class MoERuntimeConfig(BaseConfig): dispatch: MoEDispatchConfig = TorchMoEDispatchConfig() +class InitConfig(BaseConfig): + """Random weight initialization: every weight matrix and the embedding from N(0, std), scaled per role under + CompleteP (see ``[parametrization]``). Norm weights start at 1, biases at 0.""" + + std: float = Field(0.02, gt=0) + """Base standard deviation.""" + + seed: int = 0 + """Seed for the initialization. The values depend only on the seed, not on the number of GPUs.""" + + +class ParametrizationConfig(BaseConfig): + """CompleteP (muP extended to depth, Dey et al. 2025, alpha = 1) relative to a base shape: init std, learning + rate, Adam eps and weight decay scale per parameter role with ``m_N = hidden_size / base_hidden_size`` and + ``m_L = num_layers / base_num_layers``, so hyperparameters tuned on the base shape carry over.""" + + base_hidden_size: int = Field(ge=1) + """Hidden size of the base shape.""" + + base_num_layers: int = Field(ge=1) + """Number of layers of the base shape.""" + + class ModelConfig(BaseModelConfig): conversion_dir: Path | None = None """Directory for the auto-converted weights (written to a `prime`/`hf` subdirectory). If not set, we write into the model snapshot directory.""" diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py index 3f55b9fe73b..00d5a80ade0 100644 --- a/src/prime_rl/entrypoints/pretrain.py +++ b/src/prime_rl/entrypoints/pretrain.py @@ -27,34 +27,6 @@ PRETRAIN_SBATCH = "pretrain.sbatch" -def write_model_dir(config: PretrainConfig) -> None: - """Write the config built from ``[model.arch]`` and the tokenizer to ``/configs/model``. The vocabulary size - (rounded up to a multiple of 128) and the special token ids come from the tokenizer.""" - from transformers import AutoTokenizer - from transformers.models.auto.configuration_auto import CONFIG_MAPPING - - import prime_rl.trainer.models # noqa: F401 - registers the custom model configs with AutoConfig - - arch = dict(config.model.arch) - config_class = CONFIG_MAPPING[arch.pop("model_type")] - # PretrainedConfig keeps unknown kwargs as attributes, so a typo would silently do nothing. - if unknown := arch.keys() - config_class().to_dict().keys(): - raise ValueError(f"[model.arch] fields {sorted(unknown)} do not exist on {config_class.__name__}") - tokenizer = AutoTokenizer.from_pretrained( - config.tokenizer.name, trust_remote_code=config.tokenizer.trust_remote_code - ) - hf_config = config_class( - **arch, - vocab_size=-(-len(tokenizer) // 128) * 128, - bos_token_id=tokenizer.bos_token_id, - eos_token_id=tokenizer.eos_token_id, - pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id, - ) - model_dir = Path(config.model.name) - hf_config.save_pretrained(model_dir) - tokenizer.save_pretrained(model_dir) - - def pretrain_slurm(config: PretrainConfig, config_path: Path, log_dir: Path) -> None: logger = get_logger() script_path = get_launcher_dir(config.run_dir) / PRETRAIN_SBATCH @@ -90,7 +62,9 @@ def pretrain(config: PretrainConfig): write_launch_artifacts(config_dir, "pretrain") if not config.dry_run: - write_model_dir(config) + from prime_rl.trainer.model import pre_download_model + + pre_download_model(config.model.name, skip_weights=True) config_path = config_dir / PRETRAIN_CONFIG exclude = {"dry_run", "clean"} diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 313d0e9a3a6..59db2584d2a 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -176,7 +176,7 @@ async def init(self, config: BaseConfig | None = None, output_dir: Path | None = # Pretraining trains on token data from scratch: no environments, the model is its architecture. init_kwargs = dict( name=self.config.name, - model=config.model.arch["model_type"], + model=config.model.name, environments=[], training=pr.TrainingSpec( max_steps=config.max_steps, diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 4ff32e3bb5d..3c1a8fe53b0 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -31,8 +31,10 @@ ActivationCheckpointConfig, CompileConfig, FP8Config, + InitConfig, ModelConfig, MXFP8Config, + ParametrizationConfig, TokenizerConfig, ) from prime_rl.multimodal import ForwardPolicy @@ -54,6 +56,7 @@ write_back_loaded_packed_parameters, ) from prime_rl.trainer.models.glm_moe_dsa.sparse_mla_attention import Indexer +from prime_rl.trainer.models.init import init_weights, multipliers from prime_rl.trainer.models.layers.fp8_linear import replace_linear_with_fp8_blockwise_linear from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head from prime_rl.trainer.models.layers.moe import MoE, TokenChoiceTopKRouter @@ -942,7 +945,12 @@ def setup_model( config: ModelConfig, parallel_dims: ParallelDims, loading_from_checkpoint_later: bool = False, + init: InitConfig | None = None, + parametrization: ParametrizationConfig | None = None, ) -> nn.Module: + """Build the sharded model. Its weights come from ``config.name``, or later from a checkpoint + (``loading_from_checkpoint_later``); with ``init``, they are initialized from scratch (with the CompleteP + multipliers of ``parametrization``) instead.""" resolve_auto_attn(config) if config.attn == "flash_attention_3" and not is_flash_attn_3_available(): @@ -1009,9 +1017,11 @@ def setup_model( setup_fsdp(model, config, parallel_dims) - if loading_from_checkpoint_later: + if loading_from_checkpoint_later or init is not None: logger.warning( "Skipping loading weights. Initializing an empty model on device, loading from checkpoint later." + if init is None + else "Initializing the weights from scratch." ) device = "cuda" model.to_empty(device=device) @@ -1021,6 +1031,8 @@ def setup_model( load_dcp_from_hf(model, config, parallel_dims) _reset_runtime_moe_buffers(model) + if init is not None: + init_weights(model, init, *multipliers(parametrization, model)) return model diff --git a/src/prime_rl/trainer/init.py b/src/prime_rl/trainer/models/init.py similarity index 86% rename from src/prime_rl/trainer/init.py rename to src/prime_rl/trainer/models/init.py index 0f0fc6a5e87..b64323e0e58 100644 --- a/src/prime_rl/trainer/init.py +++ b/src/prime_rl/trainer/models/init.py @@ -7,8 +7,8 @@ from torch.distributed.tensor import DTensor from torch.distributed.tensor._utils import compute_local_shape_and_global_offset -from prime_rl.configs.pretrain import InitConfig -from prime_rl.trainer.model import iter_moe_blocks +from prime_rl.configs.trainer import InitConfig, ParametrizationConfig +from prime_rl.trainer.models.layers.moe import TokenChoiceTopKRouter from prime_rl.utils.weights import resolve_fqn ParamRole = Literal["embedding", "output", "residual_out", "hidden", "norm", "bias"] @@ -44,6 +44,14 @@ def role_std(role: ParamRole, std: float, m_n: float, m_l: float) -> float: return std * scale[role] +def multipliers(config: ParametrizationConfig | None, model: nn.Module) -> tuple[float, float]: + """Width and depth multipliers ``(m_N, m_L)`` of the model relative to the CompleteP base shape; 1 and 1 in + the standard parametrization.""" + if config is None: + return 1.0, 1.0 + return model.config.hidden_size / config.base_hidden_size, model.config.num_hidden_layers / config.base_num_layers + + def block_seed(seed: int, name: str, block: int) -> int: digest = hashlib.sha256(f"{seed}:{name}:{block}".encode()).digest() return int.from_bytes(digest[:8], "little") & ((1 << 63) - 1) @@ -103,6 +111,6 @@ def init_weights(model: nn.Module, config: InitConfig, m_n: float = 1.0, m_l: fl ) local.copy_(values.mul_(role_std(role, config.std, m_n, m_l))) # MoE selection biases are buffers: they start at 0, so routing starts from the router scores alone. - for moe in iter_moe_blocks(model): - if moe.router.selection_bias is not None: - moe.router.selection_bias.zero_() + for module in model.modules(): + if isinstance(module, TokenChoiceTopKRouter) and module.selection_bias is not None: + module.selection_bias.zero_() diff --git a/src/prime_rl/trainer/pretrain/parametrization.py b/src/prime_rl/trainer/pretrain/parametrization.py index 37357f08403..4159f1f953a 100644 --- a/src/prime_rl/trainer/pretrain/parametrization.py +++ b/src/prime_rl/trainer/pretrain/parametrization.py @@ -2,19 +2,11 @@ from torch import nn -from prime_rl.configs.pretrain import ParametrizationConfig, PretrainAdamWConfig, PretrainMuonConfig -from prime_rl.trainer.init import param_role +from prime_rl.configs.pretrain import PretrainAdamWConfig, PretrainMuonConfig +from prime_rl.trainer.models.init import param_role from prime_rl.utils.weights import resolve_fqn -def multipliers(config: ParametrizationConfig | None, model: nn.Module) -> tuple[float, float]: - """Width and depth multipliers ``(m_N, m_L)`` of the model relative to the CompleteP base shape; 1 and 1 in - the standard parametrization.""" - if config is None: - return 1.0, 1.0 - return model.config.hidden_size / config.base_hidden_size, model.config.num_hidden_layers / config.base_num_layers - - def param_groups( model: nn.Module, optim: PretrainAdamWConfig | PretrainMuonConfig, m_n: float, m_l: float, ep_enabled: bool ) -> list[dict]: diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 45fd1f8999f..9205b678cee 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -32,8 +32,8 @@ from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter -from prime_rl.trainer.init import init_weights, param_role -from prime_rl.trainer.pretrain.parametrization import multipliers, param_groups +from prime_rl.trainer.models.init import multipliers, param_role +from prime_rl.trainer.pretrain.parametrization import param_groups from prime_rl.trainer.pretrain.bytes import token_byte_lengths from prime_rl.trainer.pretrain.data import ( TokenShards, @@ -145,10 +145,14 @@ def train(config: PretrainConfig): logger.info(f"Initializing model ({config.model})") # Weights come from the checkpoint when resuming and from the random init otherwise. - model = setup_model(config.model, parallel_dims, loading_from_checkpoint_later=True) + model = setup_model( + config.model, + parallel_dims, + loading_from_checkpoint_later=checkpoint_step is not None, + init=config.model.init if checkpoint_step is None else None, + parametrization=config.parametrization, + ) m_n, m_l = multipliers(config.parametrization, model) - if checkpoint_step is None: - init_weights(model, config.model.init, m_n, m_l) if parallel_dims.cp_enabled: setup_context_parallel(model, config.model, parallel_dims) num_params = count_parameters(model) diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index 61e85997be3..b3021a39c62 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -5,9 +5,10 @@ from torch.optim import AdamW from transformers import LlamaConfig, LlamaForCausalLM -from prime_rl.configs.pretrain import InitConfig, PretrainAdamWConfig, PretrainMuonConfig -from prime_rl.trainer.init import init_weights, normal_shard +from prime_rl.configs.pretrain import PretrainAdamWConfig, PretrainMuonConfig +from prime_rl.configs.trainer import InitConfig from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM +from prime_rl.trainer.models.init import init_weights, normal_shard from prime_rl.trainer.pretrain.parametrization import param_groups diff --git a/tests/unit/train/pretrain/test_launcher.py b/tests/unit/train/pretrain/test_launcher.py index 9238d16d42e..e712e4f4549 100644 --- a/tests/unit/train/pretrain/test_launcher.py +++ b/tests/unit/train/pretrain/test_launcher.py @@ -11,7 +11,7 @@ def test_multi_node_dry_run_writes_trainer_sbatch(tmp_path, monkeypatch): config = PretrainConfig.model_validate( { "output_dir": tmp_path, - "model": {"arch": {"model_type": "llama"}}, + "model": {"name": "configs/debug/fake/pretrain-model"}, "tokenizer": {"name": "PrimeIntellect/Qwen3-0.6B"}, "data": {"type": "fake", "seq_len": 8, "batch_tokens": 16, "micro_batch_tokens": 8}, "optim": {"lr": 1e-3}, diff --git a/tests/unit/train/pretrain/test_model_dir.py b/tests/unit/train/pretrain/test_model_dir.py deleted file mode 100644 index 7124c0f3ec2..00000000000 --- a/tests/unit/train/pretrain/test_model_dir.py +++ /dev/null @@ -1,34 +0,0 @@ -import pytest -from transformers import AutoConfig, AutoTokenizer - -from prime_rl.configs.pretrain import PretrainConfig -from prime_rl.entrypoints.pretrain import write_model_dir - -TOKENIZER = "PrimeIntellect/Qwen3-0.6B" - - -def make_config(tmp_path, **arch) -> PretrainConfig: - return PretrainConfig( - model={"arch": {"model_type": "glm4_moe", "hidden_size": 64, "num_hidden_layers": 2, **arch}}, - tokenizer={"name": TOKENIZER}, - data={"type": "fake", "seq_len": 8, "batch_tokens": 8, "micro_batch_tokens": 8}, - optim={"type": "muon", "lr": 1e-3}, - max_steps=1, - output_dir=tmp_path, - ) - - -def test_model_dir_takes_vocabulary_and_special_tokens_from_tokenizer(tmp_path): - config = make_config(tmp_path) - write_model_dir(config) - hf_config = AutoConfig.from_pretrained(config.model.name) - tokenizer = AutoTokenizer.from_pretrained(TOKENIZER) - assert hf_config.model_type == "glm4_moe" - assert hf_config.hidden_size == 64 - assert hf_config.vocab_size == -(-len(tokenizer) // 128) * 128 - assert hf_config.eos_token_id == tokenizer.eos_token_id - - -def test_model_dir_rejects_unknown_arch_fields(tmp_path): - with pytest.raises(ValueError, match="hiden_size"): - write_model_dir(make_config(tmp_path, hiden_size=64)) diff --git a/tests/unit/train/pretrain/test_schedules.py b/tests/unit/train/pretrain/test_schedules.py index e02f526fa74..9a57a13d232 100644 --- a/tests/unit/train/pretrain/test_schedules.py +++ b/tests/unit/train/pretrain/test_schedules.py @@ -31,7 +31,7 @@ def test_wsd_endpoints_and_shapes(): def test_schedule_phases_must_fit_max_steps(): config = PretrainConfig( - model={"arch": {"model_type": "llama"}}, + model={"name": "configs/debug/fake/pretrain-model"}, tokenizer={"name": "t"}, data={"type": "fake", "seq_len": 8, "batch_tokens": 8, "micro_batch_tokens": 8}, optim={"lr": 1e-3}, From 4b6c2f09023c92a820561c58e53f9089881edc5d Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 19:37:36 +0200 Subject: [PATCH 072/132] perf(pretrain): sync and reshard gradients once per step, not per micro-batch --- src/prime_rl/trainer/pretrain/step.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/pretrain/step.py b/src/prime_rl/trainer/pretrain/step.py index a9c306281b4..f17921de2f4 100644 --- a/src/prime_rl/trainer/pretrain/step.py +++ b/src/prime_rl/trainer/pretrain/step.py @@ -63,7 +63,7 @@ def run_micro_batches( sums = {"loss": zero.clone(), "loss_tokens": zero.clone(), "nan_count": zero.clone(), "tokens": zero.clone()} maxes = {} data_time = 0.0 - for _ in range(grad_accum_steps): + for micro_step in range(grad_accum_steps): data_start_time = time.perf_counter() micro_batch = next(dataiter) data_time += time.perf_counter() - data_start_time @@ -75,6 +75,10 @@ def run_micro_batches( sums["nan_count"] += is_nan sums["loss"] += torch.where(is_nan, 0.0, micro_loss_sum.detach()) scaled_loss = torch.where(is_nan, torch.nan_to_num(micro_loss_sum, nan=0.0), micro_loss_sum + aux_loss_sum) + is_last_micro_step = micro_step == grad_accum_steps - 1 + # Reduce-scatter (and reshard) once per step, not per micro-batch: the expert gradients cross nodes. + model.set_requires_gradient_sync(is_last_micro_step) + model.set_reshard_after_backward(is_last_micro_step) (scaled_loss / grad_accum_steps).backward() if is_moe_model: From 31d947310dad214b008fb447ff24d11c7ba238eb Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 20:15:43 +0200 Subject: [PATCH 073/132] fix(trainer): dense pretrain models with the default MoE runtime PretrainConfig's model.moe carries the balancing settings, so comparing the whole config with MoERuntimeConfig() rejected every dense pretrain model. Compare the runtime fields (compute, dispatch) only; SFT/RL use MoERuntimeConfig itself, so their check is unchanged. --- src/prime_rl/trainer/moe_runtime.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/moe_runtime.py b/src/prime_rl/trainer/moe_runtime.py index fa1dc2bdc35..dcbd60994d7 100644 --- a/src/prime_rl/trainer/moe_runtime.py +++ b/src/prime_rl/trainer/moe_runtime.py @@ -62,7 +62,8 @@ def _resolve_expert_compute(config: ModelConfig) -> ExpertCompute: def configure_moe_runtime(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims) -> None: moe_layers = [module for module in model.modules() if isinstance(module, MoE)] if not moe_layers: - if config.moe != MoERuntimeConfig(): + default = MoERuntimeConfig() + if (config.moe.compute, config.moe.dispatch) != (default.compute, default.dispatch): raise ValueError("A non-default model.moe runtime was configured, but the model has no custom MoE layers.") return From e3e4feb407e428b3ec0a2772dd60c7930a8e6d91 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 21:05:57 +0200 Subject: [PATCH 074/132] fix(pretrain): drop the stale top-level rope_theta from the glm4_moe example config The model reads rope_parameters (rope_theta 500000); the top-level 10000 was a default the config class serializes alongside it. --- examples/pretrain/glm4-moe-d512/config.json | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/pretrain/glm4-moe-d512/config.json b/examples/pretrain/glm4-moe-d512/config.json index f8e1ad8cb96..f821da8e747 100644 --- a/examples/pretrain/glm4-moe-d512/config.json +++ b/examples/pretrain/glm4-moe-d512/config.json @@ -27,7 +27,6 @@ "rope_theta": 500000.0, "rope_type": "default" }, - "rope_theta": 10000.0, "routed_scaling_factor": 2.5, "tie_word_embeddings": false, "transformers_version": "5.6.2", From 4284fe8857a8292933256ecf3cc5215d39438fab Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Oct 2026 21:22:41 +0200 Subject: [PATCH 075/132] refactor(pretrain): use the shared linear scheduler; model.seq_len follows data.seq_len Pretrain's own [schedule] mostly duplicated the shared scheduler. The shared linear scheduler gains warmup_steps/decay_steps as fractions of max_steps (rounded down) and decay_shape (linear, cosine, 1-sqrt); its defaults and SFT/RL behavior are unchanged. Pretrain uses [scheduler] (default: linear, warmup 0.1, decay 0.9 to 0) and checkpoints the scheduler state like SFT. Compared with the old [schedule], the first update uses 1e-8 * lr and each phase boundary moves one update later, as in the schedule experiments, which ran on the shared scheduler. Checkpoints from before this change hold no scheduler state: resume them with ckpt.skip_scheduler = true. model.seq_len, unused by pretrain, now equals data.seq_len in the resolved config. --- docs/pretraining.md | 16 ++--- .../src/prime_rl/configs/pretrain.py | 46 ++++---------- .../src/prime_rl/configs/trainer.py | 28 +++++++-- src/prime_rl/trainer/pretrain/schedules.py | 61 ------------------- src/prime_rl/trainer/pretrain/train.py | 22 ++++--- src/prime_rl/trainer/scheduler.py | 37 +++++++++-- tests/unit/train/pretrain/test_schedules.py | 56 ----------------- tests/unit/train/test_scheduler.py | 28 +++++++++ 8 files changed, 112 insertions(+), 182 deletions(-) delete mode 100644 src/prime_rl/trainer/pretrain/schedules.py delete mode 100644 tests/unit/train/pretrain/test_schedules.py create mode 100644 tests/unit/train/test_scheduler.py diff --git a/docs/pretraining.md b/docs/pretraining.md index 1d144fb080a..1c2ff2c3952 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -63,17 +63,17 @@ Use QK-norm (`model.arch.use_qk_norm = true`) with CompleteP. Without it, Comple ## Learning-rate schedule -`[schedule]` sets a multiplier `f(step)` on the learning rate, where `step` counts the optimizer updates completed so far (0 for the first update). Every param group uses its own learning rate times `f(step)`, so ratios between groups hold at every step. The default warms up over 10% of the steps and decays linearly to 0 over the remaining 90%. +Pretraining uses the shared `[scheduler]`. Its default here is the `linear` scheduler with a warmup over 10% of the steps and a linear decay to `min_lr = 0` over the remaining 90%. Warmup and decay take a number of steps or a fraction of `max_steps` (rounded down), and `decay_shape` is `linear`, `cosine` or `1-sqrt`: ```toml -[schedule] -warmup = 0.05 # fraction of max_steps, or a number of steps -decay = 0.2 -end_ratio = 0.1 -decay_shape = "1-sqrt" # or "linear" (default), "cosine" +[scheduler] +type = "linear" +warmup_steps = 0.05 +decay_steps = 0.2 +decay_shape = "1-sqrt" ``` -`f(0) = 0`, the warmup reaches 1 at its last update (`warmup - 1`; a warmup of 0 or 1 starts at 1), and the decay reaches `end_ratio` at the last update (`max_steps - 1`). AdamW and Muon apply weight decay as `lr * weight_decay`, so the schedule also scales the decay. The multiplier is recomputed from the step, so checkpoints hold no schedule state. Resuming with `ckpt.skip_scheduler = true` restarts the schedule at the resumed step, spanning the remaining steps. The multiplier is logged as `schedule/lr`. +The schedule scales every param group's own learning rate, so ratios between groups hold. AdamW and Muon apply weight decay as `lr * weight_decay`, so the schedule also scales the decay. The scheduler state is checkpointed; resuming with `ckpt.skip_scheduler = true` rebuilds the schedule over the remaining steps. ## MoE load balancing @@ -117,7 +117,7 @@ Batch sizes are in tokens, so changing `seq_len` keeps the tokens per step and t ## Monitoring -Every step logs the loss, learning rate and schedule multiplier, gradient norms (total and per parameter group, `optim/grad_norm/`), loss and grad-norm spike flags, per-source token counts and epochs (`progress//*`), throughput, MFU, data wait and peak memory (max and mean over ranks). MoE models add the expert balance of the step's global counts (`expert_load/*`, also per layer), `router/*` and `loss/seq_balance`. The step statistics are reduced once per step. +Every step logs the loss, learning rate, gradient norms (total and per parameter group, `optim/grad_norm/`), loss and grad-norm spike flags, per-source token counts and epochs (`progress//*`), throughput, MFU, data wait and peak memory (max and mean over ranks). MoE models add the expert balance of the step's global counts (`expert_load/*`, also per layer), `router/*` and `loss/seq_balance`. The step statistics are reduced once per step. `[diagnostics]` (default: every 100 steps; `diagnostics = null` turns it off) adds the size of the optimizer update relative to the weights per parameter group (`update/rms_ratio/*`, `param/rms/*`) and a forward pass on a fixed probe batch: attention max logit and entropy per layer (`attn/*`, FlashAttention layers) and the residual-stream size per layer (`act/*`). It does not support context parallelism yet. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 9fafb5cb757..dec2743aa0c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -21,10 +21,13 @@ AdamWConfig, CheckpointConfig, GCConfig, + LinearSchedulerConfig, ModelConfig, MoERuntimeConfig, MuonConfig, + SchedulerConfig, TokenizerConfig, + validate_scheduler, ) from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource @@ -201,34 +204,6 @@ class ValConfig(BaseConfig): windows: as many as fill whole micro-batches on every data-parallel rank.""" -StepsOrFraction: TypeAlias = Annotated[int, Field(ge=0)] | Annotated[float, Field(gt=0, lt=1)] - - -def resolve_steps(value: int | float, total_steps: int) -> int: - """Steps as given, or a fraction of ``total_steps`` rounded down, so fractions that sum to at most 1 fit.""" - return math.floor(value * total_steps) if isinstance(value, float) else value - - -class ScheduleConfig(BaseConfig): - """Learning-rate multiplier ``f(step)`` (warmup, stable at 1, decay), where ``step`` counts the updates - completed before the current one. Every param group uses its own learning rate times ``f(step)``, so - ratios between groups hold. ``f(0) = 0``, the warmup reaches 1 at its last update (``warmup - 1``; a - warmup of 0 or 1 starts at 1), and the decay reaches ``end_ratio`` at the last update of the run.""" - - warmup: StepsOrFraction = 0.1 - """Warmup updates, or a fraction of ``max_steps``.""" - - decay: StepsOrFraction = 0.9 - """Decay updates at the end of the run, or a fraction of ``max_steps``. The default decays over all - updates after the default warmup.""" - - end_ratio: float = Field(0.0, ge=0) - """Multiplier of the last update.""" - - decay_shape: Literal["linear", "cosine", "1-sqrt"] = "linear" - """Shape of the decay from 1 to ``end_ratio``.""" - - class DiagnosticsConfig(BaseConfig): """Diagnostic step: on every ``interval``-th step, the optimizer update's size per parameter group and a forward pass on a fixed probe batch (the first ``probe_tokens`` tokens of each rank's first micro-batch) @@ -312,7 +287,7 @@ class PretrainConfig(BaseConfig): optim: PretrainOptimizerConfig = PretrainMuonConfig() """Optimizer. ``optim.lr`` is required.""" - schedule: ScheduleConfig = ScheduleConfig() + scheduler: SchedulerConfig = LinearSchedulerConfig(warmup_steps=0.1, decay_steps=0.9) """Learning-rate schedule. Defaults to a warmup over 10% of the steps and a linear decay to 0 over the rest.""" ckpt: CheckpointConfig | None = CheckpointConfig() @@ -458,10 +433,11 @@ def validate_cp(self): return self @model_validator(mode="after") - def validate_schedule(self): - phases = resolve_steps(self.schedule.warmup, self.max_steps) + resolve_steps( - self.schedule.decay, self.max_steps - ) - if phases > self.max_steps: - raise ValueError(f"schedule.warmup + schedule.decay ({phases}) exceed max_steps ({self.max_steps})") + def validate_scheduler_steps(self): + validate_scheduler(self.scheduler, self.max_steps) + return self + + @model_validator(mode="after") + def auto_setup_model_seq_len(self): + self.model.seq_len = self.data.seq_len return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 3aa56dc3fae..9ec00ff55d5 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -1,3 +1,4 @@ +import math import re import warnings from pathlib import Path @@ -356,14 +357,25 @@ class ConstantSchedulerConfig(BaseConfig): type: Literal["constant"] = "constant" +StepsOrFraction: TypeAlias = Annotated[int, Field(ge=0)] | Annotated[float, Field(gt=0, lt=1)] + + +def resolve_steps(value: int | float, max_steps: int | None) -> int: + """Steps as given, or a fraction of ``max_steps`` rounded down.""" + return math.floor(value * max_steps) if isinstance(value, float) else value + + class LinearSchedulerConfig(BaseConfig): type: Literal["linear"] = "linear" - warmup_steps: int = Field(10, ge=0) - """Warmup steps for the learning rate scheduler.""" + warmup_steps: StepsOrFraction = 10 + """Warmup steps for the learning rate scheduler, or a fraction of ``max_steps``.""" + + decay_steps: StepsOrFraction = 10 + """Steps to decay the learning rate during the final portion of training, or a fraction of ``max_steps``.""" - decay_steps: int = Field(10, ge=0) - """Steps to decay the learning rate during the final portion of training.""" + decay_shape: Literal["linear", "cosine", "1-sqrt"] = "linear" + """Shape of the decay from the peak to ``min_lr``.""" min_lr: float = Field(0.0, ge=0) """Minimum learning rate to converge to.""" @@ -394,11 +406,15 @@ def validate_scheduler(scheduler: SchedulerConfig, max_steps: int | None) -> Non if scheduler.decay_steps > 0: if max_steps is None: raise ValueError("Must specify max_steps when using a linear scheduler with decay_steps > 0") - if scheduler.warmup_steps + scheduler.decay_steps > max_steps: + warmup_steps = resolve_steps(scheduler.warmup_steps, max_steps) + decay_steps = resolve_steps(scheduler.decay_steps, max_steps) + if warmup_steps + decay_steps > max_steps: raise ValueError( - f"warmup_steps ({scheduler.warmup_steps}) + decay_steps ({scheduler.decay_steps}) " + f"warmup_steps ({warmup_steps}) + decay_steps ({decay_steps}) " f"must not exceed max_steps ({max_steps})" ) + elif isinstance(scheduler.warmup_steps, float) and max_steps is None: + raise ValueError("Must specify max_steps when warmup_steps is a fraction") if isinstance(scheduler, CosineSchedulerConfig): if max_steps is None: raise ValueError("Must specify max_steps when using a cosine scheduler") diff --git a/src/prime_rl/trainer/pretrain/schedules.py b/src/prime_rl/trainer/pretrain/schedules.py deleted file mode 100644 index 2f8f0cba068..00000000000 --- a/src/prime_rl/trainer/pretrain/schedules.py +++ /dev/null @@ -1,61 +0,0 @@ -import math -from typing import Callable - -from prime_rl.configs.pretrain import ScheduleConfig, resolve_steps -from prime_rl.trainer.optim import OptimizerLike - -DECAY_SHAPES: dict[str, Callable[[float], float]] = { - "linear": lambda progress: 1 - progress, - "cosine": lambda progress: 0.5 * (1 + math.cos(math.pi * progress)), - "1-sqrt": lambda progress: 1 - math.sqrt(progress), -} - - -def wsd(step: int, total_steps: int, warmup: int, decay: int, end_ratio: float, decay_shape: str) -> float: - """Linear warmup from 0 at step 0 to 1 at step ``warmup - 1``, stable at 1, then ``decay`` updates - that reach ``end_ratio`` at the last step ``total_steps - 1``.""" - peak = max(warmup - 1, 0) - if step < peak: - return step / peak - if decay == 0: - return 1.0 - last_stable = max(total_steps - 1 - decay, peak) - if step <= last_stable: - return 1.0 - progress = (step - last_stable) / (total_steps - 1 - last_stable) - return end_ratio + (1 - end_ratio) * DECAY_SHAPES[decay_shape](progress) - - -class LRSchedule: - """Sets the learning rate of each update. It is a pure function of the step, so nothing needs to be - checkpointed: on resume it is recomputed from the step. - - The base learning rates are read from the optimizer's param groups at construction, before any - checkpoint load. Each param group keeps its own base, so ratios between groups hold at every step. - """ - - def __init__(self, config: ScheduleConfig, optimizer: OptimizerLike, total_steps: int, step_offset: int = 0): - self.config = config - self.optimizer = optimizer - self.total_steps = total_steps - self.step_offset = step_offset - # Bases by group index: loading an optimizer state dict replaces the group dicts. - self.bases = [group["lr"] for group in optimizer.param_groups] - - def multiplier(self, step: int) -> float: - config = self.config - return wsd( - step - self.step_offset, - self.total_steps, - resolve_steps(config.warmup, self.total_steps), - resolve_steps(config.decay, self.total_steps), - config.end_ratio, - config.decay_shape, - ) - - def apply(self, step: int) -> float: - """Write the learning rates of update ``step`` (updates completed so far) and return the multiplier.""" - multiplier = self.multiplier(step) - for group, base in zip(self.optimizer.param_groups, self.bases): - group["lr"] = base * multiplier - return multiplier diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 45fd1f8999f..70cfef974ac 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -43,7 +43,6 @@ source_tokens, setup_val_dataset, ) -from prime_rl.trainer.pretrain.schedules import LRSchedule from prime_rl.trainer.pretrain.diagnostics import UpdateTracker, probe_batch, run_probe from prime_rl.trainer.pretrain.suite import ( length_label, @@ -62,6 +61,7 @@ take_routing_counts, ) from prime_rl.trainer.sft.data import get_dataset_state +from prime_rl.trainer.scheduler import setup_scheduler from prime_rl.trainer.utils import ( get_dcp_group, GarbageCollection, @@ -170,10 +170,10 @@ def train(config: PretrainConfig): param_groups=param_groups(model, config.optim, m_n, m_l, parallel_dims.ep_enabled), ) - # With ckpt.skip_scheduler, the schedule restarts at the resumed step and spans the remaining steps. - restart_schedule = config.ckpt is not None and config.ckpt.skip_scheduler and checkpoint_step is not None - schedule_offset = checkpoint_step if restart_schedule else 0 - lr_schedule = LRSchedule(config.schedule, optimizer, config.max_steps - schedule_offset, schedule_offset) + # With ckpt.skip_scheduler, the schedule is rebuilt over the remaining steps. + skip_scheduler = (config.ckpt or CheckpointConfig()).skip_scheduler + scheduler_steps = config.max_steps - checkpoint_step if skip_scheduler and checkpoint_step else config.max_steps + scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) train_dataset = setup_train_dataset( config.data, config.max_steps, len(tokenizer), tokenizer.eos_token_id, non_dp_size=config.model.cp @@ -201,12 +201,15 @@ def train(config: PretrainConfig): checkpoint_step, model, [optimizer], - None, + None if skip_scheduler else scheduler, progress, dataloader=dataloader if restore_dataloader else None, path=config.resume.dir / "trainer" if config.resume.dir is not None else None, ) progress.step += 1 + if skip_scheduler: + # Loading the optimizer state replaces the param groups the scheduler was built on. + scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) logger.info(f"Resuming from step {checkpoint_step} (dataset_state={get_dataset_state(dataloader)})") # Step s trains on the global windows [(s - 1) * batch_size, s * batch_size). The tokens per source are counted @@ -361,7 +364,6 @@ def run_suite(step: int) -> None: gc_handler.run(progress.step) is_last_step = progress.step >= config.max_steps step_start_time = time.perf_counter() - lr_multiplier = lr_schedule.apply(progress.step - 1) seq_balance_sums.clear() result = run_micro_batches( @@ -403,6 +405,7 @@ def run_suite(step: int) -> None: optimizer.step() optimizer.zero_grad() current_lr = optimizer.param_groups[0]["lr"] + scheduler.step() diagnostic_metrics = update_tracker.after_step(model) if is_diagnostic_step else {} # Peak memory of the step's forward, backward and optimizer passes, over all ranks. local_peak_memory = torch.tensor(torch.cuda.max_memory_reserved() / 1024**3, device="cuda") @@ -417,7 +420,7 @@ def run_suite(step: int) -> None: disk_metrics = {} if config.ckpt and config.ckpt.interval and progress.step % config.ckpt.interval == 0 and not is_last_step: save_ckpt_start_time = time.perf_counter() - ckpt_manager.save(progress.step, model, [optimizer], None, progress, dataloader=dataloader) + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) ckpt_manager.maybe_clean() save_ckpt_time = time.perf_counter() - save_ckpt_start_time disk_metrics = get_ckpt_disk_metrics(config.run_dir) @@ -481,7 +484,6 @@ def run_suite(step: int) -> None: **result.metrics, **moe_metrics, **diagnostic_metrics, - "schedule/lr": lr_multiplier, "step": progress.step, } if grad_norm is not None: @@ -497,7 +499,7 @@ def run_suite(step: int) -> None: if config.ckpt is not None: logger.info(f"Saving final checkpoint at step {progress.step}") - ckpt_manager.save(progress.step, model, [optimizer], None, progress, dataloader=dataloader) + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) ckpt_manager.maybe_clean() logger.success("Pretrain trainer finished") asyncio.run(monitors.finalize()) diff --git a/src/prime_rl/trainer/scheduler.py b/src/prime_rl/trainer/scheduler.py index 25bfb93f6c3..7fb18fab87f 100644 --- a/src/prime_rl/trainer/scheduler.py +++ b/src/prime_rl/trainer/scheduler.py @@ -1,9 +1,12 @@ from __future__ import annotations +import math +from collections.abc import Callable + from torch.optim import Optimizer -from torch.optim.lr_scheduler import ConstantLR, CosineAnnealingLR, LinearLR, LRScheduler, SequentialLR +from torch.optim.lr_scheduler import ConstantLR, CosineAnnealingLR, LambdaLR, LinearLR, LRScheduler, SequentialLR -from prime_rl.configs.trainer import SchedulerConfig +from prime_rl.configs.trainer import SchedulerConfig, resolve_steps from prime_rl.trainer.optim import OffloadOptimizer, OptimizerLike @@ -19,8 +22,20 @@ def setup_constant_scheduler(optimizer: Optimizer) -> LRScheduler: return ConstantLR(optimizer, factor=1.0) +DECAY_SHAPES: dict[str, Callable[[float], float]] = { + "cosine": lambda progress: 0.5 * (1 + math.cos(math.pi * progress)), + "1-sqrt": lambda progress: 1 - math.sqrt(progress), +} + + def setup_linear_scheduler( - optimizer: Optimizer, max_steps: int | None, warmup_steps: int, decay_steps: int, lr: float, min_lr: float + optimizer: Optimizer, + max_steps: int | None, + warmup_steps: int, + decay_steps: int, + lr: float, + min_lr: float, + decay_shape: str = "linear", ) -> LRScheduler: """Create a linear (WSD) learning rate scheduler.""" # Create schedulers for each phase @@ -46,7 +61,16 @@ def setup_linear_scheduler( ) constant_steps = decay_start_step - warmup_steps constant_scheduler = LinearLR(optimizer, start_factor=1.0, end_factor=1.0, total_iters=constant_steps) - decay_scheduler = LinearLR(optimizer, start_factor=1.0, end_factor=min_lr_factor, total_iters=decay_steps - 1) + if decay_shape == "linear": + decay_scheduler = LinearLR( + optimizer, start_factor=1.0, end_factor=min_lr_factor, total_iters=decay_steps - 1 + ) + else: + shape = DECAY_SHAPES[decay_shape] + decay_scheduler = LambdaLR( + optimizer, + lambda step: min_lr_factor + (1 - min_lr_factor) * shape(min(step / max(decay_steps - 1, 1), 1.0)), + ) schedulers.append(constant_scheduler) schedulers.append(decay_scheduler) milestones.append(decay_start_step) @@ -104,10 +128,11 @@ def setup_scheduler( return setup_linear_scheduler( base_optimizer, max_steps=max_steps, - warmup_steps=scheduler_config.warmup_steps, - decay_steps=scheduler_config.decay_steps, + warmup_steps=resolve_steps(scheduler_config.warmup_steps, max_steps), + decay_steps=resolve_steps(scheduler_config.decay_steps, max_steps), lr=lr, min_lr=scheduler_config.min_lr, + decay_shape=scheduler_config.decay_shape, ) case "cosine": return setup_cosine_scheduler( diff --git a/tests/unit/train/pretrain/test_schedules.py b/tests/unit/train/pretrain/test_schedules.py deleted file mode 100644 index e02f526fa74..00000000000 --- a/tests/unit/train/pretrain/test_schedules.py +++ /dev/null @@ -1,56 +0,0 @@ -import math - -import pytest -import torch -from pydantic import ValidationError - -from prime_rl.configs.pretrain import PretrainConfig, ScheduleConfig -from prime_rl.trainer.pretrain.schedules import LRSchedule - - -def multipliers(config: ScheduleConfig, total_steps: int) -> list[float]: - params = [torch.nn.Parameter(torch.zeros(1))] - schedule = LRSchedule(config, torch.optim.AdamW(params, lr=1.0), total_steps) - return [schedule.multiplier(step) for step in range(total_steps)] - - -def test_wsd_endpoints_and_shapes(): - f = multipliers(ScheduleConfig(warmup=0.1, decay=0.2, end_ratio=0.05), total_steps=100) - assert f[0] == 0.0 - assert f[9] == 1.0 - assert f[79] == 1.0 - assert f[89] == pytest.approx(0.05 + 0.95 * 0.5) - assert f[99] == pytest.approx(0.05) - - for shape, midpoint in [("cosine", 0.5), ("1-sqrt", 1 - math.sqrt(0.5))]: - f = multipliers(ScheduleConfig(warmup=0, decay=20, decay_shape=shape), total_steps=100) - assert f[0] == 1.0 - assert f[89] == pytest.approx(midpoint) - assert f[99] == 0.0 - - -def test_schedule_phases_must_fit_max_steps(): - config = PretrainConfig( - model={"arch": {"model_type": "llama"}}, - tokenizer={"name": "t"}, - data={"type": "fake", "seq_len": 8, "batch_tokens": 8, "micro_batch_tokens": 8}, - optim={"lr": 1e-3}, - max_steps=100, - ) - with pytest.raises(ValidationError, match="exceed max_steps"): - PretrainConfig(**{**config.model_dump(), "schedule": {"warmup": 50, "decay": 51}}) - # Fractions summing to 1 fit for any max_steps. - PretrainConfig(**{**config.model_dump(), "max_steps": 21095}) - - -def test_apply_keeps_group_ratios(): - params = [torch.nn.Parameter(torch.zeros(2)), torch.nn.Parameter(torch.zeros(2))] - optimizer = torch.optim.AdamW([{"params": params[:1], "lr": 1.0}, {"params": params[1:], "lr": 0.1}]) - schedule = LRSchedule(ScheduleConfig(warmup=0, decay=2, end_ratio=0.5), optimizer, total_steps=4) - assert schedule.apply(0) == 1.0 - assert [group["lr"] for group in optimizer.param_groups] == [1.0, 0.1] - - # Loading a state dict replaces the param group dicts. - optimizer.load_state_dict(optimizer.state_dict()) - assert schedule.apply(3) == 0.5 - assert [group["lr"] for group in optimizer.param_groups] == pytest.approx([0.5, 0.05]) diff --git a/tests/unit/train/test_scheduler.py b/tests/unit/train/test_scheduler.py new file mode 100644 index 00000000000..57df2dd3c46 --- /dev/null +++ b/tests/unit/train/test_scheduler.py @@ -0,0 +1,28 @@ +import pytest +import torch + +from prime_rl.configs.trainer import LinearSchedulerConfig +from prime_rl.trainer.scheduler import setup_scheduler + + +def lrs(config: LinearSchedulerConfig, max_steps: int) -> list[list[float]]: + params = [torch.nn.Parameter(torch.zeros(1)), torch.nn.Parameter(torch.zeros(1))] + optimizer = torch.optim.AdamW([{"params": params[:1], "lr": 1.0}, {"params": params[1:], "lr": 0.5}]) + scheduler = setup_scheduler(optimizer, config, max_steps, lr=1.0) + values = [] + for _ in range(max_steps): + values.append([group["lr"] for group in optimizer.param_groups]) + optimizer.step() + scheduler.step() + return values + + +def test_linear_scheduler_fractions_and_decay_shapes(): + linear = lrs(LinearSchedulerConfig(warmup_steps=0.1, decay_steps=0.9), max_steps=20) + assert linear[2] == [1.0, 0.5] and linear[-1][0] < 1e-7 + shaped = lrs(LinearSchedulerConfig(warmup_steps=2, decay_steps=0.5, decay_shape="1-sqrt"), max_steps=20) + # The decay runs over the last 10 updates: 1 at its start, 1 - sqrt(5/9) five updates in, ~0 at the end. + assert shaped[10] == [1.0, 0.5] + assert shaped[15][0] == pytest.approx(1 - (5 / 9) ** 0.5, abs=1e-6) + assert shaped[15][1] == pytest.approx(shaped[15][0] / 2) + assert shaped[-1][0] < 1e-7 From 8a4a579b6ddbb5f455aeccaaa3241c44fb326ec3 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Sun, 18 Oct 2026 15:57:40 +0200 Subject: [PATCH 076/132] feat(pretrain): compile MoE blocks inside full activation checkpoints; keep fp32 routers out of FSDP Compiling a checkpointed block makes the checkpoint a higher-order op that cannot hold a graph break, so the EP all-to-all, the router's FSDP hooks and other breaks made dynamo run every fully checkpointed block eagerly. The pretrain loop now compiles the block inside an eager full checkpoint, where breaks only split the graph. model.moe_router_fsdp = "replicated" (the pretrain default) leaves the fp32 router gates out of FSDP and averages their gradients in the pretrain step; as its own FSDP unit inside the block the router can fail the checkpoint recompute. SFT and RL keep the default "unit" and compiled checkpoints. --- .../src/prime_rl/configs/pretrain.py | 2 + .../src/prime_rl/configs/trainer.py | 3 + src/prime_rl/trainer/model.py | 55 +++++++++++++------ src/prime_rl/trainer/pretrain/step.py | 8 +++ src/prime_rl/trainer/pretrain/train.py | 1 + 5 files changed, 51 insertions(+), 18 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 04c2b54a98e..b0d834e48d8 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -61,6 +61,8 @@ class PretrainModelConfig(ModelConfig): moe: PretrainMoEConfig = PretrainMoEConfig() + moe_router_fsdp: Literal["unit", "replicated"] = "replicated" + offload: Literal["none"] = "none" """CPU offloading is not supported: pretraining's per-role parameter groups need ``none``.""" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 858d7810749..c852ab117e9 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -307,6 +307,9 @@ class ModelConfig(BaseModelConfig): moe_router_dtype: Literal["bfloat16", "float32", "auto"] = "auto" """Compute dtype for MoE router gates. ``float32`` keeps router gate weights in fp32 through forward and backward (exempt from FSDP bf16 parameter casting) and computes the gate GEMM and routing logits in fp32, matching models trained with fp32 routing (e.g. GLM-5.x via Megatron's ``--moe-router-dtype fp32``). ``bfloat16`` computes the gate GEMM in the model compute dtype. ``auto`` (default) resolves to ``float32`` for RL and ``bfloat16`` for SFT. Router score functions (sigmoid/softmax) run in fp32 regardless. A no-op for non-MoE models.""" + moe_router_fsdp: Literal["unit", "replicated"] = "unit" + """How FSDP holds fp32 MoE router gates (``moe_router_dtype = "float32"``). ``unit``: each router is its own FSDP unit inside its transformer block, so its hooks run inside the block's activation checkpoint and compiled graph: the checkpoint recompute can fail ("... not found in storage") and the compiled graph breaks there. ``replicated``: the router gates stay out of FSDP, as full fp32 copies on every rank; the training loop must average their gradients over the data-parallel ranks (the pretrain loop does).""" + reduce_scatter_max_input_buffers: int = Field(1, ge=1) """Maximum number of FSDP gradient reduce-scatter input buffers kept in flight. Raising it trades memory (about one layer's gradients per extra buffer) for overlap when reduce-scatters are exposed, e.g. MoE without EP or full activation checkpointing.""" diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 60417a996bd..e7a78147f39 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -567,6 +567,7 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None: transformer_layers = language_model.layers fullgraph = config.compile is not None and config.compile.fullgraph + replicated_router_params: set[nn.Parameter] = set() for transformer_block in transformer_layers: for module in transformer_block.modules(): if isinstance(module, NGramEmbedding) and parallel_dims.get_mesh("head").size() > 1: @@ -593,21 +594,28 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None: } if config.moe_router_dtype == "float32" and isinstance(block_mlp, MoE): - if fullgraph: - raise ValueError( - "model.compile.fullgraph=true requires model.moe_router_dtype='bfloat16': the fp32 router is " - "its own FSDP unit inside the compiled block, and dynamo cannot trace FSDP hooks." + if config.moe_router_fsdp == "replicated": + # Left out of FSDP, so no FSDP hooks run inside the (checkpointed, compiled) block. + router_params = set(block_mlp.router.parameters()) + replicated_router_params |= router_params + block_fsdp_config = {**block_fsdp_config, "ignored_params": router_params} + else: + if fullgraph: + raise ValueError( + "model.compile.fullgraph=true requires model.moe_router_dtype='bfloat16' or " + "model.moe_router_fsdp='replicated': the fp32 router is its own FSDP unit inside the compiled " + "block, and dynamo cannot trace FSDP hooks." + ) + # Own FSDP unit with an fp32 policy so the gate weight is not cast to + # bf16 for forward and its gradients reduce in fp32. + fully_shard( + block_mlp.router, + mesh=hsdp_mesh, + mp_policy=MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32), + offload_policy=offload_policy, + reshard_after_forward=config.reshard_after_forward, + shard_placement_fn=shard_placement_fn, ) - # Own FSDP unit with an fp32 policy so the gate weight is not cast to - # bf16 for forward and its gradients reduce in fp32. - fully_shard( - block_mlp.router, - mesh=hsdp_mesh, - mp_policy=MixedPrecisionPolicy(param_dtype=torch.float32, reduce_dtype=torch.float32), - offload_policy=offload_policy, - reshard_after_forward=config.reshard_after_forward, - shard_placement_fn=shard_placement_fn, - ) fully_shard( transformer_block, @@ -652,6 +660,7 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None: offload_policy=offload_policy, reshard_after_forward=config.reshard_after_forward, shard_placement_fn=shard_placement_fn, + ignored_params=replicated_router_params, ) model.set_reduce_scatter_max_input_buffers(config.reduce_scatter_max_input_buffers) @@ -849,13 +858,20 @@ def mark_dynamic_int_args() -> None: torch.compiler.config.dynamic_sources = ",".join(sources) -def apply_compile(model: nn.Module, compile_config: CompileConfig): +def apply_compile(model: nn.Module, compile_config: CompileConfig, inside_checkpoint: bool = False): + """Compile every transformer block. With ``inside_checkpoint``, compile the block inside its activation + checkpoint wrapper instead of the checkpointed block: a graph break inside a compiled checkpoint (the EP + all-to-all, FSDP hooks) makes dynamo run the whole block eagerly, while inside an eager checkpoint it only + splits the graph.""" torch._dynamo.config.capture_scalar_outputs = True mark_dynamic_int_args() language_model = get_language_model(model) for layer_id in range(len(language_model.layers)): + layer = language_model.layers[layer_id] + if inside_checkpoint: + layer = getattr(layer, "_checkpoint_wrapped_module", layer) # Doing it in-place avoids mangled fqn which can break checkpoint loading - language_model.layers[layer_id].compile(fullgraph=compile_config.fullgraph, mode=compile_config.mode) + layer.compile(fullgraph=compile_config.fullgraph, mode=compile_config.mode) get_logger().info( f"Compiled {len(language_model.layers)} layers (fullgraph={compile_config.fullgraph}, mode={compile_config.mode})" ) @@ -968,10 +984,12 @@ def setup_model( loading_from_checkpoint_later: bool = False, init: InitConfig | None = None, parametrization: ParametrizationConfig | None = None, + compile_inside_full_checkpoint: bool = False, ) -> nn.Module: """Build the sharded model. Its weights come from ``config.name``, or later from a checkpoint (``loading_from_checkpoint_later``); with ``init``, they are initialized from scratch (with the CompleteP - multipliers of ``parametrization``) instead.""" + multipliers of ``parametrization``) instead. ``compile_inside_full_checkpoint`` compiles blocks inside their + full activation checkpoints (see ``apply_compile``).""" resolve_auto_attn(config) if config.attn == "flash_attention_3" and not is_flash_attn_3_available(): @@ -1034,7 +1052,8 @@ def setup_model( if config.ac is not None: apply_ac(model, config.ac) if config.compile is not None: - apply_compile(model, config.compile) + inside_checkpoint = compile_inside_full_checkpoint and config.ac is not None and config.ac.mode == "full" + apply_compile(model, config.compile, inside_checkpoint=inside_checkpoint) setup_fsdp(model, config, parallel_dims) diff --git a/src/prime_rl/trainer/pretrain/step.py b/src/prime_rl/trainer/pretrain/step.py index f17921de2f4..c85fcca2ab7 100644 --- a/src/prime_rl/trainer/pretrain/step.py +++ b/src/prime_rl/trainer/pretrain/step.py @@ -8,6 +8,7 @@ import torch import torch.distributed as dist from torch import Tensor, nn +from torch.distributed.tensor import DTensor from prime_rl.trainer.model import is_tt_moe_model, iter_moe_blocks from prime_rl.trainer.parallel_dims import ParallelDims @@ -89,6 +90,13 @@ def run_micro_batches( sums["micro_max_mean"] = sums.get("micro_max_mean", zero) + micro_max_mean.mean() maxes["micro_max_mean_max"] = torch.maximum(maxes.get("micro_max_mean_max", zero), micro_max_mean.amax()) + if is_moe_model: + # Router gates kept out of FSDP (``model.moe_router_fsdp = "replicated"``) average like FSDP's reduction. + for moe in iter_moe_blocks(model): + for param in moe.router.parameters(): + if param.grad is not None and not isinstance(param, DTensor): + dist.all_reduce(param.grad, op=dist.ReduceOp.AVG, group=dp_cp_group) + sums["data_time"] = zero + data_time maxes["data_time_max"] = zero + data_time extra_sums = extra_sums or {} diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 303e684821c..64ff6335bc5 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -151,6 +151,7 @@ def train(config: PretrainConfig): loading_from_checkpoint_later=checkpoint_step is not None, init=config.model.init if checkpoint_step is None else None, parametrization=config.parametrization, + compile_inside_full_checkpoint=True, ) m_n, m_l = multipliers(config.parametrization, model) if parallel_dims.cp_enabled: From f16c919682411973df75135828e32c00f7f2470a Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 09:15:28 +0200 Subject: [PATCH 077/132] feat(pretrain): prime_moe, pretraining's own MoE model GLM-4.5's architecture (glm4_moe's config fields and state-dict layout) plus Gated Residual connections (hc_count, hc_lowrank) and null experts (n_null_experts, variable-k dispatch). PrimeRouter/PrimeMoE subclass the shared router and MoE runtime, so FSDP, EP, the expert compute and the dispatchers apply unchanged; the sequence-wise balance moves into PrimeMoE. Registered as model_type prime_moe. --- src/prime_rl/trainer/models/__init__.py | 3 + .../trainer/models/prime_moe/__init__.py | 4 + .../prime_moe/configuration_prime_moe.py | 111 +++++++ .../models/prime_moe/hyper_connection.py | 74 +++++ .../models/prime_moe/modeling_prime_moe.py | 272 ++++++++++++++++++ src/prime_rl/trainer/models/prime_moe/moe.py | 165 +++++++++++ 6 files changed, 629 insertions(+) create mode 100644 src/prime_rl/trainer/models/prime_moe/__init__.py create mode 100644 src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py create mode 100644 src/prime_rl/trainer/models/prime_moe/hyper_connection.py create mode 100644 src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py create mode 100644 src/prime_rl/trainer/models/prime_moe/moe.py diff --git a/src/prime_rl/trainer/models/__init__.py b/src/prime_rl/trainer/models/__init__.py index af9a0684f14..b004af3638b 100644 --- a/src/prime_rl/trainer/models/__init__.py +++ b/src/prime_rl/trainer/models/__init__.py @@ -20,6 +20,7 @@ from prime_rl.trainer.models.llama import LlamaForCausalLM from prime_rl.trainer.models.minimax_m2 import MiniMaxM2Config, MiniMaxM2ForCausalLM from prime_rl.trainer.models.nemotron_h import NemotronHConfig, NemotronHForCausalLM +from prime_rl.trainer.models.prime_moe import PrimeMoeConfig, PrimeMoeForCausalLM from prime_rl.trainer.models.qwen3 import Qwen3ForCausalLM from prime_rl.trainer.models.qwen3_5 import ( Qwen3_5Config, @@ -45,6 +46,7 @@ AutoConfig.register("laguna", LagunaConfig, exist_ok=True) AutoConfig.register("minimax_m2", MiniMaxM2Config, exist_ok=True) AutoConfig.register("nemotron_h", NemotronHConfig, exist_ok=True) +AutoConfig.register("prime_moe", PrimeMoeConfig, exist_ok=True) AutoConfig.register("qwen3_moe", Qwen3MoeConfig, exist_ok=True) AutoConfig.register("qwen3_5", Qwen3_5Config, exist_ok=True) AutoConfig.register("qwen3_5_text", Qwen3_5TextConfig, exist_ok=True) @@ -68,6 +70,7 @@ (LagunaConfig, LagunaForCausalLM), (MiniMaxM2Config, MiniMaxM2ForCausalLM), (NemotronHConfig, NemotronHForCausalLM), + (PrimeMoeConfig, PrimeMoeForCausalLM), (Qwen3MoeConfig, Qwen3MoeForCausalLM), (Qwen3_5TextConfig, Qwen3_5ForCausalLM), (Qwen3_5MoeTextConfig, Qwen3_5ForCausalLM), diff --git a/src/prime_rl/trainer/models/prime_moe/__init__.py b/src/prime_rl/trainer/models/prime_moe/__init__.py new file mode 100644 index 00000000000..00162be6fbb --- /dev/null +++ b/src/prime_rl/trainer/models/prime_moe/__init__.py @@ -0,0 +1,4 @@ +from prime_rl.trainer.models.prime_moe.configuration_prime_moe import PrimeMoeConfig +from prime_rl.trainer.models.prime_moe.modeling_prime_moe import PrimeMoeForCausalLM, PrimeMoeModel + +__all__ = ["PrimeMoeConfig", "PrimeMoeForCausalLM", "PrimeMoeModel"] diff --git a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py new file mode 100644 index 00000000000..6b7d18e9eaf --- /dev/null +++ b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py @@ -0,0 +1,111 @@ +from transformers.configuration_utils import PretrainedConfig + + +class PrimeMoeConfig(PretrainedConfig): + r""" + Configuration of the pretraining MoE (`PrimeMoeForCausalLM`): GLM-4.5's architecture (attention with optional QK + norm, sigmoid-scored token-choice experts with a selection bias, shared experts, leading dense layers) plus + Gated Residual connections and null experts. The fields shared with `Glm4MoeConfig` have the same names. + + Args: + vocab_size (`int`): Vocabulary size. + hidden_size (`int`): Width of the residual stream. + intermediate_size (`int`): Hidden size of the dense MLP layers. + num_hidden_layers (`int`): Number of decoder layers. + num_attention_heads (`int`): Number of query heads. + num_key_value_heads (`int`): Number of key/value heads (grouped-query attention). + head_dim (`int`, *optional*): Dimension of each attention head. Defaults to `hidden_size // num_attention_heads`. + partial_rotary_factor (`float`): Share of each head's dimensions that RoPE rotates. + hidden_act (`str`): Activation of the gated MLPs. + max_position_embeddings (`int`): Maximum sequence length. + initializer_range (`float`): Unused by prime-rl's initialization; kept for Hugging Face compatibility. + rms_norm_eps (`float`): Epsilon of the RMS norms. + tie_word_embeddings (`bool`): Whether the LM head shares the input embedding. + rope_theta (`float`): RoPE base period. + rope_scaling (`dict`, *optional*): RoPE scaling configuration. + attention_bias (`bool`): Whether the attention projections have biases. + attention_dropout (`float`): Attention dropout (unused in training). + moe_intermediate_size (`int`): Hidden size of each routed and shared expert. + num_experts_per_tok (`int`): Experts each token selects (`top_k`), null experts included. + n_shared_experts (`int`): Number of shared experts; their hidden sizes add up. + n_routed_experts (`int`): Number of routed experts. + routed_scaling_factor (`float`): Scale of the routed experts' output. + first_k_dense_replace (`int`): Number of leading dense layers. + norm_topk_prob (`bool`): Whether the selected experts' scores are normalized to sum to 1. + use_qk_norm (`bool`): Whether queries and keys are RMS-normalized. + hc_count (`int`): Number of residual streams. Above 1, Gated Residual connections (Qwen3.8-Next) replace the + pre-norms and the final norm. + hc_lowrank (`int`, *optional*): Rank of the Gated Residual read gates; required when `hc_count` is above 1. + n_null_experts (`int`): Zero-output experts the router can select next to the routed experts (AdaMoE); a + token that selects them uses fewer routed experts. + """ + + model_type = "prime_moe" + + def __init__( + self, + vocab_size=151552, + hidden_size=4096, + intermediate_size=10944, + num_hidden_layers=46, + num_attention_heads=96, + partial_rotary_factor=0.5, + num_key_value_heads=8, + head_dim=None, + hidden_act="silu", + max_position_embeddings=131072, + initializer_range=0.02, + rms_norm_eps=1e-5, + use_cache=True, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + moe_intermediate_size=1408, + num_experts_per_tok=8, + n_shared_experts=1, + n_routed_experts=128, + routed_scaling_factor=1.0, + first_k_dense_replace=1, + norm_topk_prob=True, + use_qk_norm=False, + hc_count=1, + hc_lowrank=None, + n_null_experts=0, + pad_token_id=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.partial_rotary_factor = partial_rotary_factor + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.pad_token_id = pad_token_id + + self.moe_intermediate_size = moe_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.routed_scaling_factor = routed_scaling_factor + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.use_qk_norm = use_qk_norm + + self.hc_count = hc_count + self.hc_lowrank = hc_lowrank + self.n_null_experts = n_null_experts + + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) diff --git a/src/prime_rl/trainer/models/prime_moe/hyper_connection.py b/src/prime_rl/trainer/models/prime_moe/hyper_connection.py new file mode 100644 index 00000000000..3082dc59768 --- /dev/null +++ b/src/prime_rl/trainer/models/prime_moe/hyper_connection.py @@ -0,0 +1,74 @@ +import torch +import torch.nn.functional as F +from torch import nn + + +class ExpandedRMSNorm(nn.Module): + """RMS-normalize each residual stream with an independent affine weight.""" + + def __init__(self, hidden_size: int, stream_count: int, eps: float) -> None: + super().__init__() + self.hidden_size = hidden_size + self.stream_count = stream_count + self.eps = eps + self.weight = nn.Parameter(torch.zeros(stream_count * hidden_size)) + + def reset_parameters(self) -> None: + nn.init.zeros_(self.weight) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + streams = hidden_states.float().unflatten(-1, (self.stream_count, self.hidden_size)) + variance = streams.square().mean(dim=-1, keepdim=True) + normalized = streams * torch.rsqrt(variance + self.eps) + return (normalized.flatten(-2) * (1.0 + self.weight.float())).to(input_dtype) + + +class HyperConnection(nn.Module): + """Learned mixing and residual injection across expanded residual streams.""" + + def __init__( + self, + *, + hidden_size: int, + stream_count: int, + low_rank: int, + norm_eps: float, + with_residual_injection: bool = True, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.stream_count = stream_count + expanded_size = stream_count * hidden_size + + self.hc_norm = ExpandedRMSNorm(hidden_size, stream_count, norm_eps) + self.input_mix_weight_down = nn.Linear(expanded_size, low_rank, bias=False) + self.input_mix_weight_up = nn.Linear(low_rank, expanded_size, bias=False) + self.block_inject_weight = ( + nn.Linear(expanded_size, stream_count, bias=False) if with_residual_injection else None + ) + + def mix(self, expanded_states: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + normalized_states = self.hc_norm(expanded_states) + mixing_weights = F.silu(self.input_mix_weight_down(normalized_states) / self.stream_count) + mixing_weights = torch.sigmoid(self.input_mix_weight_up(mixing_weights)) + mixing_weights = mixing_weights.unflatten(-1, (self.stream_count, self.hidden_size)) + normalized_streams = normalized_states.unflatten(-1, (self.stream_count, self.hidden_size)) + block_input = (mixing_weights * normalized_streams).mean(dim=-2) + return block_input.to(expanded_states.dtype), (expanded_states, normalized_states) + + def forward(self, expanded_states: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + return self.mix(expanded_states) + + def combine( + self, + block_output: torch.Tensor, + residual_state: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + expanded_states, normalized_states = residual_state + injection_weights = 2.0 * torch.sigmoid(self.block_inject_weight(normalized_states) / self.stream_count) + residual_streams = expanded_states.unflatten(-1, (self.stream_count, self.hidden_size)) + return (residual_streams + block_output.unsqueeze(-2) * injection_weights.unsqueeze(-1)).flatten(-2) + + +__all__ = ["ExpandedRMSNorm", "HyperConnection"] diff --git a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py new file mode 100644 index 00000000000..c8cd8ba6f55 --- /dev/null +++ b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py @@ -0,0 +1,272 @@ +"""The pretraining MoE: GLM-4.5's architecture (prime-rl's `glm4_moe`) with Gated Residual connections and null +experts. Module names follow `glm4_moe`, so checkpoints share its state-dict layout.""" + +from typing import Optional, Union + +import torch +from torch import Tensor, nn +from transformers.cache_utils import Cache +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_outputs import BaseModelOutputWithPast +from transformers.processing_utils import Unpack +from transformers.utils import TransformersKwargs + +from prime_rl.trainer.models.base import PreTrainedModelPrimeRL +from prime_rl.trainer.models.glm4_moe.converting_glm4_moe import conversion_chain +from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig +from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput +from prime_rl.trainer.models.layers.mlp import FeedForward +from prime_rl.trainer.models.layers.moe import GroupedExperts +from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig +from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig +from prime_rl.trainer.models.prime_moe.configuration_prime_moe import PrimeMoeConfig +from prime_rl.trainer.models.prime_moe.hyper_connection import HyperConnection +from prime_rl.trainer.models.prime_moe.moe import PrimeMoE, PrimeRouter +from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens + + +class PrimeMoeDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: PrimeMoeConfig, layer_idx: int): + super().__init__() + attn_config = AttentionConfig( + hidden_size=config.hidden_size, + head_dim=config.head_dim, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + is_causal=True, + attention_bias=config.attention_bias, + use_qk_norm=config.use_qk_norm, + rms_norm_eps=config.rms_norm_eps, + ) + self.self_attn = ATTN_IMPL2CLASS[config._attn_implementation](attn_config) + + if layer_idx >= config.first_k_dense_replace: + shared_expert = None + if config.n_shared_experts > 0: + shared_expert = FeedForward( + dim=config.hidden_size, + hidden_dim=config.moe_intermediate_size * config.n_shared_experts, + expert_type="gated", + activation=config.hidden_act, + ) + self.mlp = PrimeMoE( + router=PrimeRouter( + dim=config.hidden_size, + num_experts=config.n_routed_experts, + num_null_experts=config.n_null_experts, + top_k=config.num_experts_per_tok, + route_norm=config.norm_topk_prob, + route_scale=config.routed_scaling_factor, + ), + experts=GroupedExperts( + dim=config.hidden_size, + hidden_dim=config.moe_intermediate_size, + num_experts=config.n_routed_experts, + expert_type="gated", + activation=config.hidden_act, + ), + shared_expert=shared_expert, + ) + else: + self.mlp = FeedForward( + dim=config.hidden_size, hidden_dim=config.intermediate_size, activation=config.hidden_act + ) + + self.gated_residual = config.hc_count > 1 + if self.gated_residual: + connection_args = { + "hidden_size": config.hidden_size, + "stream_count": config.hc_count, + "low_rank": config.hc_lowrank, + "norm_eps": config.rms_norm_eps, + } + self.attn_hyper_connection = HyperConnection(**connection_args) + self.mlp_hyper_connection = HyperConnection(**connection_args) + else: + self.input_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + self.post_attention_layernorm = RMSNorm( + RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps) + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor, + max_seqlen: int, + ) -> torch.Tensor: + if self.gated_residual: + block_input, residual_state = self.attn_hyper_connection.mix(hidden_states) + block_output, _ = self.self_attn( + hidden_states=block_input, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + hidden_states = self.attn_hyper_connection.combine(block_output, residual_state) + block_input, residual_state = self.mlp_hyper_connection.mix(hidden_states) + block_output = self.mlp(block_input) + return self.mlp_hyper_connection.combine(block_output, residual_state) + + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class PrimeMoePreTrainedModel(PreTrainedModelPrimeRL): + config: PrimeMoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["PrimeMoeDecoderLayer"] + _supports_flash_attn = True + _supports_sdpa = False + _supports_flex_attn = True + _can_compile_fullgraph = False + _supports_attention_backend = True + + @classmethod + def keep_in_fp32_for_weight_transfer(cls, name: str) -> bool: + return name.endswith("mlp.router.selection_bias") + + @classmethod + def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.1.up_proj" in name or "mlp.experts.gate_up_proj" in name for name in state_dict) + + @classmethod + def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return any("mlp.experts.gate_proj" in name for name in state_dict) + + @classmethod + def conversion_chain(cls, config): + # Same expert layout as glm4_moe. + return conversion_chain(config) + + +class PrimeMoeModel(PrimeMoePreTrainedModel): + def __init__(self, config: PrimeMoeConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [PrimeMoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + if config.hc_count > 1: + self.hyper_connection_mixer = HyperConnection( + hidden_size=config.hidden_size, + stream_count=config.hc_count, + low_rank=config.hc_lowrank, + norm_eps=config.rms_norm_eps, + with_residual_injection=False, + ) + else: + self.norm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) + + if isinstance(config.rope_scaling, dict): + rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + rope_type = "default" + self.rotary_emb = RotaryEmbedding( + RotaryEmbeddingConfig( + max_position_embeddings=config.max_position_embeddings, rope_type=rope_type, model_config=config + ) + ) + self.gradient_checkpointing = False + self.post_init() + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + *, + seq_lens: torch.LongTensor, + seq_lens_are_pre_shard: bool = False, + ) -> BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens( + seq_lens.to(device=inputs_embeds.device), + total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1], + ) + torch._dynamo.mark_dynamic(cu_seqlens, 0) + + position_embeddings = self.rotary_emb(inputs_embeds, position_ids) + # Gated Residual: every residual stream starts from the token embedding. + hidden_states = inputs_embeds.repeat(1, 1, self.config.hc_count) + for decoder_layer in self.layers: + hidden_states = decoder_layer( + hidden_states, position_embeddings=position_embeddings, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen + ) + + if self.config.hc_count > 1: + hidden_states, _ = self.hyper_connection_mixer(hidden_states) + else: + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast(last_hidden_state=hidden_states) + + +class PrimeMoeForCausalLM(PrimeMoePreTrainedModel): + def __init__(self, config: PrimeMoeConfig): + super().__init__(config) + self.model = PrimeMoeModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.post_init() + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + temperature: Optional[torch.Tensor] = None, + *, + seq_lens: torch.LongTensor, + seq_lens_are_pre_shard: bool = False, + **kwargs: Unpack[TransformersKwargs], + ) -> PrimeLmOutput: + """``labels`` and ``temperature`` go to prime-rl's wrapped LM head, which then returns per-token logprobs + and entropy; without them it returns logits.""" + assert use_cache is None and past_key_values is None, "prime_moe does not support KV caching" + if position_ids is None: + tokens = inputs_embeds if inputs_embeds is not None else input_ids + position_ids = torch.arange(tokens.shape[1], device=tokens.device).unsqueeze(0) + + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + seq_lens=seq_lens, + seq_lens_are_pre_shard=seq_lens_are_pre_shard, + ) + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + return self.lm_head( + hidden_states[:, slice_indices, :], + labels[:, slice_indices] if labels is not None else None, + temperature=temperature, + ) + + def init_buffers_post_meta(self): + rotary_emb = self.model.rotary_emb + inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn(rotary_emb.config, rotary_emb.inv_freq.device) + rotary_emb.inv_freq.copy_(inv_freq) diff --git a/src/prime_rl/trainer/models/prime_moe/moe.py b/src/prime_rl/trainer/models/prime_moe/moe.py new file mode 100644 index 00000000000..4379461ca09 --- /dev/null +++ b/src/prime_rl/trainer/models/prime_moe/moe.py @@ -0,0 +1,165 @@ +import torch +import torch.nn.functional as F + +from prime_rl.trainer.distributed.token_dispatcher import TokenDispatcher +from prime_rl.trainer.models.layers.mlp import FeedForward +from prime_rl.trainer.models.layers.moe import ( + GroupedExperts, + MoE, + TokenChoiceTopKRouter, + record_moe_routing_statistics, +) + + +class PrimeRouter(TokenChoiceTopKRouter): + """Sigmoid-scored top-k router with a selection bias and ``num_null_experts`` zero-output experts next to the + routed ones (AdaMoE). Null experts are scored, selected and balanced like the others, so a token that selects + them uses fewer routed experts. Returns the scores of all experts too, for the sequence-wise balance loss.""" + + def __init__( + self, *, dim: int, num_experts: int, num_null_experts: int, top_k: int, route_norm: bool, route_scale: float + ): + super().__init__( + dim=dim, + num_experts=num_experts + num_null_experts, + top_k=top_k, + score_func="sigmoid", + route_norm=route_norm, + route_scale=route_scale, + selection_bias=True, + ) + self.num_null_experts = num_null_experts + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Like ``TokenChoiceTopKRouter.forward`` with sigmoid scores; also returns the ``(bs*slen, num_experts)`` + scores of all experts.""" + if self.fp32_gate: + gate_bias = self.gate.bias.float() if self.gate.bias is not None else None + logits = F.linear(x.float(), self.gate.weight.float(), gate_bias) + else: + logits = self.gate(x) + scores = torch.sigmoid(logits.float()) + # The selection bias only affects which experts are selected, not their scores. + _, selected_experts_indices = torch.topk( + scores + self.selection_bias, k=self.top_k, dim=1, sorted=self.topk_sorted + ) + top_scores = scores.gather(dim=1, index=selected_experts_indices) + + with torch.no_grad(): + routing_confidence_sum = (top_scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)).sum() + + if self.route_norm: + top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20) + top_scores = top_scores * self.route_scale + + num_tokens_per_expert = torch.histc( + selected_experts_indices.reshape(-1).float(), bins=self.num_experts, min=0, max=self.num_experts + ).to(torch.int64) + return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, scores + + +def sequence_balance(scores: torch.Tensor, selected_experts: torch.Tensor, num_rows: int) -> torch.Tensor: + """DeepSeek-V3's sequence-wise balance loss without its coefficient, per row: ``sum_i f_i P_i`` with + ``f_i`` the row's share of expert ``i``'s selections times ``num_experts / top_k`` (1 when balanced) and + ``P_i`` its mean score normalized over all experts. 1 for a balanced row, up to ``num_experts / top_k``. + + Args: + scores: ``(num_rows * slen, num_experts)`` scores of all experts. + selected_experts: ``(num_rows * slen, top_k)`` selected experts. + """ + num_experts = scores.shape[1] + selections = selected_experts.view(num_rows, -1) + counts = torch.zeros(num_rows, num_experts, device=scores.device).scatter_add_( + 1, selections, torch.ones_like(selections, dtype=torch.float32) + ) + f = counts * (num_experts / selections.shape[1]) + p = (scores / scores.sum(dim=-1, keepdim=True)).view(num_rows, -1, num_experts).mean(dim=1) + return (f * p).sum(dim=-1) + + +class PrimeMoE(MoE): + """The shared token-choice MoE runtime with a ``PrimeRouter``. With null experts, tokens use different numbers of + routed experts: only the (token, routed expert) pairs are dispatched, one per row. Also computes the sequence-wise + balance for its loss (``balance_rows``) and counts the routed experts per token (``routed_k_counts``).""" + + def __init__(self, *, router: PrimeRouter, experts: GroupedExperts, shared_expert: FeedForward | None) -> None: + super().__init__( + router=router, + experts=experts, + shared_expert=shared_expert, + score_before_experts=False, + load_balance_coeff=1e-3, + ) + self.variable_k = router.num_null_experts > 0 + self.set_token_dispatcher(self.token_dispatcher) + # Usage of every expert, null experts included. + self.register_buffer( + "tokens_per_expert", torch.zeros(router.num_experts, dtype=torch.float32), persistent=False + ) + # Training tokens per number of routed experts used, for variable-k routing. + self.register_buffer( + "routed_k_counts", + torch.zeros(router.top_k + 1, dtype=torch.float32) if self.variable_k else None, + persistent=False, + ) + # Sequence-wise balance (DeepSeek-V3 eq. 17-20): with balance_rows set by the training loop, which knows + # how many packed rows a sample holds, every forward with grad computes the per-row balance for its loss. + self.balance_rows: int | None = None + self.seq_balance: torch.Tensor | None = None + + def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None: + # Variable-k routing dispatches one (token, expert) pair per row. + token_dispatcher.top_k = 1 if self.variable_k else self.router.top_k + self.token_dispatcher = token_dispatcher + + def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) -> torch.Tensor: + assert routed_experts is None, "PrimeMoE does not support router replay" + bs, slen, dim = x.shape + x = x.view(-1, dim) + top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, scores = self.router(x) + + if self.balance_rows is not None and torch.is_grad_enabled(): + self.seq_balance = sequence_balance(scores, selected_experts_indices, self.balance_rows) + + with torch.no_grad(): + record_moe_routing_statistics( + self.tokens_per_expert, self.routing_confidence_sum, num_tokens_per_expert, routing_confidence_sum + ) + + if self.variable_k: + # Dispatch only the (token, routed expert) pairs, without the null experts. + routed = selected_experts_indices < self.experts.num_experts + if torch.is_grad_enabled(): + with torch.no_grad(): + self.routed_k_counts += torch.bincount(routed.sum(dim=1), minlength=self.routed_k_counts.numel()) + token_indices, slots = routed.nonzero(as_tuple=True) + pair_output = self.token_dispatcher.run( + x[token_indices], + top_scores[token_indices, slots].unsqueeze(1), + selected_experts_indices[token_indices, slots].unsqueeze(1), + self.experts, + score_before_experts=self.score_before_experts, + ) + routed_output = ( + torch.zeros(x.shape, dtype=torch.float32, device=x.device) + .index_add(0, token_indices, pair_output.float()) + .to(x.dtype) + ) + else: + routed_output = self.token_dispatcher.run( + x, top_scores, selected_experts_indices, self.experts, score_before_experts=self.score_before_experts + ) + + shared_output = self.shared_expert(x) if self.shared_expert is not None else None + self.token_dispatcher.synchronize() + if shared_output is not None: + routed_output = routed_output + shared_output + return routed_output.reshape(bs, slen, dim) + + def init_weights(self, init_std: float, buffer_device: torch.device): + super().init_weights(init_std, buffer_device) + with torch.device(buffer_device): + self.tokens_per_expert = torch.zeros(self.router.num_experts, dtype=torch.float32) + self.router.selection_bias = torch.zeros(self.router.num_experts, dtype=torch.float32) + if self.routed_k_counts is not None: + self.routed_k_counts = torch.zeros_like(self.routed_k_counts) From f39a2a93b81549a46868c9a6aef5ce8ca2aadf7f Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 09:15:28 +0200 Subject: [PATCH 078/132] feat(pretrain): train prime_moe (Gated Residual roles, init and metrics, variable-k metrics, MoE trace) Gated Residual projections get their own CompleteP role (AdamW), zero-centred norm gains start at 0 and get their own grad-norm group; router/k/* metrics for null experts; per-expert bias, router-row norm and load every 32 steps in moe_trace.npz. Pretraining MoE models must be prime_moe; the example and docs switch to it. The sign-rule bias update moves into the pretrain loop. --- docs/pretraining.md | 16 ++++--- .../{glm4-moe-d512 => moe-d512}/config.json | 2 +- .../{glm4_moe.toml => prime_moe.toml} | 6 +-- src/prime_rl/trainer/models/init.py | 4 ++ src/prime_rl/trainer/pretrain/metrics.py | 31 +++++++++++++ .../trainer/pretrain/parametrization.py | 10 ++++- src/prime_rl/trainer/pretrain/train.py | 44 ++++++++++++++++++- tests/unit/train/pretrain/test_init.py | 6 +-- tests/unit/train/pretrain/test_moe_balance.py | 30 +++++++++++++ 9 files changed, 133 insertions(+), 16 deletions(-) rename examples/pretrain/{glm4-moe-d512 => moe-d512}/config.json (96%) rename examples/pretrain/{glm4_moe.toml => prime_moe.toml} (85%) create mode 100644 tests/unit/train/pretrain/test_moe_balance.py diff --git a/docs/pretraining.md b/docs/pretraining.md index 84448f867e2..4db5699391e 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -3,22 +3,26 @@ `uv run pretrain` trains a model from scratch on pre-tokenized token shards. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, checkpointing and monitors as the SFT trainer. It runs locally on one node or, with `[slurm]`, as a SLURM job on one or more nodes (see [Multi-node](#multi-node)). ```bash -uv run pretrain @ examples/pretrain/glm4_moe.toml +uv run pretrain @ examples/pretrain/prime_moe.toml ``` ## Model -The model is set up as in SFT and RL: `model.name` is an HF model id or a local model directory with a `config.json`. Its weights are not loaded. The model starts from a random init (below) or from a checkpoint. A new architecture is a directory with only a `config.json` for a model type in `prime_rl.trainer.models`; [`examples/pretrain/glm4-moe-d512`](../examples/pretrain/glm4-moe-d512/config.json) is the one in the example. `vocab_size` and the BOS, EOS and pad ids in the config must fit the tokenizer. The checked-in configs round `len(tokenizer)` up to a multiple of 128 and use the tokenizer's special ids, with EOS as the pad id if the tokenizer has none. +The model is set up as in SFT and RL: `model.name` is an HF model id or a local model directory with a `config.json`. Its weights are not loaded. The model starts from a random init (below) or from a checkpoint. A new architecture is a directory with only a `config.json` for a model type in `prime_rl.trainer.models`; [`examples/pretrain/moe-d512`](../examples/pretrain/moe-d512/config.json) is the one in the example. `vocab_size` and the BOS, EOS and pad ids in the config must fit the tokenizer. The checked-in configs round `len(tokenizer)` up to a multiple of 128 and use the tokenizer's special ids, with EOS as the pad id if the tokenizer has none. ```toml [model] -name = "examples/pretrain/glm4-moe-d512" +name = "examples/pretrain/moe-d512" [tokenizer] name = "tokenizers/deepseek-v4-flash-base" # defaults to model.name ``` -The trainer, the HF export (`tools/convert_dcp_to_bf16.py`) and vLLM all read the same model directory and tokenizer. The trainer logs the active non-embedding (routed experts counted at `top_k / num_experts`), non-embedding, embedding (input embedding and LM head) and total parameter counts at startup, and as `params/*` at step 0. +MoE models use `prime_moe`, pretraining's own model (`prime_rl.trainer.models.prime_moe`), so architecture work does not touch the `glm4_moe` that SFT and RL use. It is GLM-4.5's architecture with the same config fields and state-dict layout, plus: +- Gated Residual connections (Qwen3.8-Next): `hc_count` residual streams (above 1) with read gates of rank `hc_lowrank` replace the pre-norms and the final norm. Their projections train with AdamW. +- Null experts (AdaMoE): `n_null_experts` zero-output experts the router can select next to the routed ones, so tokens use between 0 and `num_experts_per_tok` routed experts. Only the (token, routed expert) pairs are dispatched. The metrics are `router/k/*`. + +The trainer and the HF export (`tools/convert_dcp_to_bf16.py`) read the same model directory and tokenizer; vLLM does not support `prime_moe` yet. The trainer logs the active non-embedding (routed experts counted at `top_k / num_experts`), non-embedding, embedding (input embedding and LM head) and total parameter counts at startup, and as `params/*` at step 0. ## Initialization @@ -73,7 +77,7 @@ The schedule scales every param group's own learning rate, so ratios between gro ## MoE load balancing -`[model.moe.balance]` balances the routed experts as in DeepSeek-V3, for every MoE family whose router has a selection bias (e.g. `glm4_moe`): +`[model.moe.balance]` balances the routed experts of `prime_moe` as in DeepSeek-V3: | Field | Meaning | |---|---| @@ -81,7 +85,7 @@ The schedule scales every param group's own learning rate, so ratios between gro | `bias_freeze` | Fraction of `max_steps` at the end during which the biases stop updating. Default 0.1. | | `seq_aux_coeff` | Weight (alpha) of the sequence-wise balance loss, computed per packed row of `data.seq_len` tokens and added to the training loss. Default 1e-3; 0 turns it off. | -The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` (1 is balanced) and `router/bias/*`. +The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` (1 is balanced) and `router/bias/*`. Every 32 steps the trainer also records each layer's per-expert selection bias, router-row norm and load in `moe_trace.npz` in the run directory. ## Data diff --git a/examples/pretrain/glm4-moe-d512/config.json b/examples/pretrain/moe-d512/config.json similarity index 96% rename from examples/pretrain/glm4-moe-d512/config.json rename to examples/pretrain/moe-d512/config.json index f821da8e747..84fe7c66888 100644 --- a/examples/pretrain/glm4-moe-d512/config.json +++ b/examples/pretrain/moe-d512/config.json @@ -10,7 +10,7 @@ "initializer_range": 0.02, "intermediate_size": 2048, "max_position_embeddings": 131072, - "model_type": "glm4_moe", + "model_type": "prime_moe", "moe_intermediate_size": 256, "n_routed_experts": 64, "n_shared_experts": 2, diff --git a/examples/pretrain/glm4_moe.toml b/examples/pretrain/prime_moe.toml similarity index 85% rename from examples/pretrain/glm4_moe.toml rename to examples/pretrain/prime_moe.toml index 12154c82924..bab6151cb97 100644 --- a/examples/pretrain/glm4_moe.toml +++ b/examples/pretrain/prime_moe.toml @@ -1,14 +1,14 @@ -# GLM-4.5-style MoE (64 routed experts, 6 active, 2 shared) with d = 512 and 8 layers (glm4-moe-d512/config.json), trained from scratch on +# GLM-4.5-style MoE (64 routed experts, 6 active, 2 shared) with d = 512 and 8 layers (moe-d512/config.json), trained from scratch on # the Dolma 3 mixture tokenized with the DeepSeek-V4 tokenizer. # 4096 steps x 256K tokens = 1.1B tokens. QK-norm and the learning rate follow the MoE QK-norm sweep (optimum 2^-7 # with QK-norm at d = 512). [parametrization] with this base shape keeps the LR when the width grows. # -# uv run pretrain @ examples/pretrain/glm4_moe.toml +# uv run pretrain @ examples/pretrain/prime_moe.toml max_steps = 4096 [model] -name = "examples/pretrain/glm4-moe-d512" # the architecture: config.json only, no weights +name = "examples/pretrain/moe-d512" # the architecture: config.json only, no weights ac = "None" [tokenizer] diff --git a/src/prime_rl/trainer/models/init.py b/src/prime_rl/trainer/models/init.py index b64323e0e58..79ea50103f7 100644 --- a/src/prime_rl/trainer/models/init.py +++ b/src/prime_rl/trainer/models/init.py @@ -9,6 +9,7 @@ from prime_rl.configs.trainer import InitConfig, ParametrizationConfig from prime_rl.trainer.models.layers.moe import TokenChoiceTopKRouter +from prime_rl.trainer.models.prime_moe.hyper_connection import ExpandedRMSNorm from prime_rl.utils.weights import resolve_fqn ParamRole = Literal["embedding", "output", "residual_out", "hidden", "norm", "bias"] @@ -111,6 +112,9 @@ def init_weights(model: nn.Module, config: InitConfig, m_n: float = 1.0, m_l: fl ) local.copy_(values.mul_(role_std(role, config.std, m_n, m_l))) # MoE selection biases are buffers: they start at 0, so routing starts from the router scores alone. + # Gated Residual norm gains are zero-centred (scale 1 + w), so they start at 0. for module in model.modules(): if isinstance(module, TokenChoiceTopKRouter) and module.selection_bias is not None: module.selection_bias.zero_() + if isinstance(module, ExpandedRMSNorm): + module.weight.zero_() diff --git a/src/prime_rl/trainer/pretrain/metrics.py b/src/prime_rl/trainer/pretrain/metrics.py index 5ef34f9564f..47f698d7a9f 100644 --- a/src/prime_rl/trainer/pretrain/metrics.py +++ b/src/prime_rl/trainer/pretrain/metrics.py @@ -41,6 +41,35 @@ def take_routing_counts(model: nn.Module) -> tuple[Tensor, Tensor]: return counts, confidence +def routed_k_metrics(model: nn.Module, group: dist.ProcessGroup) -> dict[str, float]: + """Routed experts per training token since the last call, for variable-k routing (null experts), over all ranks: + mean, std, share of tokens with none and with the maximum (the mean over MoE layers), the worst layer's std, and + the busiest rank's routed pairs per token over the mean.""" + moes = [moe for moe in iter_moe_blocks(model) if moe.routed_k_counts is not None] + if not moes: + return {} + counts = torch.stack([moe.routed_k_counts for moe in moes]) + for moe in moes: + moe.routed_k_counts.zero_() + k = torch.arange(counts.shape[1], device=counts.device) + # This rank's routed pairs per token, over the global mean: the load the data puts on one rank. + rank_mean = (counts * k).sum() / counts.sum() + rank_max = rank_mean.clone() + dist.all_reduce(rank_max, op=dist.ReduceOp.MAX, group=group) + dist.all_reduce(counts, group=group) + tokens = counts.sum(dim=1) + mean = (counts * k).sum(dim=1) / tokens + std = ((counts * (k - mean[:, None]) ** 2).sum(dim=1) / tokens).sqrt() + return { + "router/k/mean": mean.mean().item(), + "router/k/std": std.mean().item(), + "router/k/std/max": std.max().item(), + "router/k/zero_frac": (counts[:, 0] / tokens).mean().item(), + "router/k/max_frac": (counts[:, -1] / tokens).mean().item(), + "router/k/rank_max_mean": (rank_max / mean.mean()).item(), + } + + def moe_layer_indices(model: nn.Module) -> list[int]: """The decoder-layer index of every MoE layer, in the order of `iter_moe_blocks`.""" return [ @@ -106,6 +135,8 @@ def param_group(name: str) -> str: return "lm_head" if any("norm" in part for part in parts): return "norm" + if any("hyper_connection" in part for part in parts): + return "hyper_connection" if "router" in parts: return "router" if "shared_expert" in parts: diff --git a/src/prime_rl/trainer/pretrain/parametrization.py b/src/prime_rl/trainer/pretrain/parametrization.py index 4159f1f953a..ef165ba8def 100644 --- a/src/prime_rl/trainer/pretrain/parametrization.py +++ b/src/prime_rl/trainer/pretrain/parametrization.py @@ -26,12 +26,16 @@ def param_groups( own shape). The router maps width to a fixed number of experts, so its spectral-norm target sqrt(num_experts / d) shrinks as m_N^-1/2: with Muon its learning rate is scaled by 1/m_N; with AdamW it follows the hidden rule. + + Gated Residual projections (``hyper_connection``) follow the hidden rule but always train with AdamW, as in + Qwen3.8-Next. """ # (lr, eps, weight decay) scale per role for AdamW. PyTorch AdamW and Dion decay by lr * weight_decay. adam_scales = { "embedding": (1.0, 1 / m_n, 1.0), "hidden": (1 / m_n, 1 / (m_n * m_l), m_n), "router": (1 / m_n, 1 / (m_n * m_l), m_n), + "hyper_connection": (1 / m_n, 1 / (m_n * m_l), m_n), "residual_out": (1 / (m_n * m_l), 1 / m_n, m_n * m_l), "output": (1 / m_n, 1.0, m_n), "norm": (1.0, 1 / m_n, 0.0), @@ -48,7 +52,11 @@ def param_groups( for name, param in model.named_parameters(): if param.requires_grad: name = resolve_fqn(model, name) - role = "router" if ".mlp.router." in name and param.ndim == 2 else param_role(name, param) + role = param_role(name, param) + if role == "hidden" and ".mlp.router." in name: + role = "router" + elif role == "hidden" and "hyper_connection" in name: + role = "hyper_connection" params[role, ".mlp.experts." in name].append(param) groups = [] diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 64ff6335bc5..93fdb1c1995 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -13,7 +13,8 @@ import numpy as np import torch import torch.distributed as dist -from torch import nn +from torch import Tensor, nn +from torch.distributed.tensor import DTensor from prime_rl import monitors from prime_rl.configs.pretrain import PretrainConfig @@ -26,7 +27,6 @@ reshard_module, setup_model, setup_tokenizer, - update_selection_biases, ) from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX from prime_rl.trainer.optim import setup_optimizer @@ -58,6 +58,7 @@ SpikeDetector, grad_norms_by_group, reduce_stats, + routed_k_metrics, take_routing_counts, ) from prime_rl.trainer.sft.data import get_dataset_state @@ -80,6 +81,21 @@ from prime_rl.utils.weights import resolve_fqn +MOE_TRACE_INTERVAL = 32 + + +@torch.no_grad() +def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor, rate: float) -> None: + """Aux-loss-free load balancing (https://arxiv.org/abs/2408.15664): move each layer's selection bias by + ``rate`` towards underloaded experts, given the step's global ``[num_moe_layers, num_experts]`` counts. + The update is mean-centred (as in Moonlight and torchtitan), which leaves the selection unchanged but keeps + the biases from drifting together.""" + for moe, load in zip(iter_moe_blocks(model), tokens_per_expert): + if moe.router.selection_bias is not None: + delta = torch.sign(load.mean() - load) + moe.router.selection_bias += rate * (delta - delta.mean()) + + def count_parameters(model: nn.Module) -> dict[str, int]: """Total, embedding (input embedding and LM head), non-embedding and active non-embedding parameter counts. Active counts the routed experts of each MoE layer at ``top_k / num_experts``, the share each token uses.""" @@ -247,7 +263,13 @@ def train(config: PretrainConfig): cp_size = parallel_dims.cp dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() is_moe_model = is_tt_moe_model(model) + if is_moe_model and model.config.model_type != "prime_moe": + raise ValueError(f"Pretraining MoE models use model_type 'prime_moe', got {model.config.model_type!r}") balance = config.model.moe.balance + for moe in iter_moe_blocks(model): + if moe.routed_k_counts is not None: + # Allocated uninitialized with the rest of the model. + moe.routed_k_counts.zero_() moes_with_seq_balance = list(iter_moe_blocks(model)) if balance.seq_aux_coeff > 0 else [] for moe in moes_with_seq_balance: # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. @@ -361,6 +383,8 @@ def run_suite(step: int) -> None: dead_experts = DeadExperts() loss_spikes, grad_norm_spikes = SpikeDetector(), SpikeDetector() last_bias_moves = None + # Per-expert selection bias, router-row norm and load of every MoE layer, every MOE_TRACE_INTERVAL steps. + moe_trace: dict[str, list] = {"step": [], "bias": [], "router_norm": [], "load": []} update_tracker = UpdateTracker() logger.info(f"Starting training loop (max_steps={config.max_steps})") while True: @@ -401,6 +425,20 @@ def run_suite(step: int) -> None: stacked = torch.stack(biases) moe_metrics["router/bias/absmax"] = stacked.abs().max().item() moe_metrics["router/bias/std"] = stacked.std(dim=1).mean().item() + if progress.step % MOE_TRACE_INTERVAL == 0: + gates = [ + moe.router.gate.weight + for moe in iter_moe_blocks(model) + if moe.router.selection_bias is not None + ] + gates = [gate.full_tensor() if isinstance(gate, DTensor) else gate for gate in gates] + moe_trace["step"].append(progress.step) + moe_trace["bias"].append(stacked.float().cpu().numpy()) + moe_trace["router_norm"].append( + torch.stack([gate.detach().float().norm(dim=1) for gate in gates]).cpu().numpy() + ) + moe_trace["load"].append(result.tokens_per_expert.float().cpu().numpy()) + moe_metrics |= routed_k_metrics(model, dp_cp_group) group_grad_norms = grad_norms_by_group(model) grad_norm = clip_grad_norm_(None, model, config.optim.max_norm) if config.optim.max_norm else None @@ -498,6 +536,8 @@ def run_suite(step: int) -> None: if heart is not None: heart.beat() + if moe_trace["step"] and world.rank == 0 and (is_last_step or progress.step % (16 * MOE_TRACE_INTERVAL) == 0): + np.savez(config.run_dir / "moe_trace.npz", **{name: np.array(values) for name, values in moe_trace.items()}) if is_last_step: break progress.step += 1 diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index b3021a39c62..e720b18483a 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -7,8 +7,8 @@ from prime_rl.configs.pretrain import PretrainAdamWConfig, PretrainMuonConfig from prime_rl.configs.trainer import InitConfig -from prime_rl.trainer.models.glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM from prime_rl.trainer.models.init import init_weights, normal_shard +from prime_rl.trainer.models.prime_moe import PrimeMoeConfig, PrimeMoeForCausalLM from prime_rl.trainer.pretrain.parametrization import param_groups @@ -91,7 +91,7 @@ def explicit_scales(name: str) -> tuple[float, float, float]: def test_completep_moe_groups_cover_experts_and_scale_the_router(): - config = Glm4MoeConfig( + config = PrimeMoeConfig( vocab_size=128, hidden_size=512, intermediate_size=1024, @@ -106,7 +106,7 @@ def test_completep_moe_groups_cover_experts_and_scale_the_router(): ) config._attn_implementation = "flash_attention_2" with torch.device("meta"): - model = Glm4MoeForCausalLM._from_config(config) + model = PrimeMoeForCausalLM._from_config(config) groups = param_groups(model, PretrainMuonConfig(lr=1.0, weight_decay=0.1), m_n=2.0, m_l=1.0, ep_enabled=False) lr = {param: group["lr"] for group in groups for param in group["params"]} assert sum(len(group["params"]) for group in groups) == len(lr) == len(list(model.parameters())) diff --git a/tests/unit/train/pretrain/test_moe_balance.py b/tests/unit/train/pretrain/test_moe_balance.py new file mode 100644 index 00000000000..ac71a015ca0 --- /dev/null +++ b/tests/unit/train/pretrain/test_moe_balance.py @@ -0,0 +1,30 @@ +import torch + +from prime_rl.trainer.models.layers.moe import MoE, MoEArgs +from prime_rl.trainer.models.prime_moe.moe import sequence_balance +from prime_rl.trainer.pretrain.train import update_selection_biases + + +def test_update_selection_biases_moves_towards_underloaded_experts(): + moe = MoE.from_args(MoEArgs(num_experts=4), dim=4, hidden_dim=8, shared_expert=None) + layer = torch.nn.Module() + layer.mlp = moe + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.layers = torch.nn.ModuleList([layer]) + + update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]]), rate=0.1) + update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]]), rate=0.1) + + # sign(mean - load) = [-1, 1, 1, 1], centred to [-1.5, 0.5, 0.5, 0.5] per step + torch.testing.assert_close(moe.router.selection_bias, torch.tensor([-0.3, 0.1, 0.1, 0.1])) + + +def test_sequence_balance(): + scores = torch.tensor([[0.9, 0.1], [0.9, 0.1], [0.9, 0.1], [0.1, 0.9]]) + selected = torch.tensor([[0], [0], [0], [1]]) + + # row 0 sends both tokens to expert 0: f = [2, 0], P = [0.9, 0.1]; row 1 is balanced: f = [1, 1], P = [0.5, 0.5] + torch.testing.assert_close(sequence_balance(scores, selected, num_rows=2), torch.tensor([1.8, 1.0])) + # the same tokens as one sequence (a packed sample of both rows) hide row 0's imbalance: f = [1.5, 0.5], P = [0.7, 0.3] + torch.testing.assert_close(sequence_balance(scores, selected, num_rows=1), torch.tensor([1.2])) From 27fdb0666b181a12143f16bdafcccaf83c48f000 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 09:15:28 +0200 Subject: [PATCH 079/132] refactor(pretrain): shared MoE layer, glm4_moe and model.py back to main The pretrain-only sequence balance, the extra router output, glm4_moe's head_dim default and the selection-bias update now live with prime_moe and the pretrain loop; SFT and RL run main's model code. --- src/prime_rl/trainer/model.py | 12 ------- .../models/glm4_moe/configuration_glm4_moe.py | 4 --- src/prime_rl/trainer/models/layers/moe.py | 35 ++----------------- tests/unit/train/models/test_moe.py | 28 +-------------- 4 files changed, 4 insertions(+), 75 deletions(-) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index e7a78147f39..d94e8c89780 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -333,18 +333,6 @@ def get_expert_load_stats(tokens_per_expert: Tensor, group: dist.ProcessGroup) - return {name: value.item() for name, value in compute_expert_load_stats(tokens_per_expert).items()} -@torch.no_grad() -def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor, rate: float) -> None: - """Aux-loss-free load balancing (https://arxiv.org/abs/2408.15664): move each layer's selection bias by - ``rate`` towards underloaded experts, given the step's global ``[num_moe_layers, num_experts]`` counts. - The update is mean-centred (as in Moonlight and torchtitan), which leaves the selection unchanged but keeps - the biases from drifting together.""" - for moe, load in zip(iter_moe_blocks(model), tokens_per_expert): - if moe.router.selection_bias is not None: - delta = torch.sign(load.mean() - load) - moe.router.selection_bias += rate * (delta - delta.mean()) - - def get_model( config: ModelConfig, device: torch.device = torch.device("cpu"), dtype: torch.dtype = torch.bfloat16 ) -> nn.Module: diff --git a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py index 30dce92cd34..2763679fa93 100644 --- a/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/configuration_glm4_moe.py @@ -29,8 +29,6 @@ class Glm4MoeConfig(PretrainedConfig): converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group. For more details, check out [this paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`. - head_dim (`int`, *optional*): - Dimension of each attention head. Defaults to `hidden_size // num_attention_heads`. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. max_position_embeddings (`int`, *optional*, defaults to 131072): @@ -146,7 +144,6 @@ def __init__( num_attention_heads=96, partial_rotary_factor=0.5, num_key_value_heads=8, - head_dim=None, hidden_act="silu", max_position_embeddings=131072, initializer_range=0.02, @@ -177,7 +174,6 @@ def __init__( self.partial_rotary_factor = partial_rotary_factor self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps diff --git a/src/prime_rl/trainer/models/layers/moe.py b/src/prime_rl/trainer/models/layers/moe.py index 11e90bfd4ce..f0494e435e4 100644 --- a/src/prime_rl/trainer/models/layers/moe.py +++ b/src/prime_rl/trainer/models/layers/moe.py @@ -180,14 +180,14 @@ def forward( self, x: torch.Tensor, routed_experts: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Args: x (torch.Tensor): Input tensor with shape ``(bs*slen, dim)``. routed_experts (torch.Tensor | None, optional): Optional tensor with shape ``(bs * slen, top_k)``. Returns: - tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - top_scores (torch.Tensor): Routing scores for selected experts with shape ``(bs*slen, top_k)``. - selected_experts_indices (torch.Tensor): @@ -196,8 +196,6 @@ def forward( Number of tokens assigned to each expert with shape ``(num_experts,)``. - routing_confidence_sum (torch.Tensor): Sum over tokens of the selected-expert probability mass before route normalization/scaling. - - scores (torch.Tensor): - Scores of all experts with shape ``(bs*slen, num_experts)``. """ # scores shape (bs*slen, num_experts) assert routed_experts is None or routed_experts.shape[-1] == self.top_k, ( @@ -268,7 +266,7 @@ def forward( max=self.num_experts, ).to(torch.int64) - return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, scores + return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum def init_weights(self, init_std: float): nn.init.trunc_normal_(self.gate.weight, mean=0.0, std=init_std) @@ -276,25 +274,6 @@ def init_weights(self, init_std: float): nn.init.zeros_(self.gate.bias) -def sequence_balance(scores: torch.Tensor, selected_experts: torch.Tensor, num_rows: int) -> torch.Tensor: - """DeepSeek-V3's sequence-wise balance loss without its coefficient, per row: ``sum_i f_i P_i`` with - ``f_i`` the row's share of expert ``i``'s selections times ``num_experts / top_k`` (1 when balanced) and - ``P_i`` its mean score normalized over all experts. 1 for a balanced row, up to ``num_experts / top_k``. - - Args: - scores: ``(num_rows * slen, num_experts)`` scores of all experts. - selected_experts: ``(num_rows * slen, top_k)`` selected experts. - """ - num_experts = scores.shape[1] - selections = selected_experts.view(num_rows, -1) - counts = torch.zeros(num_rows, num_experts, device=scores.device).scatter_add_( - 1, selections, torch.ones_like(selections, dtype=torch.float32) - ) - f = counts * (num_experts / selections.shape[1]) - p = (scores / scores.sum(dim=-1, keepdim=True)).view(num_rows, -1, num_experts).mean(dim=1) - return (f * p).sum(dim=-1) - - class MoE(nn.Module): """Token-choice MoE runtime composed from a router, grouped experts, and optional projections.""" @@ -365,10 +344,6 @@ def __init__( persistent=False, ) self.register_buffer("routing_confidence_sum", torch.tensor(0.0, dtype=torch.float32), persistent=False) - # Sequence-wise balance (DeepSeek-V3 eq. 17-20): with balance_rows set by the training loop, which knows - # how many packed rows a sample holds, every forward with grad computes the per-row balance for its loss. - self.balance_rows: int | None = None - self.seq_balance: torch.Tensor | None = None def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None: self.token_dispatcher = token_dispatcher @@ -408,12 +383,8 @@ def forward( selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, - scores, ) = self.router(x, routed_experts=routed_experts) - if self.balance_rows is not None and torch.is_grad_enabled(): - self.seq_balance = sequence_balance(scores, selected_experts_indices, self.balance_rows) - # Accumulate expert usage for selection-bias updates and metrics. with torch.no_grad(): record_moe_routing_statistics( diff --git a/tests/unit/train/models/test_moe.py b/tests/unit/train/models/test_moe.py index 5c12a8894e6..d34548de7bf 100644 --- a/tests/unit/train/models/test_moe.py +++ b/tests/unit/train/models/test_moe.py @@ -4,7 +4,7 @@ from prime_rl.configs.trainer import ModelConfig from prime_rl.trainer.distributed.token_dispatcher import LocalTokenDispatcher -from prime_rl.trainer.model import compute_expert_load_stats, is_tt_moe_model, update_selection_biases +from prime_rl.trainer.model import compute_expert_load_stats, is_tt_moe_model from prime_rl.trainer.models.fusions import fuse_gate_up_projections from prime_rl.trainer.models.layers.activations import ActivationDispatch, ClampedSilu from prime_rl.trainer.models.layers.expert_compute import BF16ExpertCompute, GroupedGemmExpertCompute @@ -13,7 +13,6 @@ GroupedExperts, MoE, MoEArgs, - sequence_balance, ) from prime_rl.trainer.models.qwen3_5 import ( Qwen3_5Config, @@ -255,28 +254,3 @@ def test_expert_load_stats(): assert stats.keys() == expected.keys() for name, value in expected.items(): assert stats[name].item() == pytest.approx(value) - - -def test_update_selection_biases_moves_towards_underloaded_experts(): - moe = MoE.from_args(MoEArgs(num_experts=4), dim=4, hidden_dim=8, shared_expert=None) - layer = torch.nn.Module() - layer.mlp = moe - model = torch.nn.Module() - model.model = torch.nn.Module() - model.model.layers = torch.nn.ModuleList([layer]) - - update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]]), rate=0.1) - update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]]), rate=0.1) - - # sign(mean - load) = [-1, 1, 1, 1], centred to [-1.5, 0.5, 0.5, 0.5] per step - torch.testing.assert_close(moe.router.selection_bias, torch.tensor([-0.3, 0.1, 0.1, 0.1])) - - -def test_sequence_balance(): - scores = torch.tensor([[0.9, 0.1], [0.9, 0.1], [0.9, 0.1], [0.1, 0.9]]) - selected = torch.tensor([[0], [0], [0], [1]]) - - # row 0 sends both tokens to expert 0: f = [2, 0], P = [0.9, 0.1]; row 1 is balanced: f = [1, 1], P = [0.5, 0.5] - torch.testing.assert_close(sequence_balance(scores, selected, num_rows=2), torch.tensor([1.8, 1.0])) - # the same tokens as one sequence (a packed sample of both rows) hide row 0's imbalance: f = [1.5, 0.5], P = [0.7, 0.3] - torch.testing.assert_close(sequence_balance(scores, selected, num_rows=1), torch.tensor([1.2])) From 27e911c07b71bf0e621acdcb3e7d2f153077427d Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 12:35:17 +0200 Subject: [PATCH 080/132] feat(pretrain): prime_moe options for a modern base: gated attention, XSA, LatentMoE, quantile balancing All default off (off = the previous code path). attn_output_gate: elementwise sigmoid gate on the attention output from the block input (Qiu et al. 2025; Qwen3-Next). exclusive_attention: XSA, each head's output without its component along the token's own value (Zhai 2026). latent_dim: LatentMoE, routed experts in a d -> latent projection with RMS norm, projected back up after the combine (Kimi K3, NVIDIA, Marin). router_bias_on_logits and [model.moe.balance] bias_update = "quantile": quantile balancing, the bias set every step from each expert's threshold (Su 2026; Kimi K3, Marin), null experts included. shared_expert_intermediate_size decouples the shared experts' hidden size from moe_intermediate_size (for parameter-matched variants). --- docs/pretraining.md | 5 ++ .../src/prime_rl/configs/pretrain.py | 11 ++- .../trainer/models/prime_moe/attention.py | 59 ++++++++++++++ .../prime_moe/configuration_prime_moe.py | 25 ++++++ .../models/prime_moe/modeling_prime_moe.py | 21 ++++- src/prime_rl/trainer/models/prime_moe/moe.py | 76 ++++++++++++++++--- src/prime_rl/trainer/pretrain/train.py | 27 ++++++- tests/unit/train/pretrain/test_moe_balance.py | 22 +++++- 8 files changed, 226 insertions(+), 20 deletions(-) create mode 100644 src/prime_rl/trainer/models/prime_moe/attention.py diff --git a/docs/pretraining.md b/docs/pretraining.md index 4db5699391e..2de17971c3e 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -21,6 +21,10 @@ name = "tokenizers/deepseek-v4-flash-base" # defaults to model.name MoE models use `prime_moe`, pretraining's own model (`prime_rl.trainer.models.prime_moe`), so architecture work does not touch the `glm4_moe` that SFT and RL use. It is GLM-4.5's architecture with the same config fields and state-dict layout, plus: - Gated Residual connections (Qwen3.8-Next): `hc_count` residual streams (above 1) with read gates of rank `hc_lowrank` replace the pre-norms and the final norm. Their projections train with AdamW. - Null experts (AdaMoE): `n_null_experts` zero-output experts the router can select next to the routed ones, so tokens use between 0 and `num_experts_per_tok` routed experts. Only the (token, routed expert) pairs are dispatched. The metrics are `router/k/*`. +- Exclusive self attention (XSA): `exclusive_attention` removes from each head's output its component along the token's own value vector. +- Gated attention (Qwen3-Next): `attn_output_gate` multiplies the attention output elementwise by `sigmoid(g_proj(x))`, computed from the block input, before `o_proj` (after XSA). +- LatentMoE: with `latent_dim`, the routed experts work in that width (projection down and RMS norm before the dispatch, projection up after the combine). The router and the shared experts stay full width. With `latent_dim = d/2`, doubling the routed experts, the null experts and `num_experts_per_tok` (and `routed_scaling_factor`, which keeps each expert's gate) keeps the expert parameters and routed FLOPs per token (Kimi K3, NVIDIA LatentMoE). +- `router_bias_on_logits` adds the selection bias to the router logits instead of the sigmoid scores (for quantile balancing). The trainer and the HF export (`tools/convert_dcp_to_bf16.py`) read the same model directory and tokenizer; vLLM does not support `prime_moe` yet. The trainer logs the active non-embedding (routed experts counted at `top_k / num_experts`), non-embedding, embedding (input embedding and LM head) and total parameter counts at startup, and as `params/*` at step 0. @@ -81,6 +85,7 @@ The schedule scales every param group's own learning rate, so ratios between gro | Field | Meaning | |---|---| +| `bias_update` | `sign` (default): the update below. `quantile`: quantile balancing. Every step, each layer's bias is set so that every expert would have been selected by `num_experts_per_tok / num_experts` of the step's tokens (null experts included), estimated per micro-batch and averaged over the micro-batches and ranks; no rate. Best with `router_bias_on_logits`. | | `bias_update_rate` | Every step, each layer's selection bias moves by this step size (gamma) towards underloaded experts, using the step's global expert counts. The update is mean-centred. Default 3e-3; 0 turns it off. The bias starts at 0. | | `bias_freeze` | Fraction of `max_steps` at the end during which the biases stop updating. Default 0.1. | | `seq_aux_coeff` | Weight (alpha) of the sequence-wise balance loss, computed per packed row of `data.seq_len` tokens and added to the training loss. Default 1e-3; 0 turns it off. | diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index b0d834e48d8..af05fd7d9b1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -38,9 +38,16 @@ class MoEBalanceConfig(BaseConfig): """MoE load balancing as in DeepSeek-V3: an aux-loss-free selection bias plus a small sequence-wise balance loss.""" + bias_update: Literal["sign", "quantile"] = "sign" + """``sign``: every step, each layer's bias moves by ``bias_update_rate`` towards underloaded experts. + ``quantile``: quantile balancing (Su 2026; Kimi K3, Marin): every step, each layer's bias is set to the value + with which every expert would have been selected by ``top_k / num_experts`` of the step's tokens (null experts + included), estimated per micro-batch and averaged over the micro-batches and ranks; no rate. Marin found it works + best on the router logits (``router_bias_on_logits`` in the model config).""" + bias_update_rate: float = Field(3e-3, ge=0) - """Step size of the selection-bias update (gamma): every step, each layer's bias moves by this towards - underloaded experts, mean-centred. 0 turns the update off.""" + """Step size of the sign update (gamma): every step, each layer's bias moves by this towards underloaded experts, + mean-centred. 0 turns the update off.""" bias_freeze: float = Field(0.1, ge=0, lt=1) """Fraction of ``max_steps`` at the end of training during which the biases stop updating, as DeepSeek-V3, diff --git a/src/prime_rl/trainer/models/prime_moe/attention.py b/src/prime_rl/trainer/models/prime_moe/attention.py new file mode 100644 index 00000000000..238596d9ed8 --- /dev/null +++ b/src/prime_rl/trainer/models/prime_moe/attention.py @@ -0,0 +1,59 @@ +import torch +import torch.nn.functional as F +from torch import nn + +from prime_rl.trainer.models.layers.attn import AttentionConfig, FlashAttention +from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb + + +class PrimeFlashAttention(FlashAttention): + """FlashAttention with two optional changes to the attention output, in this order: + + - ``exclusive`` (XSA; Zhai 2026, "Exclusive Self Attention"): remove from each head's output its component along the + token's own (normalized) value vector, so attention carries only context the residual path does not already have. + - ``output_gate`` (Qiu et al. 2025, "Gated Attention for Large Language Models", their best variant; Qwen3-Next): + multiply the output elementwise by ``sigmoid(g_proj(x))``, computed from the block input. + """ + + def __init__(self, config: AttentionConfig, flash_attn_version: int, *, exclusive: bool, output_gate: bool): + super().__init__(config, flash_attn_version=flash_attn_version) + self.exclusive = exclusive + self.g_proj = nn.Linear(config.hidden_size, self.qkv_sizes[0], bias=False) if output_gate else None + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + cu_seqlens: torch.LongTensor | None = None, + max_seqlen: int | None = None, + ) -> tuple[torch.Tensor, None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states, key_states, value_states = self.project_qkv(hidden_states) + if self.use_qk_norm and self.qk_norm_type == "per_layer": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + query_states = query_states.view(hidden_shape) + key_states = key_states.view(hidden_shape) + value_states = value_states.view(hidden_shape) + if self.use_qk_norm and self.qk_norm_type == "per_head": + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + if position_embeddings is not None: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb( + query_states.transpose(1, 2), key_states.transpose(1, 2), cos, sin + ) + query_states, key_states = query_states.transpose(1, 2), key_states.transpose(1, 2) + + out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) + if self.exclusive: + # Each query head's own value is its key/value group's. + values = F.normalize(value_states[0], dim=-1).repeat_interleave(self.num_key_value_groups, dim=1) + out = out - (out * values).sum(dim=-1, keepdim=True) * values + attn_output = out.contiguous().view(1, out.shape[0], -1) + if self.g_proj is not None: + attn_output = attn_output * torch.sigmoid(self.g_proj(hidden_states)) + return self.o_proj(attn_output), None diff --git a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py index 6b7d18e9eaf..da3c8df72fb 100644 --- a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py +++ b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py @@ -28,6 +28,8 @@ class PrimeMoeConfig(PretrainedConfig): moe_intermediate_size (`int`): Hidden size of each routed and shared expert. num_experts_per_tok (`int`): Experts each token selects (`top_k`), null experts included. n_shared_experts (`int`): Number of shared experts; their hidden sizes add up. + shared_expert_intermediate_size (`int`, *optional*): Hidden size of the shared experts together. Defaults to + `moe_intermediate_size * n_shared_experts`. n_routed_experts (`int`): Number of routed experts. routed_scaling_factor (`float`): Scale of the routed experts' output. first_k_dense_replace (`int`): Number of leading dense layers. @@ -38,6 +40,15 @@ class PrimeMoeConfig(PretrainedConfig): hc_lowrank (`int`, *optional*): Rank of the Gated Residual read gates; required when `hc_count` is above 1. n_null_experts (`int`): Zero-output experts the router can select next to the routed experts (AdaMoE); a token that selects them uses fewer routed experts. + attn_output_gate (`bool`): Elementwise sigmoid gate on the attention output, computed from the block input + (gated attention, Qwen3-Next). + exclusive_attention (`bool`): Exclusive self attention (XSA): remove from each head's output its component + along the token's own value vector. + latent_dim (`int`, *optional*): LatentMoE: the routed experts work in this width (`hidden_size -> latent_dim` + projection and RMS norm before them, `latent_dim -> hidden_size` projection after the combine), with hidden + size `moe_intermediate_size`. The router and the shared experts stay full width. + router_bias_on_logits (`bool`): Add the selection bias to the router logits instead of the sigmoid scores + (the space quantile balancing works best in). """ model_type = "prime_moe" @@ -67,12 +78,17 @@ def __init__( n_shared_experts=1, n_routed_experts=128, routed_scaling_factor=1.0, + shared_expert_intermediate_size=None, first_k_dense_replace=1, norm_topk_prob=True, use_qk_norm=False, hc_count=1, hc_lowrank=None, n_null_experts=0, + attn_output_gate=False, + exclusive_attention=False, + latent_dim=None, + router_bias_on_logits=False, pad_token_id=None, **kwargs, ): @@ -98,6 +114,11 @@ def __init__( self.moe_intermediate_size = moe_intermediate_size self.num_experts_per_tok = num_experts_per_tok self.n_shared_experts = n_shared_experts + self.shared_expert_intermediate_size = ( + shared_expert_intermediate_size + if shared_expert_intermediate_size is not None + else moe_intermediate_size * n_shared_experts + ) self.n_routed_experts = n_routed_experts self.routed_scaling_factor = routed_scaling_factor self.first_k_dense_replace = first_k_dense_replace @@ -107,5 +128,9 @@ def __init__( self.hc_count = hc_count self.hc_lowrank = hc_lowrank self.n_null_experts = n_null_experts + self.attn_output_gate = attn_output_gate + self.exclusive_attention = exclusive_attention + self.latent_dim = latent_dim + self.router_bias_on_logits = router_bias_on_logits super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) diff --git a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py index c8cd8ba6f55..ef8269fefc7 100644 --- a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py +++ b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py @@ -19,6 +19,7 @@ from prime_rl.trainer.models.layers.moe import GroupedExperts from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig +from prime_rl.trainer.models.prime_moe.attention import PrimeFlashAttention from prime_rl.trainer.models.prime_moe.configuration_prime_moe import PrimeMoeConfig from prime_rl.trainer.models.prime_moe.hyper_connection import HyperConnection from prime_rl.trainer.models.prime_moe.moe import PrimeMoE, PrimeRouter @@ -38,14 +39,22 @@ def __init__(self, config: PrimeMoeConfig, layer_idx: int): use_qk_norm=config.use_qk_norm, rms_norm_eps=config.rms_norm_eps, ) - self.self_attn = ATTN_IMPL2CLASS[config._attn_implementation](attn_config) + if config.attn_output_gate or config.exclusive_attention: + self.self_attn = PrimeFlashAttention( + attn_config, + flash_attn_version=int(config._attn_implementation.removeprefix("flash_attention_")), + exclusive=config.exclusive_attention, + output_gate=config.attn_output_gate, + ) + else: + self.self_attn = ATTN_IMPL2CLASS[config._attn_implementation](attn_config) if layer_idx >= config.first_k_dense_replace: shared_expert = None - if config.n_shared_experts > 0: + if config.shared_expert_intermediate_size > 0: shared_expert = FeedForward( dim=config.hidden_size, - hidden_dim=config.moe_intermediate_size * config.n_shared_experts, + hidden_dim=config.shared_expert_intermediate_size, expert_type="gated", activation=config.hidden_act, ) @@ -57,15 +66,19 @@ def __init__(self, config: PrimeMoeConfig, layer_idx: int): top_k=config.num_experts_per_tok, route_norm=config.norm_topk_prob, route_scale=config.routed_scaling_factor, + bias_on_logits=config.router_bias_on_logits, ), experts=GroupedExperts( - dim=config.hidden_size, + dim=config.latent_dim or config.hidden_size, hidden_dim=config.moe_intermediate_size, num_experts=config.n_routed_experts, expert_type="gated", activation=config.hidden_act, ), shared_expert=shared_expert, + dim=config.hidden_size, + latent_dim=config.latent_dim, + norm_eps=config.rms_norm_eps, ) else: self.mlp = FeedForward( diff --git a/src/prime_rl/trainer/models/prime_moe/moe.py b/src/prime_rl/trainer/models/prime_moe/moe.py index 4379461ca09..249bb5b979b 100644 --- a/src/prime_rl/trainer/models/prime_moe/moe.py +++ b/src/prime_rl/trainer/models/prime_moe/moe.py @@ -1,5 +1,6 @@ import torch import torch.nn.functional as F +from torch import nn from prime_rl.trainer.distributed.token_dispatcher import TokenDispatcher from prime_rl.trainer.models.layers.mlp import FeedForward @@ -9,15 +10,28 @@ TokenChoiceTopKRouter, record_moe_routing_statistics, ) +from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig class PrimeRouter(TokenChoiceTopKRouter): """Sigmoid-scored top-k router with a selection bias and ``num_null_experts`` zero-output experts next to the routed ones (AdaMoE). Null experts are scored, selected and balanced like the others, so a token that selects - them uses fewer routed experts. Returns the scores of all experts too, for the sequence-wise balance loss.""" + them uses fewer routed experts. Returns the scores of all experts too, for the sequence-wise balance loss. + + With ``bias_on_logits`` the selection bias is added to the logits instead of the sigmoid scores. With + ``quantile_balance`` set by the training loop, every forward with grad records each expert's quantile balancing + threshold (``quantile_sum``, ``quantile_count``), from which the loop sets the next step's bias.""" def __init__( - self, *, dim: int, num_experts: int, num_null_experts: int, top_k: int, route_norm: bool, route_scale: float + self, + *, + dim: int, + num_experts: int, + num_null_experts: int, + top_k: int, + route_norm: bool, + route_scale: float, + bias_on_logits: bool, ): super().__init__( dim=dim, @@ -29,6 +43,21 @@ def __init__( selection_bias=True, ) self.num_null_experts = num_null_experts + self.bias_on_logits = bias_on_logits + self.quantile_balance = False + self.register_buffer("quantile_sum", torch.zeros(self.num_experts, dtype=torch.float32), persistent=False) + self.register_buffer("quantile_count", torch.zeros((), dtype=torch.float32), persistent=False) + + @torch.no_grad() + def record_quantiles(self, selection_scores: torch.Tensor, biased_scores: torch.Tensor) -> None: + """Quantile balancing (Su 2026; Kimi K3, Marin): each token's threshold is its largest unselected biased score, + and each expert's new bias is minus the margin over that threshold that ``top_k / num_experts`` of the tokens + exceed, so that with it every expert would have been selected by exactly that share of these tokens.""" + thresholds = biased_scores.topk(self.top_k + 1, dim=1).values[:, -1:] + margins = selection_scores - thresholds + rank = max(1, round(margins.shape[0] * self.top_k / self.num_experts)) + self.quantile_sum += margins.topk(rank, dim=0).values[-1] + self.quantile_count += 1 def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Like ``TokenChoiceTopKRouter.forward`` with sigmoid scores; also returns the ``(bs*slen, num_experts)`` @@ -40,9 +69,11 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Te logits = self.gate(x) scores = torch.sigmoid(logits.float()) # The selection bias only affects which experts are selected, not their scores. - _, selected_experts_indices = torch.topk( - scores + self.selection_bias, k=self.top_k, dim=1, sorted=self.topk_sorted - ) + selection_scores = logits.float() if self.bias_on_logits else scores + biased_scores = selection_scores + self.selection_bias + _, selected_experts_indices = torch.topk(biased_scores, k=self.top_k, dim=1, sorted=self.topk_sorted) + if self.quantile_balance and torch.is_grad_enabled(): + self.record_quantiles(selection_scores, biased_scores) top_scores = scores.gather(dim=1, index=selected_experts_indices) with torch.no_grad(): @@ -80,9 +111,22 @@ def sequence_balance(scores: torch.Tensor, selected_experts: torch.Tensor, num_r class PrimeMoE(MoE): """The shared token-choice MoE runtime with a ``PrimeRouter``. With null experts, tokens use different numbers of routed experts: only the (token, routed expert) pairs are dispatched, one per row. Also computes the sequence-wise - balance for its loss (``balance_rows``) and counts the routed experts per token (``routed_k_counts``).""" + balance for its loss (``balance_rows``) and counts the routed experts per token (``routed_k_counts``). + + With ``latent_dim`` (LatentMoE), the routed experts work in that width: the token is projected down and RMS-normalized + before the dispatch, and the combined expert output is projected back up. The router and the shared experts read + the full-width token.""" - def __init__(self, *, router: PrimeRouter, experts: GroupedExperts, shared_expert: FeedForward | None) -> None: + def __init__( + self, + *, + router: PrimeRouter, + experts: GroupedExperts, + shared_expert: FeedForward | None, + dim: int, + latent_dim: int | None, + norm_eps: float, + ) -> None: super().__init__( router=router, experts=experts, @@ -106,6 +150,11 @@ def __init__(self, *, router: PrimeRouter, experts: GroupedExperts, shared_exper # how many packed rows a sample holds, every forward with grad computes the per-row balance for its loss. self.balance_rows: int | None = None self.seq_balance: torch.Tensor | None = None + self.latent_down = self.latent_norm = self.latent_up = None + if latent_dim is not None: + self.latent_down = nn.Linear(dim, latent_dim, bias=False) + self.latent_norm = RMSNorm(RMSNormConfig(hidden_size=latent_dim, eps=norm_eps)) + self.latent_up = nn.Linear(latent_dim, dim, bias=False) def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None: # Variable-k routing dispatches one (token, expert) pair per row. @@ -126,6 +175,7 @@ def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) - self.tokens_per_expert, self.routing_confidence_sum, num_tokens_per_expert, routing_confidence_sum ) + expert_input = x if self.latent_down is None else self.latent_norm(self.latent_down(x)) if self.variable_k: # Dispatch only the (token, routed expert) pairs, without the null experts. routed = selected_experts_indices < self.experts.num_experts @@ -134,24 +184,30 @@ def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) - self.routed_k_counts += torch.bincount(routed.sum(dim=1), minlength=self.routed_k_counts.numel()) token_indices, slots = routed.nonzero(as_tuple=True) pair_output = self.token_dispatcher.run( - x[token_indices], + expert_input[token_indices], top_scores[token_indices, slots].unsqueeze(1), selected_experts_indices[token_indices, slots].unsqueeze(1), self.experts, score_before_experts=self.score_before_experts, ) routed_output = ( - torch.zeros(x.shape, dtype=torch.float32, device=x.device) + torch.zeros(expert_input.shape, dtype=torch.float32, device=x.device) .index_add(0, token_indices, pair_output.float()) .to(x.dtype) ) else: routed_output = self.token_dispatcher.run( - x, top_scores, selected_experts_indices, self.experts, score_before_experts=self.score_before_experts + expert_input, + top_scores, + selected_experts_indices, + self.experts, + score_before_experts=self.score_before_experts, ) shared_output = self.shared_expert(x) if self.shared_expert is not None else None self.token_dispatcher.synchronize() + if self.latent_up is not None: + routed_output = self.latent_up(routed_output) if shared_output is not None: routed_output = routed_output + shared_output return routed_output.reshape(bs, slen, dim) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 93fdb1c1995..c32699d0407 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -96,6 +96,20 @@ def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor, rate: f moe.router.selection_bias += rate * (delta - delta.mean()) +@torch.no_grad() +def quantile_update_selection_biases(model: nn.Module, group: dist.ProcessGroup) -> None: + """Quantile balancing: set each layer's selection bias to minus the step's balancing thresholds of its experts + (``PrimeRouter.record_quantiles``), averaged over the micro-batches and ranks and mean-centred.""" + routers = [moe.router for moe in iter_moe_blocks(model)] + sums = torch.stack([torch.cat([router.quantile_sum, router.quantile_count[None]]) for router in routers]) + dist.all_reduce(sums, group=group) + for router, row in zip(routers, sums): + thresholds = row[:-1] / row[-1] + router.selection_bias.copy_(thresholds.mean() - thresholds) + router.quantile_sum.zero_() + router.quantile_count.zero_() + + def count_parameters(model: nn.Module) -> dict[str, int]: """Total, embedding (input embedding and LM head), non-embedding and active non-embedding parameter counts. Active counts the routed experts of each MoE layer at ``top_k / num_experts``, the share each token uses.""" @@ -267,9 +281,12 @@ def train(config: PretrainConfig): raise ValueError(f"Pretraining MoE models use model_type 'prime_moe', got {model.config.model_type!r}") balance = config.model.moe.balance for moe in iter_moe_blocks(model): + # Allocated uninitialized with the rest of the model. if moe.routed_k_counts is not None: - # Allocated uninitialized with the rest of the model. moe.routed_k_counts.zero_() + moe.router.quantile_sum.zero_() + moe.router.quantile_count.zero_() + moe.router.quantile_balance = balance.bias_update == "quantile" moes_with_seq_balance = list(iter_moe_blocks(model)) if balance.seq_aux_coeff > 0 else [] for moe in moes_with_seq_balance: # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. @@ -412,9 +429,13 @@ def run_suite(step: int) -> None: biases = [ moe.router.selection_bias for moe in iter_moe_blocks(model) if moe.router.selection_bias is not None ] - if balance.bias_update_rate > 0 and progress.step <= bias_freeze_step: + updates_bias = balance.bias_update == "quantile" or balance.bias_update_rate > 0 + if updates_bias and progress.step <= bias_freeze_step: previous_biases = [bias.clone() for bias in biases] - update_selection_biases(model, result.tokens_per_expert, balance.bias_update_rate) + if balance.bias_update == "quantile": + quantile_update_selection_biases(model, dp_cp_group) + else: + update_selection_biases(model, result.tokens_per_expert, balance.bias_update_rate) if biases: # Share of experts whose bias moved in the opposite direction to the previous step's move. moves = torch.stack([bias - previous for bias, previous in zip(biases, previous_biases)]).sign() diff --git a/tests/unit/train/pretrain/test_moe_balance.py b/tests/unit/train/pretrain/test_moe_balance.py index ac71a015ca0..bab38a58b71 100644 --- a/tests/unit/train/pretrain/test_moe_balance.py +++ b/tests/unit/train/pretrain/test_moe_balance.py @@ -1,7 +1,7 @@ import torch from prime_rl.trainer.models.layers.moe import MoE, MoEArgs -from prime_rl.trainer.models.prime_moe.moe import sequence_balance +from prime_rl.trainer.models.prime_moe.moe import PrimeRouter, sequence_balance from prime_rl.trainer.pretrain.train import update_selection_biases @@ -28,3 +28,23 @@ def test_sequence_balance(): torch.testing.assert_close(sequence_balance(scores, selected, num_rows=2), torch.tensor([1.8, 1.0])) # the same tokens as one sequence (a packed sample of both rows) hide row 0's imbalance: f = [1.5, 0.5], P = [0.7, 0.3] torch.testing.assert_close(sequence_balance(scores, selected, num_rows=1), torch.tensor([1.2])) + + +def test_quantile_balancing_balances_the_recorded_batch(): + torch.manual_seed(0) + router = PrimeRouter( + dim=16, num_experts=6, num_null_experts=2, top_k=2, route_norm=True, route_scale=1.0, bias_on_logits=True + ) + with torch.no_grad(): + router.gate.weight[0] += 0.5 # expert 0 is preferred + router.quantile_balance = True + x = torch.randn(400, 16) + counts = router(x)[2] + assert counts.max() / counts.float().mean() > 1.4 + + # One quantile update brings every expert, null experts included, close to 400 * 2 / 8 = 100 tokens. + thresholds = router.quantile_sum / router.quantile_count + router.selection_bias.copy_(thresholds.mean() - thresholds) + with torch.no_grad(): + counts = router(x)[2] + assert counts.max() / counts.float().mean() < 1.15 From 2769909d22e8fe197d30600c7e9b1ef700554561 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 12:35:17 +0200 Subject: [PATCH 081/132] test(pretrain): GPU integration test for prime_moe pretraining Tiny prime_moe with the base components (Gated Residual, null experts, gated attention, XSA, LatentMoE, quantile balancing) in FP8 on 2 GPUs, on Reverse-Text-SFT tokenized into one datatrove shard by the test (data and tokenizer the SFT tests already use): start for 20 steps, resume to 30; no error and the loss goes down in each. Runs in gpu_tests behind a pretrain path filter. --- .github/workflows/gpu_tests.yaml | 11 +- .../pretrain-moe/model/config.json | 39 +++++++ .../ci/integration/pretrain-moe/resume.toml | 45 ++++++++ .../ci/integration/pretrain-moe/start.toml | 42 +++++++ tests/integration/test_pretrain_moe.py | 106 ++++++++++++++++++ 5 files changed, 242 insertions(+), 1 deletion(-) create mode 100644 configs/ci/integration/pretrain-moe/model/config.json create mode 100644 configs/ci/integration/pretrain-moe/resume.toml create mode 100644 configs/ci/integration/pretrain-moe/start.toml create mode 100644 tests/integration/test_pretrain_moe.py diff --git a/.github/workflows/gpu_tests.yaml b/.github/workflows/gpu_tests.yaml index 94257465c60..f6394282380 100644 --- a/.github/workflows/gpu_tests.yaml +++ b/.github/workflows/gpu_tests.yaml @@ -69,6 +69,13 @@ jobs: - 'tests/integration/test_alphabet_sort.py' - 'configs/ci/integration/reverse-text{,-lora,-moe,-rl-opd,-rl-sft}/**' - 'configs/ci/integration/alphabet_sort.toml' + pretrain: + - *runtime + - 'packages/prime-rl-configs/src/prime_rl/configs/{pretrain,monitors}.py' + - 'src/prime_rl/entrypoints/pretrain.py' + - 'src/prime_rl/monitors/**' + - 'tests/integration/test_pretrain_moe.py' + - 'configs/ci/integration/pretrain-moe/**' eval: - *runtime - 'packages/prime-rl-configs/src/prime_rl/configs/{eval,orchestrator,inference,env_server,algorithm,monitors}.py' @@ -83,9 +90,11 @@ jobs: RUN_ALL: ${{ github.event_name != 'pull_request' }} SFT: ${{ steps.filter.outputs.sft || 'false' }} RL: ${{ steps.filter.outputs.rl || 'false' }} + PRETRAIN: ${{ steps.filter.outputs.pretrain || 'false' }} run: | - tests=$(jq -cn --argjson all "$RUN_ALL" --argjson sft "$SFT" --argjson rl "$RL" '[ + tests=$(jq -cn --argjson all "$RUN_ALL" --argjson sft "$SFT" --argjson rl "$RL" --argjson pretrain "$PRETRAIN" '[ if $all or $sft then "reverse_text_sft", "reverse_text_sft_lora" else empty end, + if $all or $pretrain then "pretrain_moe" else empty end, if $all or $rl then "reverse_text", "reverse_text_rl_opd", "reverse_text_rl_sft", "alphabet_sort", "reverse_text_lora", "reverse_text_moe" diff --git a/configs/ci/integration/pretrain-moe/model/config.json b/configs/ci/integration/pretrain-moe/model/config.json new file mode 100644 index 00000000000..178c5251779 --- /dev/null +++ b/configs/ci/integration/pretrain-moe/model/config.json @@ -0,0 +1,39 @@ +{ + "model_type": "prime_moe", + "vocab_size": 151680, + "bos_token_id": null, + "eos_token_id": 151645, + "pad_token_id": 151643, + "hidden_size": 256, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 64, + "intermediate_size": 1024, + "hidden_act": "silu", + "max_position_embeddings": 4096, + "rms_norm_eps": 1e-06, + "partial_rotary_factor": 1.0, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10000.0, + "partial_rotary_factor": 1.0 + }, + "tie_word_embeddings": false, + "attention_bias": false, + "use_qk_norm": true, + "first_k_dense_replace": 1, + "n_routed_experts": 16, + "n_null_experts": 8, + "num_experts_per_tok": 6, + "moe_intermediate_size": 128, + "n_shared_experts": 1, + "norm_topk_prob": true, + "routed_scaling_factor": 2.5, + "hc_count": 2, + "hc_lowrank": 32, + "attn_output_gate": true, + "exclusive_attention": true, + "latent_dim": 128, + "router_bias_on_logits": true +} diff --git a/configs/ci/integration/pretrain-moe/resume.toml b/configs/ci/integration/pretrain-moe/resume.toml new file mode 100644 index 00000000000..73893a0d83d --- /dev/null +++ b/configs/ci/integration/pretrain-moe/resume.toml @@ -0,0 +1,45 @@ +# Tiny prime_moe with the base components (Gated Residual, null experts, gated attention, XSA, LatentMoE, quantile +# balancing) in FP8 on 2 GPUs. The test writes the token shards (Reverse-Text-SFT, Qwen3 tokenizer) and passes their +# folder as --data.sources.reverse.path. Resumes the start run from its step-20 checkpoint. +max_steps = 30 + +[model] +name = "configs/ci/integration/pretrain-moe/model" +ac = "None" + +[model.quantization] +type = "fp8" + +[model.moe.compute] +type = "deepgemm_fp8" + +[model.moe.balance] +bias_update = "quantile" + +[tokenizer] +name = "PrimeIntellect/Qwen3-0.6B" + +[data] +type = "tokens" +seq_len = 2048 +batch_tokens = 16384 +micro_batch_tokens = 4096 +max_epochs = 100 + +[data.sources] +reverse = { path = "set-by-the-test" } + +[[data.phases]] +weights = { reverse = 1.0 } + +[optim] +type = "muon" +lr = 0.004 + +[scheduler] +type = "constant" + +[ckpt] + +[resume] +step = 20 diff --git a/configs/ci/integration/pretrain-moe/start.toml b/configs/ci/integration/pretrain-moe/start.toml new file mode 100644 index 00000000000..084a2bdfec8 --- /dev/null +++ b/configs/ci/integration/pretrain-moe/start.toml @@ -0,0 +1,42 @@ +# Tiny prime_moe with the base components (Gated Residual, null experts, gated attention, XSA, LatentMoE, quantile +# balancing) in FP8 on 2 GPUs. The test writes the token shards (Reverse-Text-SFT, Qwen3 tokenizer) and passes their +# folder as --data.sources.reverse.path. +max_steps = 20 + +[model] +name = "configs/ci/integration/pretrain-moe/model" +ac = "None" + +[model.quantization] +type = "fp8" + +[model.moe.compute] +type = "deepgemm_fp8" + +[model.moe.balance] +bias_update = "quantile" + +[tokenizer] +name = "PrimeIntellect/Qwen3-0.6B" + +[data] +type = "tokens" +seq_len = 2048 +batch_tokens = 16384 +micro_batch_tokens = 4096 +max_epochs = 100 + +[data.sources] +reverse = { path = "set-by-the-test" } + +[[data.phases]] +weights = { reverse = 1.0 } + +[optim] +type = "muon" +lr = 0.004 + +[scheduler] +type = "constant" + +[ckpt] diff --git a/tests/integration/test_pretrain_moe.py b/tests/integration/test_pretrain_moe.py new file mode 100644 index 00000000000..d94756f87fb --- /dev/null +++ b/tests/integration/test_pretrain_moe.py @@ -0,0 +1,106 @@ +from pathlib import Path +from typing import Callable + +import numpy as np +import pytest +from datasets import load_dataset +from transformers import AutoTokenizer + +from tests.conftest import ProcessResult +from tests.utils import check_loss_goes_down, strip_escape_codes + +pytestmark = [pytest.mark.slow, pytest.mark.gpu] + +RUN_NAME = "pretrain-moe" +TIMEOUT = 300 # 5 minutes + + +@pytest.fixture(scope="module") +def run_dir(output_dir: Path) -> Path: + return output_dir / RUN_NAME + + +@pytest.fixture(scope="module") +def wandb_name(branch_name: str) -> str: + return f"test-pretrain-moe:{branch_name}" + + +@pytest.fixture(scope="module") +def token_dir(output_dir: Path) -> Path: + """Reverse-Text-SFT as one datatrove token shard: every conversation one document, ending with EOS.""" + tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") + documents = [ + tokenizer.encode("\n".join(message["content"] for message in row["prompt"] + row["completion"])) + + [tokenizer.eos_token_id] + for row in load_dataset("PrimeIntellect/Reverse-Text-SFT", split="train") + ] + path = output_dir / "pretrain-tokens" + path.mkdir(parents=True, exist_ok=True) + np.array([token for document in documents for token in document], dtype=" list[str]: + return [ + "uv", + "run", + "pretrain", + "@", + f"configs/ci/integration/pretrain-moe/{config}", + "--data.sources.reverse.path", + token_dir.as_posix(), + "--deployment.num-train-gpus", + "2", + "--monitors.wandb.project", + wandb_project, + "--monitors.wandb.name", + wandb_name, + "--output-dir", + output_dir.as_posix(), + "--run.name", + RUN_NAME, + ] + + +@pytest.fixture(scope="module") +def pretrain_process( + run_process: Callable[..., ProcessResult], token_dir: Path, wandb_project: str, wandb_name: str, output_dir: Path +) -> ProcessResult: + cmd = pretrain_command("start.toml", token_dir, wandb_project, wandb_name, output_dir) + return run_process([*cmd, "--clean"], timeout=TIMEOUT) + + +@pytest.fixture(scope="module") +def pretrain_resume_process( + pretrain_process: ProcessResult, + run_process: Callable[..., ProcessResult], + token_dir: Path, + wandb_project: str, + wandb_name: str, + output_dir: Path, +) -> ProcessResult: + cmd = pretrain_command("resume.toml", token_dir, wandb_project, f"{wandb_name}-resume", output_dir) + return run_process(cmd, timeout=TIMEOUT) + + +def trainer_log(run_dir: Path) -> list[str]: + with open(run_dir / "logs" / "latest" / "trainer.log") as f: + return strip_escape_codes(f.read()).splitlines() + + +def test_no_error(pretrain_process: ProcessResult): + assert pretrain_process.returncode == 0, f"Process has non-zero return code ({pretrain_process})" + + +def test_loss_goes_down(pretrain_process: ProcessResult, run_dir: Path): + check_loss_goes_down(trainer_log(run_dir)) + + +def test_no_error_resume(pretrain_resume_process: ProcessResult): + assert pretrain_resume_process.returncode == 0, f"Process has non-zero return code ({pretrain_resume_process})" + + +def test_loss_goes_down_resume(pretrain_resume_process: ProcessResult, run_dir: Path): + check_loss_goes_down(trainer_log(run_dir)) From df9284733754047b79884d30ce44cfe09af13985 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 12:44:40 +0200 Subject: [PATCH 082/132] feat(pretrain): prime_moe options for Kimi K3's Stable LatentMoE: latent_norm and SiTU-GLU latent_norm = "output" puts LatentMoE's RMS norm on the combined expert output before the up-projection (K3 eq. 11) instead of on the projected expert input ("input", the default, Marin). expert_activation = "situ_glu" uses K3's Sigmoid Tanh Unit GLU (eq. 12, caps 4 and 25) in the routed and shared experts. Defaults unchanged. --- docs/pretraining.md | 3 +- .../prime_moe/configuration_prime_moe.py | 8 +++++ .../models/prime_moe/modeling_prime_moe.py | 8 +++-- src/prime_rl/trainer/models/prime_moe/moe.py | 30 ++++++++++++++++--- 4 files changed, 41 insertions(+), 8 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 2de17971c3e..f1c9f87ee9c 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -23,7 +23,8 @@ MoE models use `prime_moe`, pretraining's own model (`prime_rl.trainer.models.pr - Null experts (AdaMoE): `n_null_experts` zero-output experts the router can select next to the routed ones, so tokens use between 0 and `num_experts_per_tok` routed experts. Only the (token, routed expert) pairs are dispatched. The metrics are `router/k/*`. - Exclusive self attention (XSA): `exclusive_attention` removes from each head's output its component along the token's own value vector. - Gated attention (Qwen3-Next): `attn_output_gate` multiplies the attention output elementwise by `sigmoid(g_proj(x))`, computed from the block input, before `o_proj` (after XSA). -- LatentMoE: with `latent_dim`, the routed experts work in that width (projection down and RMS norm before the dispatch, projection up after the combine). The router and the shared experts stay full width. With `latent_dim = d/2`, doubling the routed experts, the null experts and `num_experts_per_tok` (and `routed_scaling_factor`, which keeps each expert's gate) keeps the expert parameters and routed FLOPs per token (Kimi K3, NVIDIA LatentMoE). +- LatentMoE: with `latent_dim`, the routed experts work in that width (projection down and RMS norm before the dispatch, projection up after the combine). The router and the shared experts stay full width. With `latent_dim = d/2`, doubling the routed experts, the null experts and `num_experts_per_tok` (and `routed_scaling_factor`, which keeps each expert's gate) keeps the expert parameters and routed FLOPs per token (Kimi K3, NVIDIA LatentMoE). `latent_norm` puts its RMS norm on the expert input (`"input"`, Marin) or on the combined expert output before the up-projection (`"output"`, Kimi K3's Stable LatentMoE). +- `expert_activation = "situ_glu"`: Kimi K3's SiTU-GLU in the routed and shared experts, SwiGLU with both factors soft-capped (`4 tanh(g/4) sigmoid(g) * 25 tanh(u/25)`). - `router_bias_on_logits` adds the selection bias to the router logits instead of the sigmoid scores (for quantile balancing). The trainer and the HF export (`tools/convert_dcp_to_bf16.py`) read the same model directory and tokenizer; vLLM does not support `prime_moe` yet. The trainer logs the active non-embedding (routed experts counted at `top_k / num_experts`), non-embedding, embedding (input embedding and LM head) and total parameter counts at startup, and as `params/*` at step 0. diff --git a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py index da3c8df72fb..5f57cd4584d 100644 --- a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py +++ b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py @@ -47,6 +47,10 @@ class PrimeMoeConfig(PretrainedConfig): latent_dim (`int`, *optional*): LatentMoE: the routed experts work in this width (`hidden_size -> latent_dim` projection and RMS norm before them, `latent_dim -> hidden_size` projection after the combine), with hidden size `moe_intermediate_size`. The router and the shared experts stay full width. + latent_norm (`str`): Where LatentMoE's RMS norm sits: `"input"` on the projected expert input (Marin) or + `"output"` on the combined expert output before the up-projection (Kimi K3's Stable LatentMoE). + expert_activation (`str`, *optional*): Activation of the routed and shared experts: `"situ_glu"` (Kimi K3's + soft-capped SwiGLU). Defaults to `hidden_act`. router_bias_on_logits (`bool`): Add the selection bias to the router logits instead of the sigmoid scores (the space quantile balancing works best in). """ @@ -88,6 +92,8 @@ def __init__( attn_output_gate=False, exclusive_attention=False, latent_dim=None, + latent_norm="input", + expert_activation=None, router_bias_on_logits=False, pad_token_id=None, **kwargs, @@ -131,6 +137,8 @@ def __init__( self.attn_output_gate = attn_output_gate self.exclusive_attention = exclusive_attention self.latent_dim = latent_dim + self.latent_norm = latent_norm + self.expert_activation = expert_activation self.router_bias_on_logits = router_bias_on_logits super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) diff --git a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py index ef8269fefc7..a4ad4a41613 100644 --- a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py +++ b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py @@ -22,7 +22,7 @@ from prime_rl.trainer.models.prime_moe.attention import PrimeFlashAttention from prime_rl.trainer.models.prime_moe.configuration_prime_moe import PrimeMoeConfig from prime_rl.trainer.models.prime_moe.hyper_connection import HyperConnection -from prime_rl.trainer.models.prime_moe.moe import PrimeMoE, PrimeRouter +from prime_rl.trainer.models.prime_moe.moe import PrimeMoE, PrimeRouter, SituGlu from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens @@ -50,13 +50,14 @@ def __init__(self, config: PrimeMoeConfig, layer_idx: int): self.self_attn = ATTN_IMPL2CLASS[config._attn_implementation](attn_config) if layer_idx >= config.first_k_dense_replace: + expert_activation = SituGlu() if config.expert_activation == "situ_glu" else config.hidden_act shared_expert = None if config.shared_expert_intermediate_size > 0: shared_expert = FeedForward( dim=config.hidden_size, hidden_dim=config.shared_expert_intermediate_size, expert_type="gated", - activation=config.hidden_act, + activation=expert_activation, ) self.mlp = PrimeMoE( router=PrimeRouter( @@ -73,11 +74,12 @@ def __init__(self, config: PrimeMoeConfig, layer_idx: int): hidden_dim=config.moe_intermediate_size, num_experts=config.n_routed_experts, expert_type="gated", - activation=config.hidden_act, + activation=expert_activation, ), shared_expert=shared_expert, dim=config.hidden_size, latent_dim=config.latent_dim, + latent_norm=config.latent_norm, norm_eps=config.rms_norm_eps, ) else: diff --git a/src/prime_rl/trainer/models/prime_moe/moe.py b/src/prime_rl/trainer/models/prime_moe/moe.py index 249bb5b979b..3ed046c6cff 100644 --- a/src/prime_rl/trainer/models/prime_moe/moe.py +++ b/src/prime_rl/trainer/models/prime_moe/moe.py @@ -1,8 +1,11 @@ +from typing import Literal + import torch import torch.nn.functional as F from torch import nn from prime_rl.trainer.distributed.token_dispatcher import TokenDispatcher +from prime_rl.trainer.models.layers.activations import Activation from prime_rl.trainer.models.layers.mlp import FeedForward from prime_rl.trainer.models.layers.moe import ( GroupedExperts, @@ -13,6 +16,16 @@ from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig +class SituGlu(Activation): + """Sigmoid Tanh Unit GLU (Kimi K3 report, eq. 12): SwiGLU with both factors soft-capped, + ``(4 tanh(gate / 4) * sigmoid(gate)) * (25 tanh(up / 25))``; close to SwiGLU near 0 and bounded by 100.""" + + @staticmethod + def apply(gate: torch.Tensor | None, up: torch.Tensor) -> torch.Tensor: + assert gate is not None + return 4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate) * (25.0 * torch.tanh(up / 25.0)) + + class PrimeRouter(TokenChoiceTopKRouter): """Sigmoid-scored top-k router with a selection bias and ``num_null_experts`` zero-output experts next to the routed ones (AdaMoE). Null experts are scored, selected and balanced like the others, so a token that selects @@ -113,9 +126,10 @@ class PrimeMoE(MoE): routed experts: only the (token, routed expert) pairs are dispatched, one per row. Also computes the sequence-wise balance for its loss (``balance_rows``) and counts the routed experts per token (``routed_k_counts``). - With ``latent_dim`` (LatentMoE), the routed experts work in that width: the token is projected down and RMS-normalized - before the dispatch, and the combined expert output is projected back up. The router and the shared experts read - the full-width token.""" + With ``latent_dim`` (LatentMoE), the routed experts work in that width: the token is projected down before the + dispatch and the combined expert output is projected back up, with an RMS norm on the expert input + (``latent_norm = "input"``, Marin) or on the combined output before the up-projection (``"output"``, Kimi K3's + Stable LatentMoE). The router and the shared experts read the full-width token.""" def __init__( self, @@ -125,6 +139,7 @@ def __init__( shared_expert: FeedForward | None, dim: int, latent_dim: int | None, + latent_norm: Literal["input", "output"], norm_eps: float, ) -> None: super().__init__( @@ -151,6 +166,7 @@ def __init__( self.balance_rows: int | None = None self.seq_balance: torch.Tensor | None = None self.latent_down = self.latent_norm = self.latent_up = None + self.norm_latent_output = latent_norm == "output" if latent_dim is not None: self.latent_down = nn.Linear(dim, latent_dim, bias=False) self.latent_norm = RMSNorm(RMSNormConfig(hidden_size=latent_dim, eps=norm_eps)) @@ -175,7 +191,11 @@ def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) - self.tokens_per_expert, self.routing_confidence_sum, num_tokens_per_expert, routing_confidence_sum ) - expert_input = x if self.latent_down is None else self.latent_norm(self.latent_down(x)) + expert_input = x + if self.latent_down is not None: + expert_input = self.latent_down(x) + if not self.norm_latent_output: + expert_input = self.latent_norm(expert_input) if self.variable_k: # Dispatch only the (token, routed expert) pairs, without the null experts. routed = selected_experts_indices < self.experts.num_experts @@ -207,6 +227,8 @@ def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) - shared_output = self.shared_expert(x) if self.shared_expert is not None else None self.token_dispatcher.synchronize() if self.latent_up is not None: + if self.norm_latent_output: + routed_output = self.latent_norm(routed_output) routed_output = self.latent_up(routed_output) if shared_output is not None: routed_output = routed_output + shared_output From cfc999c7db5f376101193a1d60221ab1b178caa0 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 13:50:18 +0200 Subject: [PATCH 083/132] docs(pretrain): entrypoint table: pretrain runs multi-node via [slurm] --- docs/training.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/training.md b/docs/training.md index 2a5b11080fa..04aa1e9b82b 100644 --- a/docs/training.md +++ b/docs/training.md @@ -35,7 +35,7 @@ This page covers everything you need to launch, observe, checkpoint, and recover |---|---|---| | `uv run rl` | Wraps the trainer, orchestrator, and inference server in one launch from a merged TOML. | The default for any RL run. Runs locally for single-node experiments; submits to SLURM for single- or multi-node when `[slurm]` is set (see [Scaling § SLURM](scaling.md#slurm)). | | `uv run sft` | Supervised fine-tuning on a HF dataset. | Launches torchrun internally; never call torchrun directly. | -| `uv run pretrain` | Pretraining from scratch on token shards. | Single node; see [Pretraining](pretraining.md). | +| `uv run pretrain` | Pretraining from scratch on token shards. | Runs locally on one node; submits to SLURM for single- or multi-node when `[slurm]` is set (see [Pretraining § Multi-node](pretraining.md#multi-node)). | | `uv run inference` | vLLM server. | Always use this entrypoint over `vllm serve` — it adds `/update_weights`, `/load_lora_adapter`, and `/init_broadcaster`. | | `uv run orchestrator` | Standalone orchestrator process. | Pair with a separately-launched inference server and one `env-server` per source. | | `uv run eval` | Multi-env evals against a live inference server. | One epoch per source, pinned (or adaptive) concurrency, cursor checkpoints + `--resume`, dashboard + optional platform upload; see [Eval](eval.md). | From 9fc940ce37f38fc9813445981587a45093343d34 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 13:51:35 +0200 Subject: [PATCH 084/132] docs(pretrain): entrypoint table: drop the deployment note, as for sft --- docs/training.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/training.md b/docs/training.md index 04aa1e9b82b..4a59c355056 100644 --- a/docs/training.md +++ b/docs/training.md @@ -35,7 +35,7 @@ This page covers everything you need to launch, observe, checkpoint, and recover |---|---|---| | `uv run rl` | Wraps the trainer, orchestrator, and inference server in one launch from a merged TOML. | The default for any RL run. Runs locally for single-node experiments; submits to SLURM for single- or multi-node when `[slurm]` is set (see [Scaling § SLURM](scaling.md#slurm)). | | `uv run sft` | Supervised fine-tuning on a HF dataset. | Launches torchrun internally; never call torchrun directly. | -| `uv run pretrain` | Pretraining from scratch on token shards. | Runs locally on one node; submits to SLURM for single- or multi-node when `[slurm]` is set (see [Pretraining § Multi-node](pretraining.md#multi-node)). | +| `uv run pretrain` | Pretraining from scratch on token shards. | See [Pretraining](pretraining.md). | | `uv run inference` | vLLM server. | Always use this entrypoint over `vllm serve` — it adds `/update_weights`, `/load_lora_adapter`, and `/init_broadcaster`. | | `uv run orchestrator` | Standalone orchestrator process. | Pair with a separately-launched inference server and one `env-server` per source. | | `uv run eval` | Multi-env evals against a live inference server. | One epoch per source, pinned (or adaptive) concurrency, cursor checkpoints + `--resume`, dashboard + optional platform upload; see [Eval](eval.md). | From e433e1d8e01dd5db2b725faf9800e2b685c5bc5b Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 13:53:19 +0200 Subject: [PATCH 085/132] docs(pretrain): match the other entrypoint pages: table of contents, Launch section, title-case headings; list the page in the overview --- docs/overview.md | 1 + docs/pretraining.md | 28 +++++++++++++++++++++++----- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/docs/overview.md b/docs/overview.md index 7df75dac171..78334758fd1 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -38,6 +38,7 @@ The `rl` entrypoint reads `examples/basic/reverse-text/rl.toml`, splits it into - **[Configuration](configuration.md)** — TOML composition, CLI overrides, dry-run. - **[Training](training.md)** — Launch and observe RL and SFT runs. +- **[Pretraining](pretraining.md)** — Train a model from scratch on token shards. - **[Inference](inference.md)** — vLLM-backed server (or fleet) holding the current policy. - **[Scaling](scaling.md)** — Single-GPU through multi-node clusters via FSDP / EP / CP and SLURM. - **[Algorithms](algorithms.md)** — Async semantics, loss / advantage / filter plugins, trajectory merging. diff --git a/docs/pretraining.md b/docs/pretraining.md index f1c9f87ee9c..bce9d088017 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -1,11 +1,29 @@ # Pretraining -`uv run pretrain` trains a model from scratch on pre-tokenized token shards. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, checkpointing and monitors as the SFT trainer. It runs locally on one node or, with `[slurm]`, as a SLURM job on one or more nodes (see [Multi-node](#multi-node)). +`uv run pretrain` trains a model from scratch on pre-tokenized token shards. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, checkpointing and monitors as the SFT trainer. + +## Table of Contents + +- [Launch](#launch) +- [Model](#model) +- [Initialization](#initialization) +- [Parametrization](#parametrization) +- [Learning-Rate Schedule](#learning-rate-schedule) +- [MoE Load Balancing](#moe-load-balancing) +- [Data](#data) +- [Monitoring](#monitoring) +- [Evals](#evals) +- [Multi-Node](#multi-node) + +## Launch ```bash -uv run pretrain @ examples/pretrain/prime_moe.toml +uv run pretrain @ examples/pretrain/prime_moe.toml # MoE example on token shards +uv run pretrain @ configs/debug/fake/pretrain.toml # 5-step smoke test on random tokens ``` +It runs locally on one node or, with `[slurm]`, as a SLURM job on one or more nodes (see [Multi-Node](#multi-node)). + ## Model The model is set up as in SFT and RL: `model.name` is an HF model id or a local model directory with a `config.json`. Its weights are not loaded. The model starts from a random init (below) or from a checkpoint. A new architecture is a directory with only a `config.json` for a model type in `prime_rl.trainer.models`; [`examples/pretrain/moe-d512`](../examples/pretrain/moe-d512/config.json) is the one in the example. `vocab_size` and the BOS, EOS and pad ids in the config must fit the tokenizer. The checked-in configs round `len(tokenizer)` up to a multiple of 128 and use the tokenizer's special ids, with EOS as the pad id if the tokenizer has none. @@ -66,7 +84,7 @@ In MoE models, the expert matrices follow the hidden and `down_proj` rules per e Use QK-norm (`"use_qk_norm": true` in the model config) with CompleteP. Without it, CompleteP's attention collapses one or two LR steps above the optimum (the norm gains get no weight decay), and the cliff moves down with width. With QK-norm, CompleteP and the standard parametrization reach the same loss, and CompleteP's optimal LR stays fixed across widths (MoE QK-norm sweep). The per-role groups need AdamW or Muon and `model.offload = "none"`. The standard parametrization is CompleteP with the base shape equal to the model shape. Keep `head_dim` and the ratios `intermediate_size / hidden_size` (and `moe_intermediate_size / hidden_size`) and `num_attention_heads / num_key_value_heads` fixed across shapes. -## Learning-rate schedule +## Learning-Rate Schedule Pretraining uses the shared `[scheduler]`. Its default here is the `linear` scheduler with a warmup over 10% of the steps and a linear decay to `min_lr = 0` over the remaining 90%. Warmup and decay take a number of steps or a fraction of `max_steps` (rounded down), and `decay_shape` is `linear`, `cosine` or `1-sqrt`: @@ -80,7 +98,7 @@ decay_shape = "1-sqrt" The schedule scales every param group's own learning rate, so ratios between groups hold. AdamW and Muon apply weight decay as `lr * weight_decay`, so the schedule also scales the decay. The scheduler state is checkpointed; resuming with `ckpt.skip_scheduler = true` rebuilds the schedule over the remaining steps. -## MoE load balancing +## MoE Load Balancing `[model.moe.balance]` balances the routed experts of `prime_moe` as in DeepSeek-V3: @@ -149,7 +167,7 @@ sets = ["fwedu-holdout", "lambada", "math"] seq_lens = [2048, 4096, 16384] ``` -## Multi-node +## Multi-Node `uv run pretrain` takes the same `[deployment]` and `[slurm]` blocks as `uv run sft` (see [Scaling § SLURM](scaling.md#slurm)) and submits the same bundled templates. The launcher writes the model dir on the submitting node, then submits one job in which every node runs `torchrun` on the pretrain trainer: From 926c13ad24eddbe4284602428f2b60c97989bd1e Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 14:55:04 +0200 Subject: [PATCH 086/132] refactor(pretrain): one data mixture (data.weights) instead of mixture phases TokensDataConfig takes one weights dict (non-negative, sum 1, known sources); MixturePhaseConfig and the phase schedule are gone. planned_epochs counts the run's tokens times each weight. --- .../ci/integration/pretrain-moe/resume.toml | 4 +- .../ci/integration/pretrain-moe/start.toml | 4 +- docs/pretraining.md | 10 +--- examples/pretrain/prime_moe.toml | 4 +- .../src/prime_rl/configs/pretrain.py | 30 +++--------- src/prime_rl/trainer/pretrain/data.py | 49 +++++++------------ tests/unit/train/pretrain/test_data.py | 21 +++----- 7 files changed, 37 insertions(+), 85 deletions(-) diff --git a/configs/ci/integration/pretrain-moe/resume.toml b/configs/ci/integration/pretrain-moe/resume.toml index 73893a0d83d..fba58f9aa1e 100644 --- a/configs/ci/integration/pretrain-moe/resume.toml +++ b/configs/ci/integration/pretrain-moe/resume.toml @@ -25,13 +25,11 @@ seq_len = 2048 batch_tokens = 16384 micro_batch_tokens = 4096 max_epochs = 100 +weights = { reverse = 1.0 } [data.sources] reverse = { path = "set-by-the-test" } -[[data.phases]] -weights = { reverse = 1.0 } - [optim] type = "muon" lr = 0.004 diff --git a/configs/ci/integration/pretrain-moe/start.toml b/configs/ci/integration/pretrain-moe/start.toml index 084a2bdfec8..85b22407f68 100644 --- a/configs/ci/integration/pretrain-moe/start.toml +++ b/configs/ci/integration/pretrain-moe/start.toml @@ -25,13 +25,11 @@ seq_len = 2048 batch_tokens = 16384 micro_batch_tokens = 4096 max_epochs = 100 +weights = { reverse = 1.0 } [data.sources] reverse = { path = "set-by-the-test" } -[[data.phases]] -weights = { reverse = 1.0 } - [optim] type = "muon" lr = 0.004 diff --git a/docs/pretraining.md b/docs/pretraining.md index bce9d088017..48fe1dabab3 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -115,7 +115,7 @@ The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` `data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. Each epoch reads every token of a source once: in blocks of 2^20 tokens in a random order and from a random start offset, both drawn from `data.seed`, the source and the epoch. Within an epoch, the rows of each group of 4,096 consecutive rows are read in a seeded random order, so the pieces of one long document spread over many steps. So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole source, whatever the order of the files. Documents continue across rows, and attention restarts at every document, every row and every block boundary. Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. -Each row comes from one source, drawn with the weights of the current mixture phase. Phases start at fractions of `max_steps`. At startup the trainer computes how many passes over each source the run makes and fails above `data.max_epochs`. +Each row comes from one source, drawn with the mixture `data.weights`. At startup the trainer computes how many passes over each source the run makes and fails above `data.max_epochs`. ```toml [data] @@ -124,17 +124,11 @@ seq_len = 4096 batch_tokens = 262144 # per optimizer step, over all GPUs micro_batch_tokens = 65536 # per forward pass on one GPU max_epochs = 3 +weights = { web = 0.8, code = 0.2 } [data.sources] web = { path = "data/web" } code = { path = "data/code" } - -[[data.phases]] -weights = { web = 0.8, code = 0.2 } - -[[data.phases]] -start = 0.9 -weights = { web = 0.5, code = 0.5 } ``` Batch sizes are in tokens, so changing `seq_len` keeps the tokens per step and the number of steps. Resume restores the data position on the same GPU layout and `data.num_workers`; resuming on a different layout is not supported yet. `data.type = "fake"` trains on random tokens, for smoke and throughput tests. diff --git a/examples/pretrain/prime_moe.toml b/examples/pretrain/prime_moe.toml index bab6151cb97..67b99354bd0 100644 --- a/examples/pretrain/prime_moe.toml +++ b/examples/pretrain/prime_moe.toml @@ -20,6 +20,7 @@ seq_len = 4096 batch_tokens = 262144 micro_batch_tokens = 65536 max_epochs = 3 +weights = { web = 0.60, pdfs = 0.20, code = 0.12, math = 0.05, reference = 0.03 } [data.sources] web = { path = "data/dolma3-50b-dsv4/web" } @@ -28,9 +29,6 @@ code = { path = "data/dolma3-50b-dsv4/code" } math = { path = "data/dolma3-50b-dsv4/math" } reference = { path = "data/dolma3-50b-dsv4/reference" } -[[data.phases]] -weights = { web = 0.60, pdfs = 0.20, code = 0.12, math = 0.05, reference = 0.03 } - [optim] type = "muon" lr = 0.0078125 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index af05fd7d9b1..bdca39efbf9 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -82,20 +82,6 @@ class TokenSourceConfig(BaseConfig): """Tokens available from this source: the first ``max_tokens`` of its stream. Defaults to all tokens in the shards.""" -class MixturePhaseConfig(BaseConfig): - start: float = Field(0.0, ge=0, lt=1) - """Fraction of ``max_steps`` at which this phase starts.""" - - weights: dict[str, float] - """Token share of each source in this phase. Shares are non-negative and sum to 1; unlisted sources get 0.""" - - @model_validator(mode="after") - def validate_weights(self): - if any(weight < 0 for weight in self.weights.values()) or not math.isclose(sum(self.weights.values()), 1.0): - raise ValueError(f"Phase weights must be non-negative and sum to 1, got {self.weights}") - return self - - class BaseBatchConfig(BaseConfig): seq_len: int = Field(4096, ge=1) """Tokens per row. Documents continue across rows; attention restarts at every document and every row.""" @@ -145,20 +131,18 @@ class TokensDataConfig(BaseBatchConfig): sources: dict[str, TokenSourceConfig] """Token sources by name. The name labels the source's metrics.""" - phases: list[MixturePhaseConfig] - """Mixture phases in order. The first starts at 0.""" + weights: dict[str, float] + """Token share of each source. Shares are non-negative and sum to 1; unlisted sources get 0.""" max_epochs: float = Field(1.0, gt=0) """Most passes over a source's available tokens that the run (``max_steps``) may make. Checked at startup.""" @model_validator(mode="after") - def validate_phases(self): - starts = [phase.start for phase in self.phases] - if not starts or starts[0] != 0 or starts != sorted(set(starts)): - raise ValueError(f"Phase starts must begin at 0 and increase, got {starts}") - for phase in self.phases: - if unknown := phase.weights.keys() - self.sources.keys(): - raise ValueError(f"Phase weights name unknown sources {sorted(unknown)}") + def validate_weights(self): + if any(weight < 0 for weight in self.weights.values()) or not math.isclose(sum(self.weights.values()), 1.0): + raise ValueError(f"data.weights must be non-negative and sum to 1, got {self.weights}") + if unknown := self.weights.keys() - self.sources.keys(): + raise ValueError(f"data.weights name unknown sources {sorted(unknown)}") return self diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index 0d60d68cc94..138aa50e00e 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -1,4 +1,3 @@ -import bisect import zlib from pathlib import Path @@ -6,7 +5,7 @@ import torch from torchdata.stateful_dataloader import StatefulDataLoader -from prime_rl.configs.pretrain import MixturePhaseConfig, PretrainDataConfig, TokenSourceConfig +from prime_rl.configs.pretrain import PretrainDataConfig, TokenSourceConfig from prime_rl.trainer.sft.data import StatefulIterableDataset from prime_rl.utils.logger import get_logger @@ -141,8 +140,7 @@ def window_start(self, position: int, seq_len: int) -> int: class TokenDataset(StatefulIterableDataset): """``seq_len`` windows of token sources, each window from one source of a weighted mixture. - Window ``w`` belongs to step ``w // batch_size``; its source is drawn from that step's phase weights - with an RNG keyed by ``(seed, w)``, and it continues the source's stream where that source's last + Window ``w``'s source is drawn from the mixture ``weights`` with an RNG keyed by ``(seed, w)``, and it continues the source's stream where that source's last window ended. Every data rank walks the whole window sequence but reads only its own windows; the state is the window counter and the tokens read per source. Each window reads one token past its end for the last target. With ``num_windows`` the dataset stops @@ -151,10 +149,8 @@ class TokenDataset(StatefulIterableDataset): def __init__( self, sources: dict[str, TokenShards | RandomTokens], - phases: list[MixturePhaseConfig], + weights: dict[str, float], seq_len: int, - batch_size: int, - max_steps: int, seed: int = 0, non_dp_size: int = 1, num_windows: int | None = None, @@ -162,10 +158,9 @@ def __init__( super().__init__(non_dp_size) self.sources = sources self.names = list(sources) - self.phase_starts = [round(phase.start * max_steps) for phase in phases] - self.phase_cdfs = [np.cumsum([phase.weights.get(name, 0.0) for name in self.names]) for phase in phases] + self.weights = np.array([weights.get(name, 0.0) for name in self.names]) + self.cdf = np.cumsum(self.weights) self.seq_len = seq_len - self.batch_size = batch_size self.seed = seed self.num_windows = num_windows self.tokens = dict.fromkeys(self.names, 0) @@ -180,23 +175,18 @@ def load_state_dict(self, state_dict: dict): self.step = state_dict["dataset"]["step"] self.tokens.update(state_dict["dataset"]["tokens"]) - def planned_epochs(self, max_steps: int) -> dict[str, float]: - """Expected passes over each source's tokens in ``max_steps`` steps.""" - ends = [*self.phase_starts[1:], max_steps] - step_tokens = self.batch_size * self.seq_len - tokens = sum( - (end - start) * step_tokens * np.diff(cdf, prepend=0.0) - for start, end, cdf in zip(self.phase_starts, ends, self.phase_cdfs) - ) - return {name: float(tokens[i]) / self.sources[name].num_tokens for i, name in enumerate(self.names)} + def planned_epochs(self, num_tokens: int) -> dict[str, float]: + """Expected passes over each source's tokens when ``num_tokens`` tokens are read.""" + return { + name: num_tokens * float(weight) / self.sources[name].num_tokens + for name, weight in zip(self.names, self.weights) + } def source(self, window: int) -> str: if len(self.names) == 1: return self.names[0] - phase = bisect.bisect_right(self.phase_starts, window // self.batch_size) - 1 - cdf = self.phase_cdfs[phase] - draw = np.random.default_rng([self.seed, window]).random() * cdf[-1] - return self.names[min(int(np.searchsorted(cdf, draw, side="right")), len(self.names) - 1)] + draw = np.random.default_rng([self.seed, window]).random() * self.cdf[-1] + return self.names[min(int(np.searchsorted(self.cdf, draw, side="right")), len(self.names) - 1)] def __iter__(self): self._setup_world_info() @@ -251,16 +241,14 @@ def setup_train_dataset( ) -> TokenDataset: if config.type == "fake": sources = {"fake": RandomTokens(vocab_size, config.seed)} - phases = [MixturePhaseConfig(weights={"fake": 1.0})] + weights = {"fake": 1.0} else: sources = {name: TokenShards(source, seed=(config.seed, name)) for name, source in config.sources.items()} check_eos(sources, eos_token_id) - phases = config.phases - dataset = TokenDataset( - sources, phases, config.seq_len, config.batch_size, max_steps, seed=config.seed, non_dp_size=non_dp_size - ) + weights = config.weights + dataset = TokenDataset(sources, weights, config.seq_len, seed=config.seed, non_dp_size=non_dp_size) if config.type == "tokens": - epochs = dataset.planned_epochs(max_steps) + epochs = dataset.planned_epochs(max_steps * config.batch_tokens) get_logger().info(f"Planned epochs per source: {epochs}") if over := {name: epoch for name, epoch in epochs.items() if epoch > config.max_epochs}: raise ValueError(f"Planned epochs {over} exceed data.max_epochs={config.max_epochs}") @@ -269,8 +257,7 @@ def setup_train_dataset( def setup_val_dataset(source: TokenShards, seq_len: int, num_windows: int, non_dp_size: int) -> TokenDataset: """One pass over the first ``num_windows`` windows of a held-out source.""" - phases = [MixturePhaseConfig(weights={"val": 1.0})] - return TokenDataset({"val": source}, phases, seq_len, 1, 1, non_dp_size=non_dp_size, num_windows=num_windows) + return TokenDataset({"val": source}, {"val": 1.0}, seq_len, non_dp_size=non_dp_size, num_windows=num_windows) def check_eos(sources: dict[str, TokenShards], eos_token_id: int) -> None: diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index 127240f5ec4..fc93bfffc6a 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -3,7 +3,7 @@ import numpy as np -from prime_rl.configs.pretrain import MixturePhaseConfig, TokenSourceConfig +from prime_rl.configs.pretrain import TokenSourceConfig from prime_rl.trainer.pretrain import data from prime_rl.trainer.pretrain.data import TokenDataset, TokenShards, collate_rows, source_tokens @@ -57,23 +57,18 @@ def test_window_shuffle_reads_every_window_of_an_epoch_once(tmp_path, monkeypatc assert sorted(starts[:10]) == list(range(0, 70, 7)) and sorted(starts[10:]) == list(range(70, 140, 7)) -def test_token_dataset_resumes_and_follows_phases(tmp_path): +def test_token_dataset_resumes_and_continues_each_source(tmp_path): write_shard(tmp_path / "a" / "00.ds", [list(range(100, 200))]) write_shard(tmp_path / "b" / "00.ds", [list(range(200, 300))]) sources = {name: TokenShards(TokenSourceConfig(path=tmp_path / name)) for name in "ab"} - phases = [MixturePhaseConfig(weights={"a": 1.0}), MixturePhaseConfig(start=0.5, weights={"b": 1.0})] def dataset(): - return TokenDataset(sources, phases, seq_len=4, batch_size=2, max_steps=4) + return TokenDataset(sources, {"a": 0.5, "b": 0.5}, seq_len=4) full = [window["tokens"].tolist() for window in islice(dataset(), 8)] - assert full[:4] == [ - [100, 101, 102, 103, 104], - [104, 105, 106, 107, 108], - [108, 109, 110, 111, 112], - [112, 113, 114, 115, 116], - ] - assert full[4][0] == 200 + for start in (100, 200): + windows = [window for window in full if start <= window[0] < start + 100] + assert windows and windows == [list(range(start + 4 * i, start + 4 * i + 5)) for i in range(len(windows))] first = dataset() iterator = iter(first) @@ -88,9 +83,7 @@ def test_source_tokens_count_every_window_of_the_mixture(tmp_path): for name in "ab": write_shard(tmp_path / name / "00.ds", [list(range(100))]) sources = {name: TokenShards(TokenSourceConfig(path=tmp_path / name)) for name in "ab"} - dataset = TokenDataset( - sources, [MixturePhaseConfig(weights={"a": 0.3, "b": 0.7})], seq_len=4, batch_size=8, max_steps=4 - ) + dataset = TokenDataset(sources, {"a": 0.3, "b": 0.7}, seq_len=4) list(islice(dataset, 32)) # The training loop's global per-source counts equal the tokens the dataset itself read. assert source_tokens(dataset, 0, 32) == dataset.tokens From 8a7d83d102b881f10930cfd79d5e178967d91322 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 14:56:25 +0200 Subject: [PATCH 087/132] refactor(pretrain): remove the diagnostics step DiagnosticsConfig, [diagnostics] and trainer/pretrain/diagnostics.py (update sizes, attention and activation probe) are gone; the per-step metrics stay. --- docs/pretraining.md | 2 - .../src/prime_rl/configs/pretrain.py | 17 -- src/prime_rl/trainer/pretrain/diagnostics.py | 162 ------------------ src/prime_rl/trainer/pretrain/train.py | 18 -- tests/unit/train/pretrain/test_diagnostics.py | 22 --- 5 files changed, 221 deletions(-) delete mode 100644 src/prime_rl/trainer/pretrain/diagnostics.py delete mode 100644 tests/unit/train/pretrain/test_diagnostics.py diff --git a/docs/pretraining.md b/docs/pretraining.md index 48fe1dabab3..6b527492f4d 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -137,8 +137,6 @@ Batch sizes are in tokens, so changing `seq_len` keeps the tokens per step and t Every step logs the loss, learning rate, gradient norms (total and per parameter group, `optim/grad_norm/`), loss and grad-norm spike flags, per-source token counts and epochs (`progress//*`), throughput, MFU, data wait and peak memory (max and mean over ranks). MoE models add the expert balance of the step's global counts (`expert_load/*`, also per layer), `router/*` and `loss/seq_balance`. The step statistics are reduced once per step. -`[diagnostics]` (default: every 100 steps; `diagnostics = null` turns it off) adds the size of the optimizer update relative to the weights per parameter group (`update/rms_ratio/*`, `param/rms/*`) and a forward pass on a fixed probe batch: attention max logit and entropy per layer (`attn/*`, FlashAttention layers) and the residual-stream size per layer (`act/*`). It does not support context parallelism yet. - ## Evals `[val]` computes the held-out loss (nats per token), perplexity and bits per UTF-8 byte (comparable across tokenizers) of every source in `val.sources`, through the same loss code as training, so it also works with context parallelism. Each eval reads the same leading windows of every source: as many as fill whole micro-batches on every data-parallel rank. Evals run after the optimizer update, so the metrics at step N describe the same weights as the step-N checkpoint. With `eval_on_start` (default), they also run before the first step (logged as step 0), and they always run at the last step. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index bdca39efbf9..ff4e04f282d 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -163,18 +163,6 @@ class ValConfig(BaseConfig): windows: as many as fill whole micro-batches on every data-parallel rank.""" -class DiagnosticsConfig(BaseConfig): - """Diagnostic step: on every ``interval``-th step, the optimizer update's size per parameter group and a - forward pass on a fixed probe batch (the first ``probe_tokens`` tokens of each rank's first micro-batch) - that records attention max logits and entropy per head and the residual-stream size per layer.""" - - interval: int = Field(100, ge=1) - """Run the diagnostics every N steps.""" - - probe_tokens: int = Field(2048, ge=64) - """Tokens per rank in the probe batch. Fixed across sequence lengths so attention entropy stays comparable.""" - - class EvalsConfig(BaseConfig): """The eval suite, scored inside the trainer: forward passes without gradients on the training model, every ``interval`` steps. The sets are log-likelihood evals in the format of the pretrain-evals offline suite.""" @@ -239,9 +227,6 @@ class PretrainConfig(BaseConfig): val: ValConfig | None = None """Held-out loss on token shards (``val/*``). If None, it never runs.""" - diagnostics: DiagnosticsConfig | None = DiagnosticsConfig() - """Diagnostic step (update sizes, attention and activation probe). If None, it never runs.""" - evals: EvalsConfig | None = None """The eval suite, scored inside the trainer. If None, it never runs.""" @@ -381,8 +366,6 @@ def validate_cp(self): raise ValueError("data.seq_len must be divisible by model.cp") if self.data.micro_batch_size != 1: raise ValueError("data.micro_batch_tokens must equal data.seq_len with context parallelism") - if self.diagnostics is not None: - raise ValueError("[diagnostics] does not support context parallelism yet; set diagnostics = null") if self.evals is not None: raise ValueError("[evals] does not support context parallelism yet") return self diff --git a/src/prime_rl/trainer/pretrain/diagnostics.py b/src/prime_rl/trainer/pretrain/diagnostics.py deleted file mode 100644 index fdfaa3f5023..00000000000 --- a/src/prime_rl/trainer/pretrain/diagnostics.py +++ /dev/null @@ -1,162 +0,0 @@ -from functools import cache - -import torch -from torch import Tensor, nn - -from prime_rl.trainer.model import forward -from prime_rl.trainer.models.layers.attn import FlashAttention -from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX -from prime_rl.trainer.pretrain.metrics import group_norms, param_group, reduce_stats -from prime_rl.utils.logger import get_logger -from prime_rl.utils.vlm import get_language_model - -MIN_VISIBLE_KEYS = 64 -"""Attention entropy is averaged over queries that see at least this many keys.""" - -COLLAPSED_ENTROPY = 0.5 -"""A head whose mean attention entropy (nats) is below this counts as collapsed.""" - - -def probe_batch(micro_batch: dict[str, Tensor], num_tokens: int) -> dict[str, Tensor]: - """The first ``num_tokens`` tokens of a training micro-batch, on the GPU: the fixed batch of the diagnostic - forward.""" - lengths, total = [], 0 - for length in micro_batch["seq_lens"].tolist(): - lengths.append(min(length, num_tokens - total)) - total += lengths[-1] - if total == num_tokens: - break - batch = {key: micro_batch[key][:, :total].to("cuda") for key in ("input_ids", "position_ids", "target_ids")} - batch["seq_lens"] = torch.tensor(lengths, device="cuda") - return batch - - -def attention_stats(q: Tensor, k: Tensor, cu_seqlens: Tensor, scale: float) -> tuple[Tensor, Tensor, Tensor]: - """Per head: max attention logit (q·k times the softmax scale, after RoPE and QK-norm, over causal pairs - within a document), summed attention entropy (nats) over the queries that see at least `MIN_VISIBLE_KEYS` - keys, and the number of those queries. ``q`` is `[tokens, heads, head_dim]`, ``k`` `[tokens, kv_heads, head_dim]`.""" - num_tokens, num_heads, _ = q.shape - positions = torch.arange(num_tokens, device=q.device) - document = torch.searchsorted(cu_seqlens[1:].long(), positions, right=True) - mask = (positions[None, :] <= positions[:, None]) & (document[None, :] == document[:, None]) - valid = mask.sum(dim=1) >= MIN_VISIBLE_KEYS - groups = num_heads // k.shape[1] - max_logits, entropy_sums = [], [] - for head in range(num_heads): - scores = (q[:, head].float() @ k[:, head // groups].float().T * scale).masked_fill(~mask, float("-inf")) - max_logits.append(scores.amax()) - entropy = torch.special.entr(scores.softmax(dim=-1)).sum(dim=-1) - entropy_sums.append(entropy[valid].sum()) - return torch.stack(max_logits), torch.stack(entropy_sums), valid.sum() - - -@cache -def warn_no_attention_stats() -> None: - get_logger().warning( - "The diagnostic probe records attention statistics only for FlashAttention layers; none found, skipping them" - ) - - -@torch.no_grad() -def run_probe(model: nn.Module, batch: dict[str, Tensor], group) -> dict[str, float]: - """Forward the probe batch with hooks: attention max logit and entropy per layer and head (FlashAttention - layers), residual-stream RMS and max per layer, and the probe loss. Compiled layers run eagerly for the probe, - so the hooks see every call. Every rank probes its own batch; the statistics are reduced over ``group``.""" - layers = get_language_model(model).layers - attentions = [module for module in layers.modules() if isinstance(module, FlashAttention)] - if not attentions: - warn_no_attention_stats() - attention, activations = [], [] - - def capture(module: FlashAttention): - compute = module._compute_attention - - def compute_and_record(q, k, v, cu_seqlens, max_seqlen): - attention.append(attention_stats(q, k, cu_seqlens, module.scaling)) - return compute(q, k, v, cu_seqlens, max_seqlen) - - return compute_and_record - - def record_activations(module, args, output): - hidden = (output[0] if isinstance(output, tuple) else output).float() - activations.append( - (hidden.square().sum(), torch.tensor(hidden.numel(), device=hidden.device), hidden.abs().amax()) - ) - - for module in attentions: - module._compute_attention = capture(module) - hooks = [layer.register_forward_hook(record_activations) for layer in layers] - labels = batch["target_ids"] - with torch.compiler.set_stance("force_eager"): - out = forward(model, batch["input_ids"], batch["position_ids"], seq_lens=batch["seq_lens"], labels=labels) - for module in attentions: - del module._compute_attention - for hook in hooks: - hook.remove() - - square_sum, count, absmax = (torch.stack(values) for values in zip(*activations)) - sums = { - "square_sum": square_sum, - "count": count.double(), - "loss": out["loss"].double(), - "tokens": (labels != IGNORE_INDEX).sum().double(), - } - maxes = {"absmax": absmax} - if attention: - max_logit, entropy_sum, queries = (torch.stack(values) for values in zip(*attention)) - sums |= {"entropy_sum": entropy_sum, "queries": queries.double()} - maxes["max_logit"] = max_logit - stats = reduce_stats(sums, maxes, group) - - rms = (stats["square_sum"] / stats["count"]).sqrt() - metrics = { - "act/rms/max": rms.max().item(), - "act/absmax/max": stats["absmax"].max().item(), - "probe/loss": (stats["loss"] / stats["tokens"]).item(), - } - for layer in range(len(rms)): - metrics[f"act/rms/layer_{layer}"] = rms[layer].item() - metrics[f"act/absmax/layer_{layer}"] = stats["absmax"][layer].item() - if attention: - head_entropy = stats["entropy_sum"] / stats["queries"][:, None] # [layers, heads] - layer_max_logit = stats["max_logit"].amax(dim=1) - layer_entropy = head_entropy.mean(dim=1) - metrics |= { - "attn/max_logit/max": layer_max_logit.max().item(), - "attn/entropy/min": layer_entropy.min().item(), - "attn/entropy_min_head": head_entropy.min().item(), - "attn/collapsed_heads": int((head_entropy < COLLAPSED_ENTROPY).sum()), - } - for layer in range(len(layer_max_logit)): - metrics[f"attn/max_logit/layer_{layer}"] = layer_max_logit[layer].item() - metrics[f"attn/entropy/layer_{layer}"] = layer_entropy[layer].item() - return metrics - - -def norm_role(name: str) -> str: - """Normalization layers by role (`input_layernorm`, `post_attention_layernorm`, `q_norm`, `k_norm`, `norm`), - everything else by `param_group`.""" - module = name.split(".")[-2] - return f"norm_gain/{module}" if "norm" in module else param_group(name) - - -class UpdateTracker: - """Size of one optimizer update relative to the weights, per parameter group: ||W_after - W_before|| / - ||W_before||, plus the weight RMS per group and of the norm gains by role.""" - - def before_step(self, model: nn.Module) -> None: - self.before = {name: param.detach().clone() for name, param in model.named_parameters()} - - def after_step(self, model: nn.Module) -> dict[str, float]: - params = dict(model.named_parameters()) - weight = group_norms(((name, self.before[name]) for name in params), norm_role) - update = group_norms(((name, params[name].detach() - self.before[name]) for name in params), norm_role) - numel = {} - for name, param in params.items(): - numel[norm_role(name)] = numel.get(norm_role(name), 0) + param.numel() - del self.before - metrics = {} - for group, norm in weight.items(): - metrics[f"param/rms/{group}"] = (norm / numel[group] ** 0.5).item() - metrics[f"update/rms_ratio/{group}"] = (update[group] / norm).item() - return metrics diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index c32699d0407..7d57e3ae9dc 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -2,7 +2,6 @@ import asyncio import json -import itertools import math import time from datetime import timedelta @@ -43,7 +42,6 @@ source_tokens, setup_val_dataset, ) -from prime_rl.trainer.pretrain.diagnostics import UpdateTracker, probe_batch, run_probe from prime_rl.trainer.pretrain.suite import ( length_label, prepare_multi_length_set, @@ -264,12 +262,6 @@ def train(config: PretrainConfig): # Create the iterator only after a potential resume: iter() forks the workers with the # dataset's current state. dataiter = iter(dataloader) - probe = None - if config.diagnostics is not None: - # The diagnostic forward's fixed batch: the start of the first micro-batch, which is then trained on as usual. - first_micro_batch = next(dataiter) - probe = probe_batch(first_micro_batch, config.diagnostics.probe_tokens) - dataiter = itertools.chain([first_micro_batch], dataiter) cp_enabled = parallel_dims.cp_enabled cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() if cp_enabled else 0 @@ -402,7 +394,6 @@ def run_suite(step: int) -> None: last_bias_moves = None # Per-expert selection bias, router-row norm and load of every MoE layer, every MOE_TRACE_INTERVAL steps. moe_trace: dict[str, list] = {"step": [], "bias": [], "router_norm": [], "load": []} - update_tracker = UpdateTracker() logger.info(f"Starting training loop (max_steps={config.max_steps})") while True: torch.cuda.reset_peak_memory_stats() @@ -463,14 +454,10 @@ def run_suite(step: int) -> None: group_grad_norms = grad_norms_by_group(model) grad_norm = clip_grad_norm_(None, model, config.optim.max_norm) if config.optim.max_norm else None - is_diagnostic_step = config.diagnostics is not None and progress.step % config.diagnostics.interval == 0 - if is_diagnostic_step: - update_tracker.before_step(model) optimizer.step() optimizer.zero_grad() current_lr = optimizer.param_groups[0]["lr"] scheduler.step() - diagnostic_metrics = update_tracker.after_step(model) if is_diagnostic_step else {} # Peak memory of the step's forward, backward and optimizer passes, over all ranks. local_peak_memory = torch.tensor(torch.cuda.max_memory_reserved() / 1024**3, device="cuda") peak_memory = reduce_stats({"sum": local_peak_memory}, {"max": local_peak_memory}, dp_cp_group) @@ -489,10 +476,6 @@ def run_suite(step: int) -> None: save_ckpt_time = time.perf_counter() - save_ckpt_start_time disk_metrics = get_ckpt_disk_metrics(config.run_dir) - if is_diagnostic_step: - diagnostic_metrics |= run_probe(model, probe, dp_cp_group) - finish_no_grad_forwards() - eval_start_time = time.perf_counter() if config.val is not None and (is_last_step or progress.step % config.val.interval == 0): run_val(progress.step) @@ -547,7 +530,6 @@ def run_suite(step: int) -> None: **peak_memory_metrics, **result.metrics, **moe_metrics, - **diagnostic_metrics, "step": progress.step, } if grad_norm is not None: diff --git a/tests/unit/train/pretrain/test_diagnostics.py b/tests/unit/train/pretrain/test_diagnostics.py deleted file mode 100644 index c06ade505ce..00000000000 --- a/tests/unit/train/pretrain/test_diagnostics.py +++ /dev/null @@ -1,22 +0,0 @@ -import math - -import torch - -from prime_rl.trainer.pretrain.diagnostics import MIN_VISIBLE_KEYS, attention_stats - - -def test_attention_stats_respect_documents_and_causality(): - # Two documents; zero queries give uniform attention, so each query's entropy is log(visible keys). - lengths = (MIN_VISIBLE_KEYS + 6, MIN_VISIBLE_KEYS + 2) - num_tokens = sum(lengths) - cu_seqlens = torch.tensor([0, lengths[0], num_tokens], dtype=torch.int32) - q = torch.zeros(num_tokens, 2, 4) - k = torch.randn(num_tokens, 1, 4) - q[-1, 1] = 1.0 # one non-zero query in the second head, last token of the second document - max_logit, entropy_sum, queries = attention_stats(q, k, cu_seqlens, scale=0.5) - - visible = [i + 1 for length in lengths for i in range(length)] - assert int(queries) == sum(v >= MIN_VISIBLE_KEYS for v in visible) - assert torch.isclose(entropy_sum[0], torch.tensor(sum(math.log(v) for v in visible if v >= MIN_VISIBLE_KEYS))) - second_document_keys = k[lengths[0] :, 0] - assert torch.isclose(max_logit[1], ((second_document_keys @ q[-1, 1]).max() * 0.5).clamp(min=0)) # other queries: 0 From a75d809a55c83d8bc0d4c936b86480ee209d289a Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 15:00:45 +0200 Subject: [PATCH 088/132] feat(pretrain): read each epoch's windows in one global seeded order Each epoch of a token source reads all num_tokens // seq_len windows once, in a seeded permutation over the whole source (uint32, materialized per epoch) from a seeded random offset; the wrap from the end of the stream to its start begins a new document and the short tail is skipped. Replaces the 1M-token block shuffle and the window shuffle within 4,096-window groups. Held-out sources stay in order. --- docs/pretraining.md | 2 +- src/prime_rl/trainer/pretrain/data.py | 76 +++++++++----------------- tests/unit/train/pretrain/test_data.py | 38 +++++-------- 3 files changed, 39 insertions(+), 77 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 6b527492f4d..53d4a0525e6 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -113,7 +113,7 @@ The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` ## Data -`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. Each epoch reads every token of a source once: in blocks of 2^20 tokens in a random order and from a random start offset, both drawn from `data.seed`, the source and the epoch. Within an epoch, the rows of each group of 4,096 consecutive rows are read in a seeded random order, so the pieces of one long document spread over many steps. So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole source, whatever the order of the files. Documents continue across rows, and attention restarts at every document, every row and every block boundary. Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. +`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of a source are read as one stream. Each epoch reads every row of a source once, in a random order over the whole source and from a random start offset, both drawn from `data.seed`, the source and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole source, whatever the order of the files. Documents continue across rows, and attention restarts at every document and every row. Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. Each row comes from one source, drawn with the mixture `data.weights`. At startup the trainer computes how many passes over each source the run makes and fails above `data.max_epochs`. diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index 138aa50e00e..5f77ea0bc24 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -9,26 +9,19 @@ from prime_rl.trainer.sft.data import StatefulIterableDataset from prime_rl.utils.logger import get_logger -# Unit of the per-epoch shuffle of a token source. A block boundary inside a document splits it, about once -# per 256 windows of 4,096 tokens. -BLOCK_TOKENS = 1 << 20 -# Unit of the window shuffle of a token source: consecutive windows, so the pieces of one long document spread -# over many steps. -WINDOW_GROUP = 4096 - class TokenShards: - """Datatrove tokenized shards in one folder, the ``num_tokens`` first tokens of the files in name order. + """Datatrove tokenized shards in one folder, the ``num_tokens`` first tokens of the files in name order, read as + one cyclic token stream. ``.ds`` holds the token ids (little-endian, 2 or 4 bytes as recorded on the first line of ``.ds.metadata``) and ``.ds.index`` the uint64 end offset of every document. File boundaries are document boundaries. - Stream position ``p`` belongs to epoch ``p // num_tokens``, and every epoch reads every token once. With a - ``seed``, each epoch reads the tokens in blocks of ``BLOCK_TOKENS`` in a seeded random order, starting at a - seeded random offset, so epochs differ in order and window boundaries, and part of an epoch is a sample of - the whole source, and the windows are read shuffled within groups (``window_start``). Without one - (held-out sets), the stream is the files in name order.""" + With a ``seed``, each epoch reads all ``num_tokens // seq_len`` windows of the source once, in a seeded random + order and from a seeded random offset, so epochs differ in order and window boundaries, and part of an epoch is a + sample of the whole source (``window_start``). The last ``num_tokens % seq_len`` tokens of an epoch are skipped. + Without one (held-out sets), the windows are read in order.""" def __init__(self, config: TokenSourceConfig, seed: tuple[int, str] | None = None): self.files = sorted(config.path.glob("*.ds")) @@ -43,8 +36,8 @@ def __init__(self, config: TokenSourceConfig, seed: tuple[int, str] | None = Non self.num_tokens = total if config.max_tokens is None else config.max_tokens self.seed = None if seed is None else [seed[0], zlib.crc32(seed[1].encode())] self.memmaps: dict[int, tuple[np.memmap, np.memmap]] = {} - self.layouts: dict[int, tuple[np.ndarray, np.ndarray, int]] = {} - self.window_order: tuple[tuple[int, int] | None, np.ndarray | None] = (None, None) + # The current epoch, its window order and its offset. + self.epoch_order: tuple[int | None, np.ndarray | None, int] = (None, None, 0) def open(self, file_index: int) -> tuple[np.memmap, np.memmap]: # Opened lazily, so every dataloader worker maps the files itself. @@ -55,33 +48,18 @@ def open(self, file_index: int) -> tuple[np.memmap, np.memmap]: self.memmaps[file_index] = tokens, doc_ends return self.memmaps[file_index] - def layout(self, epoch: int) -> tuple[np.ndarray, np.ndarray, int]: - """An epoch's block order, the stream offsets where its blocks start, and its start offset.""" - if epoch not in self.layouts: - rng = np.random.default_rng([*self.seed, epoch]) - order = rng.permutation(-(-self.num_tokens // BLOCK_TOKENS)) - lengths = np.minimum(BLOCK_TOKENS, self.num_tokens - order * BLOCK_TOKENS) - self.layouts[epoch] = order, np.cumsum([0, *lengths]), int(rng.integers(self.num_tokens)) - return self.layouts[epoch] - def locate(self, position: int, length: int) -> tuple[int, int, int]: """File index, offset in the file and length of the contiguous file run at stream ``position``, at most ``length`` tokens.""" - epoch, position = divmod(position, self.num_tokens) + position %= self.num_tokens length = min(length, self.num_tokens - position) - if self.seed is not None: - order, block_starts, offset = self.layout(epoch) - position = (position + offset) % self.num_tokens - block = int(np.searchsorted(block_starts, position, side="right")) - 1 - length = min(length, int(block_starts[block + 1]) - position) - position = int(order[block]) * BLOCK_TOKENS + position - int(block_starts[block]) file_index = int(np.searchsorted(self.starts, position, side="right")) - 1 length = min(length, int(self.starts[file_index + 1]) - position) return file_index, position - int(self.starts[file_index]), length def read(self, start: int, length: int) -> tuple[np.ndarray, list[int]]: - """Tokens ``[start, start + length)`` of the stream and the offsets in them where a document starts. - A block cut through a document starts a new one.""" + """Tokens ``[start, start + length)`` of the cyclic stream and the offsets in them where a document starts. + The wrap from the end of the stream to its start begins a new document.""" pieces, doc_starts = [], [] offset = 0 while offset < length: @@ -96,20 +74,17 @@ def read(self, start: int, length: int) -> tuple[np.ndarray, list[int]]: return np.concatenate(pieces).astype(np.int64), doc_starts def window_start(self, position: int, seq_len: int) -> int: - """Stream position of the window to read when ``position`` tokens have been read. With a seed, each epoch's - windows are read in a seeded random order within groups of ``WINDOW_GROUP`` consecutive windows.""" + """Stream offset of the window to read when ``position`` tokens have been read.""" if self.seed is None: - return position - epoch = position // self.num_tokens - group = (position - epoch * self.num_tokens) // (WINDOW_GROUP * seq_len) - group_start = epoch * self.num_tokens + group * WINDOW_GROUP * seq_len - group_end = min((epoch + 1) * self.num_tokens, group_start + WINDOW_GROUP * seq_len) - # The windows that start in [group_start, group_end). - first, last = -(-group_start // seq_len), -(-group_end // seq_len) - if self.window_order[0] != (epoch, group): - order = np.random.default_rng([*self.seed, epoch, group, 1]).permutation(last - first) - self.window_order = (epoch, group), order - return (first + int(self.window_order[1][position // seq_len - first])) * seq_len + return position % self.num_tokens + num_windows = self.num_tokens // seq_len + epoch, window = divmod(position // seq_len, num_windows) + if self.epoch_order[0] != epoch: + rng = np.random.default_rng([*self.seed, epoch]) + order = rng.permutation(num_windows).astype(np.uint32) + self.epoch_order = epoch, order, int(rng.integers(self.num_tokens)) + _, order, offset = self.epoch_order + return (offset + int(order[window]) * seq_len) % self.num_tokens def last_token_of_first_document(self) -> int: tokens, doc_ends = self.open(0) @@ -140,11 +115,10 @@ def window_start(self, position: int, seq_len: int) -> int: class TokenDataset(StatefulIterableDataset): """``seq_len`` windows of token sources, each window from one source of a weighted mixture. - Window ``w``'s source is drawn from the mixture ``weights`` with an RNG keyed by ``(seed, w)``, and it continues the source's stream where that source's last - window ended. Every data rank walks the whole window sequence but reads only its own windows; the - state is the window counter and the tokens read per source. Each - window reads one token past its end for the last target. With ``num_windows`` the dataset stops - after that many windows.""" + Window ``w``'s source is drawn from the mixture ``weights`` with an RNG keyed by ``(seed, w)``, and it is that + source's next window (``TokenShards.window_start``). Every data rank walks the whole window sequence but reads only + its own windows; the state is the window counter and the tokens read per source. Each window reads one token past + its end for the last target. With ``num_windows`` the dataset stops after that many windows.""" def __init__( self, diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index fc93bfffc6a..b2fe0018202 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -4,7 +4,6 @@ import numpy as np from prime_rl.configs.pretrain import TokenSourceConfig -from prime_rl.trainer.pretrain import data from prime_rl.trainer.pretrain.data import TokenDataset, TokenShards, collate_rows, source_tokens @@ -29,32 +28,21 @@ def test_token_shards_read_across_files_and_wrap(tmp_path): assert doc_starts == [2] -def test_shuffled_epochs_read_every_token_once_in_a_new_order(tmp_path, monkeypatch): - monkeypatch.setattr(data, "BLOCK_TOKENS", 4) +def test_shuffled_epochs_read_every_window_once_in_a_new_order(tmp_path): write_shard(tmp_path / "a" / "00.ds", [list(range(0, 30)), list(range(30, 50))]) - write_shard(tmp_path / "a" / "01.ds", [list(range(50, 70))]) + write_shard(tmp_path / "a" / "01.ds", [list(range(50, 72))]) shards = TokenShards(TokenSourceConfig(path=tmp_path / "a"), seed=(0, "a")) - epochs = [shards.read(epoch * 70, 70)[0].tolist() for epoch in range(3)] - assert all(sorted(epoch) == list(range(70)) for epoch in epochs) - assert len({tuple(epoch) for epoch in epochs}) == 3 - # Reading in windows gives the same stream, so resuming at any position continues it. - windows = [shards.read(start, 7)[0].tolist() for start in range(70, 140, 7)] - assert sum(windows, []) == epochs[1] - - -def test_window_shuffle_reads_every_window_of_an_epoch_once(tmp_path, monkeypatch): - monkeypatch.setattr(data, "WINDOW_GROUP", 5) - write_shard(tmp_path / "a" / "00.ds", [list(range(70))]) - shards = TokenShards(TokenSourceConfig(path=tmp_path / "a"), seed=(0, "a")) - starts = [shards.window_start(position, 7) for position in range(0, 140, 7)] - groups = [starts[i : i + 5] for i in range(0, 20, 5)] - assert [sorted(group) for group in groups] == [list(range(i * 7, i * 7 + 35, 7)) for i in range(0, 20, 5)] - assert starts != list(range(0, 140, 7)) - # A group never reaches into the next epoch. - monkeypatch.setattr(data, "WINDOW_GROUP", 100) - shards = TokenShards(TokenSourceConfig(path=tmp_path / "a"), seed=(0, "a")) - starts = [shards.window_start(position, 7) for position in range(0, 140, 7)] - assert sorted(starts[:10]) == list(range(0, 70, 7)) and sorted(starts[10:]) == list(range(70, 140, 7)) + # 72 tokens hold 10 windows of 7; the 2-token tail of each epoch is skipped. + orders = [] + for epoch in range(3): + starts = [shards.window_start(position, 7) for position in range(epoch * 70, epoch * 70 + 70, 7)] + tokens = [token for start in starts for token in shards.read(start, 7)[0].tolist()] + assert len(set(tokens)) == 70 + orders.append(tuple(starts)) + assert len(set(orders)) == 3 + # Another reader (a dataloader worker, a resumed run) draws the same order. + again = TokenShards(TokenSourceConfig(path=tmp_path / "a"), seed=(0, "a")) + assert tuple(again.window_start(position, 7) for position in range(70, 140, 7)) == orders[1] def test_token_dataset_resumes_and_continues_each_source(tmp_path): From 190533a5b200d85a4b4c2d69b63699e248d75553 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 15:03:45 +0200 Subject: [PATCH 089/132] fix(pretrain): eval requests start with BOS only if the training documents do The training shards decide: if the first document of every training source starts with the tokenizer's BOS, eval requests get it; otherwise ppl documents, ctxcont contexts and pairs sentences start at position 0 (ppl and pairs score from the second token; ppl bpb leaves out the first token's bytes). Before, every request started with a BOS the model never saw in training. --- docs/pretraining.md | 2 +- src/prime_rl/trainer/pretrain/data.py | 12 +++++ src/prime_rl/trainer/pretrain/suite.py | 58 +++++++++++++++---------- src/prime_rl/trainer/pretrain/train.py | 6 ++- tests/unit/train/pretrain/test_suite.py | 37 +++++++++++++--- 5 files changed, 83 insertions(+), 32 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 53d4a0525e6..f3574532a9d 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -149,7 +149,7 @@ interval = 500 web = { path = "data/web-val" } ``` -**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requires `model.cp = 1`. +**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training source at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. ```toml [evals] diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index 5f77ea0bc24..ab61dd2549d 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -86,6 +86,9 @@ def window_start(self, position: int, seq_len: int) -> int: _, order, offset = self.epoch_order return (offset + int(order[window]) * seq_len) % self.num_tokens + def first_token(self) -> int: + return int(self.open(0)[0][0]) + def last_token_of_first_document(self) -> int: tokens, doc_ends = self.open(0) return int(tokens[int(doc_ends[0]) - 1]) @@ -244,3 +247,12 @@ def check_eos(sources: dict[str, TokenShards], eos_token_id: int) -> None: f"The first document of source {name!r} ends with token {last}, not the tokenizer's EOS " f"{eos_token_id}. Were the shards tokenized with tokenizer.name?" ) + + +def starts_with_bos(sources: dict[str, TokenShards], bos_token_id: int | None) -> bool: + """Whether the documents of the sources start with the tokenizer's BOS, read off the first document of every + source. The eval suite formats its requests the same way.""" + starts = {name: source.first_token() == bos_token_id for name, source in sources.items()} + if len(set(starts.values())) > 1: + raise ValueError(f"Only some token sources start their documents with BOS: {starts}") + return all(starts.values()) diff --git a/src/prime_rl/trainer/pretrain/suite.py b/src/prime_rl/trainer/pretrain/suite.py index 55a97c5eee8..b173419fbeb 100644 --- a/src/prime_rl/trainer/pretrain/suite.py +++ b/src/prime_rl/trainer/pretrain/suite.py @@ -2,12 +2,15 @@ (`pretrain-evals` offline suite). A set lives in ``//`` as ``test.parquet`` plus ``meta.json`` (``kind``, ``parquet_sha256``). Kinds: -- ``ppl``: column ``text``. Each document is BOS + text + EOS, cut into non-overlapping windows of ``seq_len`` - scored tokens (later windows start cold). Every token after BOS is scored. +Requests start with BOS only if the training documents do (``bos``), so evals see the training format. + +- ``ppl``: column ``text``. Each document is [BOS +] text + EOS, cut into non-overlapping windows of ``seq_len`` + scored tokens (later windows start cold). Every token after the first is scored; without BOS, bpb leaves out + the bytes of the first text token. - ``ctxcont``: ``context``/``continuation`` (+ ``choices``/``gold``, optional ``choice_contexts``). The - continuation is scored given BOS + context; context and continuation are tokenized jointly and the + continuation is scored given [BOS +] context; context and continuation are tokenized jointly and the continuation span is found by character offsets. -- ``pairs``: ``good``/``bad`` sentences, each scored as BOS + sentence. +- ``pairs``: ``good``/``bad`` sentences, each scored as [BOS +] sentence from its second token. A set is tokenized once (`prepare_set`); every round scores its requests (`score_requests`, the model's LM head replaced by `scoring_head`) and summarizes them (`summarize`).""" @@ -162,12 +165,12 @@ def score_sets( return per_set -def continuation_ids(tokenizer: PreTrainedTokenizer, context: str, continuation: str) -> Request: - """BOS + jointly tokenized context + continuation, and the index of the first continuation token (the +def continuation_ids(tokenizer: PreTrainedTokenizer, context: str, continuation: str, prefix: list[int]) -> Request: + """``prefix`` + jointly tokenized context + continuation, and the index of the first continuation token (the first token that ends after the context; a token straddling the boundary counts as continuation).""" encoding = tokenizer(context + continuation, add_special_tokens=False, return_offsets_mapping=True) first = next(i for i, (_, end) in enumerate(encoding["offset_mapping"]) if end > len(context)) - return [tokenizer.bos_token_id, *encoding["input_ids"]], first + 1 + return [*prefix, *encoding["input_ids"]], first + len(prefix) def truncate_left(ids: list[int], start: int, seq_len: int) -> Request: @@ -186,32 +189,39 @@ def requests_of(row: dict) -> list[tuple[str, str]]: return [(row["context"], row["continuation"])] -def prepare_set(name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_len: int) -> EvalSet: - """Load and tokenize one eval set.""" +def prepare_set(name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_len: int, bos: bool) -> EvalSet: + """Load and tokenize one eval set; requests start with BOS if ``bos``.""" meta = json.loads((data_dir / name / "meta.json").read_text()) table = pq.read_table(data_dir / name / "test.parquet") kind, requests, context = meta["kind"], [], {} + prefix = [tokenizer.bos_token_id] if bos else [] if kind == "ppl": texts = table.column("text").to_pylist() - owner = [] - for doc, text_ids in enumerate(tokenizer(texts, add_special_tokens=False)["input_ids"]): - ids = [tokenizer.bos_token_id, *text_ids, tokenizer.eos_token_id] + encodings = tokenizer(texts, add_special_tokens=False, return_offsets_mapping=True) + owner, num_bytes = [], [] + for doc, (text, text_ids) in enumerate(zip(texts, encodings["input_ids"])): + ids = [*prefix, *text_ids, tokenizer.eos_token_id] for begin in range(0, len(ids) - 1, seq_len): requests.append((ids[begin : begin + seq_len + 1], 1)) owner.append(doc) - context = {"texts": texts, "owner": owner} + # Without BOS the first text token is only context. + scored_text = text if bos else text[encodings["offset_mapping"][doc][0][1] :] + num_bytes.append(len(scored_text.encode())) + context = {"owner": owner, "bytes": np.array(num_bytes)} elif kind == "ctxcont": rows = table.to_pylist() options = [requests_of(row) for row in rows] requests = [ - truncate_left(*continuation_ids(tokenizer, ctx, cont), seq_len) for item in options for ctx, cont in item + truncate_left(*continuation_ids(tokenizer, ctx, cont, prefix), seq_len) + for item in options + for ctx, cont in item ] context = {"rows": rows, "options": options} elif kind == "pairs": good, bad = table.column("good").to_pylist(), table.column("bad").to_pylist() encoded = tokenizer(good + bad, add_special_tokens=False)["input_ids"] # Cut to the window like the other kinds, so every request fits in a row of ``token_budget`` tokens. - requests = [([tokenizer.bos_token_id, *ids][: seq_len + 1], 1) for ids in encoded] + requests = [([*prefix, *ids][: seq_len + 1], 1) for ids in encoded] context = {"num_good": len(good)} else: raise ValueError(f"Unknown eval set kind {kind!r} ({name})") @@ -232,9 +242,9 @@ class MultiLengthSet: def prepare_multi_length_set( - name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_lens: list[int] + name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_lens: list[int], bos: bool ) -> MultiLengthSet: - versions = {seq_len: prepare_set(name, data_dir, tokenizer, seq_len) for seq_len in seq_lens} + versions = {seq_len: prepare_set(name, data_dir, tokenizer, seq_len, bos) for seq_len in seq_lens} distinct: dict[tuple[tuple[int, ...], int], int] = {} index = { seq_len: [distinct.setdefault((tuple(ids), start), len(distinct)) for ids, start in version.requests] @@ -257,27 +267,27 @@ def length_label(seq_len: int) -> str: def summarize(eval_set: EvalSet, logprobs: list[np.ndarray], hits: list[np.ndarray]) -> tuple[dict, dict]: """Summary metrics and per-document / per-item arrays of one scored set.""" if eval_set.kind == "ppl": - return summarize_ppl(eval_set.context["texts"], eval_set.context["owner"], logprobs) + return summarize_ppl(eval_set.context["bytes"], eval_set.context["owner"], logprobs) if eval_set.kind == "ctxcont": return summarize_ctxcont(eval_set.context["rows"], eval_set.context["options"], logprobs, hits) return summarize_pairs(eval_set.context["num_good"], logprobs) -def summarize_ppl(texts: list[str], owner: list[int], logprobs: list[np.ndarray]) -> tuple[dict, dict]: +def summarize_ppl(num_bytes: np.ndarray, owner: list[int], logprobs: list[np.ndarray]) -> tuple[dict, dict]: + """``num_bytes``: the bytes of each document's scored text.""" bins = list(zip(POSITION_BINS[:-1], POSITION_BINS[1:])) - nll = np.zeros(len(texts)) - tokens = np.zeros(len(texts), dtype=np.int64) - nll_bins = np.zeros((len(texts), len(bins))) + nll = np.zeros(len(num_bytes)) + tokens = np.zeros(len(num_bytes), dtype=np.int64) + nll_bins = np.zeros((len(num_bytes), len(bins))) for doc, window in zip(owner, logprobs): nll[doc] -= window.sum() tokens[doc] += len(window) for b, (lo, hi) in enumerate(bins): nll_bins[doc, b] -= window[lo:hi].sum() - num_bytes = np.array([len(text.encode()) for text in texts]) summary = { "loss": nll.sum() / tokens.sum(), "bpb": nll.sum() / (num_bytes.sum() * math.log(2)), - "docs": len(texts), + "docs": len(num_bytes), "tokens": int(tokens.sum()), } for b, (lo, hi) in enumerate(bins): diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 7d57e3ae9dc..7da813c699d 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -41,6 +41,7 @@ setup_train_dataset, source_tokens, setup_val_dataset, + starts_with_bos, ) from prime_rl.trainer.pretrain.suite import ( length_label, @@ -351,11 +352,12 @@ def run_val(step: int) -> None: eval_sets = [] if config.evals is not None: start_time = time.perf_counter() + bos = config.data.type == "tokens" and starts_with_bos(train_dataset.sources, tokenizer.bos_token_id) eval_sets = [ - prepare_multi_length_set(name, config.evals.data_dir, tokenizer, config.evals.seq_lens) + prepare_multi_length_set(name, config.evals.data_dir, tokenizer, config.evals.seq_lens, bos) for name in config.evals.sets ] - logger.info(f"Prepared {len(eval_sets)} eval sets in {time.perf_counter() - start_time:.0f}s") + logger.info(f"Prepared {len(eval_sets)} eval sets (BOS: {bos}) in {time.perf_counter() - start_time:.0f}s") @torch.no_grad() def run_suite(step: int) -> None: diff --git a/tests/unit/train/pretrain/test_suite.py b/tests/unit/train/pretrain/test_suite.py index 15b7f87004c..dd2ad8cf61e 100644 --- a/tests/unit/train/pretrain/test_suite.py +++ b/tests/unit/train/pretrain/test_suite.py @@ -49,15 +49,23 @@ def test_summarize_ctxcont_scores_the_gold_choice(): class CharTokenizer: bos_token_id, eos_token_id = 1, 2 - def __call__(self, texts, add_special_tokens=False): - return {"input_ids": [[ord(char) for char in text] for text in texts]} + def __call__(self, texts, add_special_tokens=False, return_offsets_mapping=False): + if isinstance(texts, str): + return { + "input_ids": [ord(char) for char in texts], + "offset_mapping": [(i, i + 1) for i in range(len(texts))], + } + return { + "input_ids": [[ord(char) for char in text] for text in texts], + "offset_mapping": [[(i, i + 1) for i in range(len(text))] for text in texts], + } def test_multi_length_set_scores_shared_requests_once(tmp_path): (tmp_path / "docs").mkdir() pq.write_table(pa.table({"text": ["short", "a much longer document"]}), tmp_path / "docs" / "test.parquet") (tmp_path / "docs" / "meta.json").write_text(json.dumps({"kind": "ppl", "parquet_sha256": "x"})) - multi = prepare_multi_length_set("docs", tmp_path, CharTokenizer(), [8, 32]) + multi = prepare_multi_length_set("docs", tmp_path, CharTokenizer(), [8, 32], bos=True) # "short" (5 chars + BOS + EOS) is one window at both lengths; the long document is 3 windows at 8, 1 at 32. assert len(multi.versions[8].requests) == 4 and len(multi.versions[32].requests) == 2 assert len(multi.requests) == 5 @@ -89,7 +97,7 @@ def test_score_sets_matches_scoring_each_set_alone(tmp_path): (tmp_path / name).mkdir() pq.write_table(pa.table({"text": texts}), tmp_path / name / "test.parquet") (tmp_path / name / "meta.json").write_text(json.dumps({"kind": "ppl", "parquet_sha256": "x"})) - eval_sets = [prepare_multi_length_set(name, tmp_path, CharTokenizer(), [8, 32]) for name in "ab"] + eval_sets = [prepare_multi_length_set(name, tmp_path, CharTokenizer(), [8, 32], bos=True) for name in "ab"] together = score_sets(TokenPositionModel(), eval_sets, token_budget=40, group=dist.group.WORLD) for eval_set, (logprobs, hits) in zip(eval_sets, together): alone_logprobs, alone_hits = score_requests( @@ -106,8 +114,27 @@ def test_pairs_requests_fit_the_window(tmp_path): table = pa.table({"good": ["a long good sentence", "ok"], "bad": ["a long bad sentence!", "no"]}) pq.write_table(table, tmp_path / "pairs" / "test.parquet") (tmp_path / "pairs" / "meta.json").write_text(json.dumps({"kind": "pairs", "parquet_sha256": "x"})) - eval_set = prepare_set("pairs", tmp_path, CharTokenizer(), seq_len=8) + eval_set = prepare_set("pairs", tmp_path, CharTokenizer(), seq_len=8, bos=True) assert all(len(ids) - 1 <= 8 for ids, _ in eval_set.requests) logprobs = [np.full(len(ids) - start, -1.0) for ids, start in eval_set.requests] summary, _ = summarize(eval_set, logprobs, [np.ones_like(lp, dtype=bool) for lp in logprobs]) assert set_metrics("pairs", 8, summary)["eval/pairs/acc@8"] == 0.0 # equal scores: good never beats bad + + +def test_requests_without_bos_start_at_the_text(tmp_path): + for name, kind, columns in [ + ("docs", "ppl", {"text": ["hello"]}), + ("cc", "ctxcont", {"context": ["Q:"], "continuation": [" yes"]}), + ("pairs", "pairs", {"good": ["ab"], "bad": ["ba"]}), + ]: + (tmp_path / name).mkdir() + pq.write_table(pa.table(columns), tmp_path / name / "test.parquet") + (tmp_path / name / "meta.json").write_text(json.dumps({"kind": kind, "parquet_sha256": "x"})) + docs = prepare_set("docs", tmp_path, CharTokenizer(), seq_len=32, bos=False) + assert docs.requests == [([ord(char) for char in "hello"] + [2], 1)] + assert docs.context["bytes"].tolist() == [4] # the first character is context only + with_bos = prepare_set("docs", tmp_path, CharTokenizer(), seq_len=32, bos=True) + assert with_bos.requests[0][0][0] == 1 and with_bos.context["bytes"].tolist() == [5] + ids, start = prepare_set("cc", tmp_path, CharTokenizer(), seq_len=32, bos=False).requests[0] + assert ids == [ord(char) for char in "Q: yes"] and start == 2 + assert prepare_set("pairs", tmp_path, CharTokenizer(), seq_len=32, bos=False).requests[0] == ([97, 98], 1) From 0c156e76aa3a99218923b64bcc3cfc5715a800a7 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 15:04:59 +0200 Subject: [PATCH 090/132] refactor(pretrain): derive the eval row length from the training micro-batch EvalsConfig.token_budget is gone: the suite packs rows of data.micro_batch_tokens tokens, or the longest eval context if that is longer. --- docs/pretraining.md | 2 +- .../prime-rl-configs/src/prime_rl/configs/pretrain.py | 9 --------- src/prime_rl/trainer/pretrain/train.py | 4 +++- 3 files changed, 4 insertions(+), 11 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index f3574532a9d..0045ca28480 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -149,7 +149,7 @@ interval = 500 web = { path = "data/web-val" } ``` -**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training source at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. +**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows of `data.micro_batch_tokens` tokens (or the longest of `seq_lens`, if longer) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training source at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. ```toml [evals] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index ff4e04f282d..f3739b96767 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -185,18 +185,9 @@ class EvalsConfig(BaseConfig): ``ctxcont`` sets. Metrics are named ``eval//@`` (e.g. ``bpb@4k``). Requests that are the same at several lengths are scored once.""" - token_budget: int = Field(32768, ge=1) - """Tokens per packed row (one forward). Must exceed the longest context length.""" - save_arrays: bool = False """Write per-document arrays (npz) and summaries to ``/evals/step_N/``.""" - @model_validator(mode="after") - def validate_token_budget(self): - if self.token_budget <= max(self.seq_lens): - raise ValueError("token_budget must exceed the longest of seq_lens") - return self - class PretrainAdamWConfig(AdamWConfig): eps: float = Field(1e-8, gt=0) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 7da813c699d..3e621fbb40b 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -358,6 +358,8 @@ def run_val(step: int) -> None: for name in config.evals.sets ] logger.info(f"Prepared {len(eval_sets)} eval sets (BOS: {bos}) in {time.perf_counter() - start_time:.0f}s") + # Packed rows as long as a training micro-batch, or the longest eval context if that is longer. + eval_token_budget = max(config.data.micro_batch_tokens, *config.evals.seq_lens) @torch.no_grad() def run_suite(step: int) -> None: @@ -367,7 +369,7 @@ def run_suite(step: int) -> None: metrics = {} # Eager, also for compiled layers: the packed rows vary in length and would recompile at every forward. with scoring_head_active(model), torch.compiler.set_stance("force_eager"): - scored = score_sets(model, eval_sets, config.evals.token_budget, get_dcp_group()) + scored = score_sets(model, eval_sets, eval_token_budget, get_dcp_group()) for eval_set, (logprobs, hits) in zip(eval_sets, scored or []): for seq_len, version in eval_set.versions.items(): index = eval_set.index[seq_len] From 63428346efe14d5b97db5ae1390be1bb52acd899 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 16:13:05 +0200 Subject: [PATCH 091/132] docs(pretrain): note the read cost of the global window order --- docs/pretraining.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index 0045ca28480..ec2a273cf12 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -113,7 +113,7 @@ The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` ## Data -`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of a source are read as one stream. Each epoch reads every row of a source once, in a random order over the whole source and from a random start offset, both drawn from `data.seed`, the source and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole source, whatever the order of the files. Documents continue across rows, and attention restarts at every document and every row. Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. +`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of a source are read as one stream. Each epoch reads every row of a source once, in a random order over the whole source and from a random start offset, both drawn from `data.seed`, the source and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole source, whatever the order of the files. Every row is a random read (about 33 ms each from a cold cache on CephFS); at the screening shapes, `data.num_workers = 2` with the default prefetch hides this (`time/data` stays at 0). Check `time/data` on other storage. Documents continue across rows, and attention restarts at every document and every row. Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. Each row comes from one source, drawn with the mixture `data.weights`. At startup the trainer computes how many passes over each source the run makes and fails above `data.max_epochs`. From 51bf2dfc89c96f187405280351347adb91890a76 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 18:38:05 +0200 Subject: [PATCH 092/132] refactor(pretrain): describe prime_moe on its own terms; trim unused config fields and forward args The model docs, example and docstrings list what prime_moe has instead of framing it as GLM-4.5's architecture; its module names and state-dict layout still match glm4_moe. PrimeMoeConfig drops use_cache, initializer_range and attention_dropout (read nowhere; old config.json keys still load) and documents pad_token_id. PrimeMoeForCausalLM.forward drops past_key_values, use_cache, logits_to_keep and inputs_embeds, and PrimeMoeModel.forward the inputs_embeds path: no caller passes them. --- docs/pretraining.md | 2 +- examples/pretrain/prime_moe.toml | 2 +- .../prime_moe/configuration_prime_moe.py | 16 +++----- .../models/prime_moe/modeling_prime_moe.py | 39 +++++-------------- 4 files changed, 17 insertions(+), 42 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index ec2a273cf12..e6de571a339 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -36,7 +36,7 @@ name = "examples/pretrain/moe-d512" name = "tokenizers/deepseek-v4-flash-base" # defaults to model.name ``` -MoE models use `prime_moe`, pretraining's own model (`prime_rl.trainer.models.prime_moe`), so architecture work does not touch the `glm4_moe` that SFT and RL use. It is GLM-4.5's architecture with the same config fields and state-dict layout, plus: +MoE models use `prime_moe`, pretraining's own model (`prime_rl.trainer.models.prime_moe`), so architecture work does not touch the shared MoE code that SFT and RL use. Its base is a decoder of grouped-query attention (optional QK norm, partial RoPE) and MoE layers with sigmoid-scored token-choice routed experts, a selection bias and shared experts, after leading dense layers. Its module names and state-dict layout match `glm4_moe`'s, so checkpoints convert the same way. On top of the base: - Gated Residual connections (Qwen3.8-Next): `hc_count` residual streams (above 1) with read gates of rank `hc_lowrank` replace the pre-norms and the final norm. Their projections train with AdamW. - Null experts (AdaMoE): `n_null_experts` zero-output experts the router can select next to the routed ones, so tokens use between 0 and `num_experts_per_tok` routed experts. Only the (token, routed expert) pairs are dispatched. The metrics are `router/k/*`. - Exclusive self attention (XSA): `exclusive_attention` removes from each head's output its component along the token's own value vector. diff --git a/examples/pretrain/prime_moe.toml b/examples/pretrain/prime_moe.toml index 67b99354bd0..72848dec750 100644 --- a/examples/pretrain/prime_moe.toml +++ b/examples/pretrain/prime_moe.toml @@ -1,4 +1,4 @@ -# GLM-4.5-style MoE (64 routed experts, 6 active, 2 shared) with d = 512 and 8 layers (moe-d512/config.json), trained from scratch on +# prime_moe with 64 routed experts (6 active) and 2 shared experts, d = 512 and 8 layers (moe-d512/config.json), trained from scratch on # the Dolma 3 mixture tokenized with the DeepSeek-V4 tokenizer. # 4096 steps x 256K tokens = 1.1B tokens. QK-norm and the learning rate follow the MoE QK-norm sweep (optimum 2^-7 # with QK-norm at d = 512). [parametrization] with this base shape keeps the LR when the width grows. diff --git a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py index 5f57cd4584d..e8c86efd6ed 100644 --- a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py +++ b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py @@ -3,9 +3,10 @@ class PrimeMoeConfig(PretrainedConfig): r""" - Configuration of the pretraining MoE (`PrimeMoeForCausalLM`): GLM-4.5's architecture (attention with optional QK - norm, sigmoid-scored token-choice experts with a selection bias, shared experts, leading dense layers) plus - Gated Residual connections and null experts. The fields shared with `Glm4MoeConfig` have the same names. + Configuration of the pretraining MoE (`PrimeMoeForCausalLM`): grouped-query attention (optional QK norm, gated + attention, exclusive self attention), MoE layers with sigmoid-scored token-choice routed experts, a selection bias, + shared experts and optional null experts or LatentMoE, after leading dense layers, and optional Gated Residual + connections. Fields that also exist in the Hugging Face MoE configs keep their names there. Args: vocab_size (`int`): Vocabulary size. @@ -18,13 +19,11 @@ class PrimeMoeConfig(PretrainedConfig): partial_rotary_factor (`float`): Share of each head's dimensions that RoPE rotates. hidden_act (`str`): Activation of the gated MLPs. max_position_embeddings (`int`): Maximum sequence length. - initializer_range (`float`): Unused by prime-rl's initialization; kept for Hugging Face compatibility. rms_norm_eps (`float`): Epsilon of the RMS norms. tie_word_embeddings (`bool`): Whether the LM head shares the input embedding. rope_theta (`float`): RoPE base period. rope_scaling (`dict`, *optional*): RoPE scaling configuration. attention_bias (`bool`): Whether the attention projections have biases. - attention_dropout (`float`): Attention dropout (unused in training). moe_intermediate_size (`int`): Hidden size of each routed and shared expert. num_experts_per_tok (`int`): Experts each token selects (`top_k`), null experts included. n_shared_experts (`int`): Number of shared experts; their hidden sizes add up. @@ -53,6 +52,7 @@ class PrimeMoeConfig(PretrainedConfig): soft-capped SwiGLU). Defaults to `hidden_act`. router_bias_on_logits (`bool`): Add the selection bias to the router logits instead of the sigmoid scores (the space quantile balancing works best in). + pad_token_id (`int`, *optional*): Padding token id, the input embedding's `padding_idx`. """ model_type = "prime_moe" @@ -69,14 +69,11 @@ def __init__( head_dim=None, hidden_act="silu", max_position_embeddings=131072, - initializer_range=0.02, rms_norm_eps=1e-5, - use_cache=True, tie_word_embeddings=False, rope_theta=10000.0, rope_scaling=None, attention_bias=False, - attention_dropout=0.0, moe_intermediate_size=1408, num_experts_per_tok=8, n_shared_experts=1, @@ -108,13 +105,10 @@ def __init__( self.num_key_value_heads = num_key_value_heads self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads self.hidden_act = hidden_act - self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps - self.use_cache = use_cache self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.attention_bias = attention_bias - self.attention_dropout = attention_dropout self.pad_token_id = pad_token_id self.moe_intermediate_size = moe_intermediate_size diff --git a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py index a4ad4a41613..d0ae449239f 100644 --- a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py +++ b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py @@ -1,11 +1,10 @@ -"""The pretraining MoE: GLM-4.5's architecture (prime-rl's `glm4_moe`) with Gated Residual connections and null -experts. Module names follow `glm4_moe`, so checkpoints share its state-dict layout.""" +"""The pretraining MoE (`prime_moe`), pretraining's own model, so architecture work leaves the shared MoE models +unchanged. Module names follow `glm4_moe`, so checkpoints keep that state-dict layout and its conversion.""" -from typing import Optional, Union +from typing import Optional import torch from torch import Tensor, nn -from transformers.cache_utils import Cache from transformers.modeling_layers import GradientCheckpointingLayer from transformers.modeling_outputs import BaseModelOutputWithPast from transformers.processing_utils import Unpack @@ -164,7 +163,7 @@ def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: @classmethod def conversion_chain(cls, config): - # Same expert layout as glm4_moe. + # The expert state-dict layout of glm4_moe, so its conversion applies. return conversion_chain(config) @@ -203,17 +202,13 @@ def __init__(self, config: PrimeMoeConfig): def forward( self, - input_ids: Optional[torch.LongTensor] = None, - position_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor, *, seq_lens: torch.LongTensor, seq_lens_are_pre_shard: bool = False, ) -> BaseModelOutputWithPast: - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if inputs_embeds is None: - inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + inputs_embeds = self.embed_tokens(input_ids) cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens( seq_lens.to(device=inputs_embeds.device), @@ -246,13 +241,9 @@ def __init__(self, config: PrimeMoeConfig): def forward( self, - input_ids: Optional[torch.LongTensor] = None, + input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[Cache] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - logits_to_keep: Union[int, torch.Tensor] = 0, temperature: Optional[torch.Tensor] = None, *, seq_lens: torch.LongTensor, @@ -261,25 +252,15 @@ def forward( ) -> PrimeLmOutput: """``labels`` and ``temperature`` go to prime-rl's wrapped LM head, which then returns per-token logprobs and entropy; without them it returns logits.""" - assert use_cache is None and past_key_values is None, "prime_moe does not support KV caching" if position_ids is None: - tokens = inputs_embeds if inputs_embeds is not None else input_ids - position_ids = torch.arange(tokens.shape[1], device=tokens.device).unsqueeze(0) - + position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) outputs = self.model( input_ids=input_ids, position_ids=position_ids, - inputs_embeds=inputs_embeds, seq_lens=seq_lens, seq_lens_are_pre_shard=seq_lens_are_pre_shard, ) - hidden_states = outputs.last_hidden_state - slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep - return self.lm_head( - hidden_states[:, slice_indices, :], - labels[:, slice_indices] if labels is not None else None, - temperature=temperature, - ) + return self.lm_head(outputs.last_hidden_state, labels, temperature=temperature) def init_buffers_post_meta(self): rotary_emb = self.model.rotary_emb From 96bf017ec761b6caf8bf84787045ff7579166c6c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 19:21:29 +0200 Subject: [PATCH 093/132] refactor(models): share HyperConnection and ExpandedRMSNorm in layers/ prime_moe's and qwen3_8_flash_next's hyper_connection.py were byte-identical. One copy now lives in models/layers/hyper_connection.py; qwen3_8_flash_next (including its re-export), prime_moe and models/init.py import it from there. --- src/prime_rl/trainer/models/init.py | 2 +- .../{prime_moe => layers}/hyper_connection.py | 0 .../models/prime_moe/modeling_prime_moe.py | 2 +- .../models/qwen3_8_flash_next/__init__.py | 2 +- .../qwen3_8_flash_next/hyper_connection.py | 74 ------------------- .../modeling_qwen3_8_flash_next.py | 2 +- .../qwen3_8_flash_next/position_learning.py | 2 +- 7 files changed, 5 insertions(+), 79 deletions(-) rename src/prime_rl/trainer/models/{prime_moe => layers}/hyper_connection.py (100%) delete mode 100644 src/prime_rl/trainer/models/qwen3_8_flash_next/hyper_connection.py diff --git a/src/prime_rl/trainer/models/init.py b/src/prime_rl/trainer/models/init.py index 79ea50103f7..84dcb072adf 100644 --- a/src/prime_rl/trainer/models/init.py +++ b/src/prime_rl/trainer/models/init.py @@ -8,8 +8,8 @@ from torch.distributed.tensor._utils import compute_local_shape_and_global_offset from prime_rl.configs.trainer import InitConfig, ParametrizationConfig +from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm from prime_rl.trainer.models.layers.moe import TokenChoiceTopKRouter -from prime_rl.trainer.models.prime_moe.hyper_connection import ExpandedRMSNorm from prime_rl.utils.weights import resolve_fqn ParamRole = Literal["embedding", "output", "residual_out", "hidden", "norm", "bias"] diff --git a/src/prime_rl/trainer/models/prime_moe/hyper_connection.py b/src/prime_rl/trainer/models/layers/hyper_connection.py similarity index 100% rename from src/prime_rl/trainer/models/prime_moe/hyper_connection.py rename to src/prime_rl/trainer/models/layers/hyper_connection.py diff --git a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py index d0ae449239f..b78993bba44 100644 --- a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py +++ b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py @@ -13,6 +13,7 @@ from prime_rl.trainer.models.base import PreTrainedModelPrimeRL from prime_rl.trainer.models.glm4_moe.converting_glm4_moe import conversion_chain from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig +from prime_rl.trainer.models.layers.hyper_connection import HyperConnection from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput from prime_rl.trainer.models.layers.mlp import FeedForward from prime_rl.trainer.models.layers.moe import GroupedExperts @@ -20,7 +21,6 @@ from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig from prime_rl.trainer.models.prime_moe.attention import PrimeFlashAttention from prime_rl.trainer.models.prime_moe.configuration_prime_moe import PrimeMoeConfig -from prime_rl.trainer.models.prime_moe.hyper_connection import HyperConnection from prime_rl.trainer.models.prime_moe.moe import PrimeMoE, PrimeRouter, SituGlu from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/__init__.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/__init__.py index 03f2798043c..0f89c639c65 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/__init__.py +++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/__init__.py @@ -1,9 +1,9 @@ +from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm, HyperConnection from prime_rl.trainer.models.qwen3_8_flash_next.attention import IndexedGatedAttention from prime_rl.trainer.models.qwen3_8_flash_next.configuration_qwen3_8_flash_next import ( Qwen3_8FlashNextConfig, Qwen3_8FlashNextTextConfig, ) -from prime_rl.trainer.models.qwen3_8_flash_next.hyper_connection import ExpandedRMSNorm, HyperConnection from prime_rl.trainer.models.qwen3_8_flash_next.modeling_qwen3_8_flash_next import ( Qwen3_8FlashNextDecoderLayer, Qwen3_8FlashNextForCausalLM, diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/hyper_connection.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/hyper_connection.py deleted file mode 100644 index 3082dc59768..00000000000 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/hyper_connection.py +++ /dev/null @@ -1,74 +0,0 @@ -import torch -import torch.nn.functional as F -from torch import nn - - -class ExpandedRMSNorm(nn.Module): - """RMS-normalize each residual stream with an independent affine weight.""" - - def __init__(self, hidden_size: int, stream_count: int, eps: float) -> None: - super().__init__() - self.hidden_size = hidden_size - self.stream_count = stream_count - self.eps = eps - self.weight = nn.Parameter(torch.zeros(stream_count * hidden_size)) - - def reset_parameters(self) -> None: - nn.init.zeros_(self.weight) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - streams = hidden_states.float().unflatten(-1, (self.stream_count, self.hidden_size)) - variance = streams.square().mean(dim=-1, keepdim=True) - normalized = streams * torch.rsqrt(variance + self.eps) - return (normalized.flatten(-2) * (1.0 + self.weight.float())).to(input_dtype) - - -class HyperConnection(nn.Module): - """Learned mixing and residual injection across expanded residual streams.""" - - def __init__( - self, - *, - hidden_size: int, - stream_count: int, - low_rank: int, - norm_eps: float, - with_residual_injection: bool = True, - ) -> None: - super().__init__() - self.hidden_size = hidden_size - self.stream_count = stream_count - expanded_size = stream_count * hidden_size - - self.hc_norm = ExpandedRMSNorm(hidden_size, stream_count, norm_eps) - self.input_mix_weight_down = nn.Linear(expanded_size, low_rank, bias=False) - self.input_mix_weight_up = nn.Linear(low_rank, expanded_size, bias=False) - self.block_inject_weight = ( - nn.Linear(expanded_size, stream_count, bias=False) if with_residual_injection else None - ) - - def mix(self, expanded_states: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: - normalized_states = self.hc_norm(expanded_states) - mixing_weights = F.silu(self.input_mix_weight_down(normalized_states) / self.stream_count) - mixing_weights = torch.sigmoid(self.input_mix_weight_up(mixing_weights)) - mixing_weights = mixing_weights.unflatten(-1, (self.stream_count, self.hidden_size)) - normalized_streams = normalized_states.unflatten(-1, (self.stream_count, self.hidden_size)) - block_input = (mixing_weights * normalized_streams).mean(dim=-2) - return block_input.to(expanded_states.dtype), (expanded_states, normalized_states) - - def forward(self, expanded_states: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: - return self.mix(expanded_states) - - def combine( - self, - block_output: torch.Tensor, - residual_state: tuple[torch.Tensor, torch.Tensor], - ) -> torch.Tensor: - expanded_states, normalized_states = residual_state - injection_weights = 2.0 * torch.sigmoid(self.block_inject_weight(normalized_states) / self.stream_count) - residual_streams = expanded_states.unflatten(-1, (self.stream_count, self.hidden_size)) - return (residual_streams + block_output.unsqueeze(-2) * injection_weights.unsqueeze(-1)).flatten(-2) - - -__all__ = ["ExpandedRMSNorm", "HyperConnection"] diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py index b3265b9d8cb..327b1f46e78 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py +++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py @@ -4,6 +4,7 @@ from prime_rl.trainer.models.base import PreTrainedModelPrimeRL from prime_rl.trainer.models.layers.gated_delta_net import GatedDeltaNet +from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm, HyperConnection from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear from prime_rl.trainer.models.layers.moe import MoE from prime_rl.trainer.models.layers.norms import ZeroCenteredRMSNorm @@ -17,7 +18,6 @@ is_hf_state_dict, is_prime_state_dict, ) -from prime_rl.trainer.models.qwen3_8_flash_next.hyper_connection import ExpandedRMSNorm, HyperConnection from prime_rl.trainer.models.qwen3_8_flash_next.moe import SigmoidOutputGatedMoE from prime_rl.trainer.models.qwen3_8_flash_next.ngram_embedding import NGramEmbedding from prime_rl.trainer.models.qwen3_8_flash_next.position_learning import PositionLearningEnhancement diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/position_learning.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/position_learning.py index f6e70f5db4a..df7ea6aeb34 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/position_learning.py +++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/position_learning.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from torch import nn -from prime_rl.trainer.models.qwen3_8_flash_next.hyper_connection import ExpandedRMSNorm +from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm from prime_rl.trainer.models.qwen3_8_flash_next.ngram_embedding import NGramEmbedding from prime_rl.utils.cp import CPContext, gather_for_cp From 6aa20ef8d419fe727b81ee2487f8eec78a750cab Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 19:33:01 +0200 Subject: [PATCH 094/132] refactor(configs): PretrainConfig and SFTConfig share a trainer launch base on BaseTrainerConfig TrainerLaunchConfig (configs/trainer.py) holds the run identity, output dir, monitors, resume, deployment and SLURM fields and validators that SFTConfig and PretrainConfig both carried; only the default run name differs (default_run_name). PretrainConfig now extends it with pretraining's autos (fp32 trainable router, weight decay 0.1) and drops its copies of BaseTrainerConfig's fields and validators; it rejects trace_path, memory_profiler_path and control, which its loop does not support. Pretraining now also names the prime monitor run. SFT's resolved config has the same values; only the order of its top-level keys changes (the launch fields come before renderer/data). --- .../src/prime_rl/configs/pretrain.py | 145 ++---------------- .../src/prime_rl/configs/sft.py | 108 ++----------- .../src/prime_rl/configs/trainer.py | 101 +++++++++++- 3 files changed, 122 insertions(+), 232 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index f3739b96767..6e82013fcc2 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -1,26 +1,13 @@ import math import uuid -import warnings from pathlib import Path from typing import Annotated, Literal, TypeAlias from pydantic import Field, model_validator -from prime_rl.configs.monitors import TrainMonitorsConfig -from prime_rl.configs.shared import ( - EnvVars, - HeartbeatConfig, - ResumeConfig, - RunConfig, - SFTDeploymentConfig, - SingleNodeDeploymentConfig, - SlurmConfig, - TrainerLogConfig, -) from prime_rl.configs.trainer import ( AdamWConfig, CheckpointConfig, - GCConfig, InitConfig, LinearSchedulerConfig, ModelConfig, @@ -28,10 +15,9 @@ MuonConfig, ParametrizationConfig, SchedulerConfig, - TokenizerConfig, - validate_scheduler, + TrainerLaunchConfig, ) -from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource +from prime_rl.utils.config import BaseConfig class MoEBalanceConfig(BaseConfig): @@ -202,17 +188,16 @@ class PretrainMuonConfig(MuonConfig): PretrainOptimizerConfig: TypeAlias = Annotated[PretrainAdamWConfig | PretrainMuonConfig, Field(discriminator="type")] -class PretrainConfig(BaseConfig): +class PretrainConfig(TrainerLaunchConfig): + # A trainable fp32 router, and weight decay 0.1. + auto_moe_router_dtype = "float32" + auto_freeze_moe_router = False + auto_weight_decay = 0.1 + model: PretrainModelConfig = PretrainModelConfig() """The model. ``model.name`` is an HF model id or a local model directory with a ``config.json`` (as in SFT and RL); its weights are not loaded: the model is initialized from scratch (``model.init``) or resumed.""" - env_vars: EnvVars = {} - """Extra environment variables for the trainer processes.""" - - tokenizer: TokenizerConfig = TokenizerConfig() - """The tokenizer. Defaults to the one in ``model.name``.""" - data: PretrainDataConfig val: ValConfig | None = None @@ -230,26 +215,6 @@ class PretrainConfig(BaseConfig): ckpt: CheckpointConfig | None = CheckpointConfig() """Checkpointing. Defaults to a checkpoint at the end of training.""" - resume: ResumeConfig | None = None - """Resume the run from a checkpoint. If None, does not resume.""" - - log: TrainerLogConfig = TrainerLogConfig() - - monitors: TrainMonitorsConfig = TrainMonitorsConfig() - """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" - - run: RunConfig = Field(default_factory=RunConfig) - """Run metadata. ``run.name`` names the run directory under ``output_dir``.""" - - output_dir: Path = Field(default_factory=default_output_dir) - """Directory that groups related runs. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``.""" - - clean: bool = False - """Delete the run directory before starting. Required to overwrite a previous run's artifacts when not resuming.""" - - matmul_precision: Literal["highest", "high", "medium"] = "high" - """Precision for float32 matrix multiplications (``torch.set_float32_matmul_precision``).""" - max_steps: int = Field(ge=1) """Training steps.""" @@ -257,38 +222,14 @@ class PretrainConfig(BaseConfig): """CompleteP relative to a base shape. If None, the standard parametrization: the same per-role parameter groups with every multiplier at 1 (norms and biases without weight decay).""" - gc: GCConfig | None = GCConfig() - """Deterministic garbage collection every N steps. Set to null for Python's default GC.""" - - dist_timeout_seconds: int = 900 - """Timeout in seconds for torch distributed ops.""" - - heartbeat: HeartbeatConfig | None = None - """BetterStack heartbeat configuration.""" - - deployment: SFTDeploymentConfig = SingleNodeDeploymentConfig() - """``single_node`` (``num_train_gpus`` on this node) or ``multi_node`` (``num_train_nodes`` x ``gpus_per_node``, needs ``[slurm]``). The online-eval fields (``num_infer_*``) are not used.""" - - slurm: SlurmConfig | None = None - """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" - - dry_run: bool = False - """Only validate and dump resolved configs, then exit.""" - - @property - def run_dir(self) -> Path: - assert self.run.dir is not None # resolved at construction - return self.output_dir / self.run.dir + def default_run_name(self) -> str: + return f"{Path(self.model.name).name}--{uuid.uuid4().hex[:8]}" @model_validator(mode="after") - def resolve_autos(self): - """Resolve the shared ``auto`` fields for pretraining: a trainable fp32 MoE router and weight decay 0.1.""" - if self.model.moe_router_dtype == "auto": - self.model.moe_router_dtype = "float32" - if self.model.freeze_moe_router == "auto": - self.model.freeze_moe_router = False - if self.optim.weight_decay == "auto": - self.optim.weight_decay = 0.1 + def validate_unsupported(self): + unsupported = [name for name in ("trace_path", "memory_profiler_path", "control") if getattr(self, name)] + if unsupported: + raise ValueError(f"Pretraining does not support {', '.join(unsupported)} yet") return self @model_validator(mode="after") @@ -297,59 +238,6 @@ def validate_lr(self): raise ValueError("optim.lr is required") return self - @model_validator(mode="after") - def auto_setup_run_identity(self): - if self.run.name is None: - self.run.name = f"{Path(self.model.name).name}--{uuid.uuid4().hex[:8]}" - if self.run.dir is None: - self.run.dir = self.run.name - if self.monitors.wandb is not None and self.monitors.wandb.name is None: - self.monitors.wandb.name = self.run.name - return self - - @model_validator(mode="after") - def auto_setup_tokenizer(self): - if self.tokenizer.name is None: - self.tokenizer.name = self.model.name - if self.tokenizer.trust_remote_code is None: - self.tokenizer.trust_remote_code = self.model.trust_remote_code - return self - - @model_validator(mode="after") - def auto_setup_deployment(self): - """A multi-node deployment needs SLURM and sets ``model.dp_replicate`` from ``nodes_per_fsdp_group``; the - SLURM job name defaults to the run name and the template to the bundled single- or multi-node template.""" - deployment = self.deployment - if deployment.type == "multi_node": - if self.slurm is None: - raise ValueError("Must use SLURM for multi-node deployment.") - if deployment.nodes_per_fsdp_group is not None: - if deployment.num_train_nodes % deployment.nodes_per_fsdp_group != 0: - raise ValueError( - f"deployment.num_train_nodes ({deployment.num_train_nodes}) must be divisible by " - f"deployment.nodes_per_fsdp_group ({deployment.nodes_per_fsdp_group})" - ) - self.model.dp_replicate = deployment.num_train_nodes // deployment.nodes_per_fsdp_group - if self.slurm is not None: - if self.slurm.job_name is None: - self.slurm.job_name = self.run.name - if self.slurm.template_path is None: - templates_dir = find_package_resource("templates") - if templates_dir is not None: - self.slurm.template_path = templates_dir / f"{deployment.type}_sft.sbatch.j2" - return self - - @model_validator(mode="after") - def deepep_disables_grad_clipping(self): - if self.model.ep != 1 and self.model.moe.dispatch.type == "deepep" and self.optim.max_norm is not None: - warnings.warn( - "Gradient clipping is not compatible with DeepEP. " - "Automatically setting optim.max_norm to None (disabled).", - stacklevel=1, - ) - self.optim.max_norm = None - return self - @model_validator(mode="after") def validate_cp(self): if self.model.cp > 1: @@ -361,11 +249,6 @@ def validate_cp(self): raise ValueError("[evals] does not support context parallelism yet") return self - @model_validator(mode="after") - def validate_scheduler_steps(self): - validate_scheduler(self.scheduler, self.max_steps) - return self - @model_validator(mode="after") def auto_setup_model_seq_len(self): self.model.seq_len = self.data.seq_len diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index ecc5de730ea..f9109e9d086 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -1,6 +1,5 @@ import uuid import warnings -from pathlib import Path from typing import Annotated, Literal, TypeAlias from urllib.parse import urlparse @@ -11,19 +10,13 @@ from prime_rl.configs.eval import SFTOnlineEvalConfig from prime_rl.configs.inference import InferenceConfig from prime_rl.configs.inference import WeightBroadcastConfig as InferenceWeightBroadcastConfig -from prime_rl.configs.monitors import TrainMonitorsConfig from prime_rl.configs.shared import ( FileSystemWeightBroadcastConfig, NCCLWeightBroadcastConfig, - ResumeConfig, - RunConfig, - SFTDeploymentConfig, - SingleNodeDeploymentConfig, - SlurmConfig, WeightBroadcastConfig, ) -from prime_rl.configs.trainer import BaseTrainerConfig -from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource +from prime_rl.configs.trainer import TrainerLaunchConfig +from prime_rl.utils.config import BaseConfig class BaseDataConfig(BaseConfig): @@ -165,7 +158,7 @@ class SFTValConfig(BaseConfig): DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig, Field(discriminator="type")] -class SFTConfig(BaseTrainerConfig): +class SFTConfig(TrainerLaunchConfig): # bf16 routing skips the fp32 gate GEMM and its fp32 FSDP unit. auto_moe_router_dtype = "bfloat16" auto_freeze_moe_router = False @@ -192,69 +185,19 @@ class SFTConfig(BaseTrainerConfig): """Trainer-to-inference weight transport for online evals. Defaults to NCCL. LoRA and external inference use filesystem broadcast.""" - resume: ResumeConfig | None = None - """Resume the run from a checkpoint (point at it with the previous run's ``run.name``). Without ``[ckpt]`` the run loads the checkpoint but saves no new ones. If None, does not resume.""" - - monitors: TrainMonitorsConfig = TrainMonitorsConfig() - """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" - - run: RunConfig = Field(default_factory=RunConfig) - """Run metadata. ``run.name`` names the run directory under ``output_dir``.""" - - output_dir: Path = Field(default_factory=default_output_dir) - """Directory that groups related runs. Each run writes its artifacts (checkpoints, logs, ...) to ``output_dir / run.name``. Should be a persistent directory with enough disk space. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``.""" - - clean: bool = False - """Delete the run directory (``output_dir / run.name``) before starting training. Required to overwrite a run directory that contains artifacts from a previous run when not resuming.""" - - @property - def run_dir(self) -> Path: - assert self.run.dir is not None # resolved at construction - return self.output_dir / self.run.dir - - @model_validator(mode="after") - def auto_setup_run_identity(self): - """Auto-generate the run name (``----``) when unset and - default the run directory, W&B run name and platform run name to it when not - set explicitly.""" - if self.run.name is None: - dataset = str(getattr(self.data, "name", "")).split("/")[-1] - model = self.model.name.split("/")[-1] - parts = [part for part in (dataset, model) if part] - self.run.name = "--".join([*parts, uuid.uuid4().hex[:8]]).lower() - if self.run.dir is None: - self.run.dir = self.run.name - if self.monitors.wandb is not None and self.monitors.wandb.name is None: - self.monitors.wandb.name = self.run.name - if self.monitors.prime is not None and self.monitors.prime.name is None: - self.monitors.prime.name = self.run.name - return self - - deployment: SFTDeploymentConfig = SingleNodeDeploymentConfig() - - slurm: SlurmConfig | None = None - """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" + def default_run_name(self) -> str: + """``----``.""" + dataset = str(getattr(self.data, "name", "")).split("/")[-1] + model = self.model.name.split("/")[-1] + parts = [part for part in (dataset, model) if part] + return "--".join([*parts, uuid.uuid4().hex[:8]]).lower() dashboard: bool = True """Make sure a local dashboard daemon serves this run's output dir (started on demand in interactive sessions; an already-running daemon's URL is logged).""" - dry_run: bool = False - """Only validate and dump resolved configs, then exit early.""" - ### Pre-validation normalization - @model_validator(mode="before") - @classmethod - def normalize_deployment(cls, data): - if not isinstance(data, dict): - return data - deployment = data.get("deployment") - if isinstance(deployment, dict) and deployment.get("type") == "multi_node": - for key in ("num_train_gpus", "num_infer_gpus"): - deployment.pop(key, None) - return data - @model_validator(mode="before") @classmethod def propagate_model_to_inference(cls, data): @@ -273,12 +216,6 @@ def propagate_model_to_inference(cls, data): ### Validate configs (e.g. raise for unsupported (combinations of) configs) - @model_validator(mode="after") - def validate_deployment(self): - if self.deployment.type == "multi_node" and self.slurm is None: - raise ValueError("Must use SLURM for multi-node deployment.") - return self - @model_validator(mode="after") def auto_setup_online_eval(self): """Wire online evals and the trainer-to-inference weight transport.""" @@ -463,30 +400,3 @@ def validate_vlm_constraints(self): if self.val is not None and self.val.data.micro_batch_size != 1: raise ValueError("VLM SFT requires val.data.micro_batch_size = 1.") return self - - ### Auto-setup and validate shared configs - - @model_validator(mode="after") - def auto_setup_deployment(self): - if self.deployment.type == "multi_node": - if self.deployment.nodes_per_fsdp_group is not None: - if self.deployment.num_train_nodes % self.deployment.nodes_per_fsdp_group != 0: - raise ValueError( - f"deployment.num_train_nodes ({self.deployment.num_train_nodes}) must be divisible by " - f"deployment.nodes_per_fsdp_group ({self.deployment.nodes_per_fsdp_group})" - ) - self.model.dp_replicate = self.deployment.num_train_nodes // self.deployment.nodes_per_fsdp_group - return self - - @model_validator(mode="after") - def auto_setup_slurm_template(self): - if self.slurm is not None and self.slurm.job_name is None: - self.slurm.job_name = self.run.name - if self.slurm is not None and self.slurm.template_path is None: - templates_dir = find_package_resource("templates") - if templates_dir is not None: - if self.deployment.type == "single_node": - self.slurm.template_path = templates_dir / "single_node_sft.sbatch.j2" - else: - self.slurm.template_path = templates_dir / "multi_node_sft.sbatch.j2" - return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 90646fa552d..b8a6cdc2074 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -1,12 +1,13 @@ import math import re +import uuid import warnings from pathlib import Path from typing import Annotated, Any, ClassVar, Literal, TypeAlias from pydantic import Field, field_validator, model_validator -from prime_rl.configs.monitors import MonitorsConfig +from prime_rl.configs.monitors import MonitorsConfig, TrainMonitorsConfig from prime_rl.configs.shared import ( BaseModelConfig, ControlConfig, @@ -14,12 +15,16 @@ FileSystemWeightBroadcastConfig, HeartbeatConfig, ResumeConfig, + RunConfig, + SFTDeploymentConfig, + SingleNodeDeploymentConfig, + SlurmConfig, TrainerLogConfig, TransportConfig, WeightBroadcastConfig, ZMQTransportConfig, ) -from prime_rl.utils.config import BaseConfig, default_output_dir +from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource # -- Shared trainer configs (used by both SFT and RL trainers) -- @@ -719,6 +724,98 @@ def auto_setup_tokenizer(self): return self +class TrainerLaunchConfig(BaseTrainerConfig): + """Fields and validators of the trainers that launch on their own (SFT and pretraining): run identity, output + directory, monitors, resume, deployment and SLURM.""" + + resume: ResumeConfig | None = None + """Resume the run from a checkpoint (point at it with the previous run's ``run.name``). Without ``[ckpt]`` the run loads the checkpoint but saves no new ones. If None, does not resume.""" + + monitors: TrainMonitorsConfig = TrainMonitorsConfig() + """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" + + run: RunConfig = Field(default_factory=RunConfig) + """Run metadata. ``run.name`` names the run directory under ``output_dir``.""" + + output_dir: Path = Field(default_factory=default_output_dir) + """Directory that groups related runs. Each run writes its artifacts (checkpoints, logs, ...) to ``output_dir / run.name``. Should be a persistent directory with enough disk space. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``.""" + + clean: bool = False + """Delete the run directory (``output_dir / run.name``) before starting training. Required to overwrite a run directory that contains artifacts from a previous run when not resuming.""" + + deployment: SFTDeploymentConfig = SingleNodeDeploymentConfig() + + slurm: SlurmConfig | None = None + """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" + + dry_run: bool = False + """Only validate and dump resolved configs, then exit early.""" + + @property + def run_dir(self) -> Path: + assert self.run.dir is not None # resolved at construction + return self.output_dir / self.run.dir + + def default_run_name(self) -> str: + return f"{self.model.name.split('/')[-1]}--{uuid.uuid4().hex[:8]}".lower() + + @model_validator(mode="before") + @classmethod + def normalize_deployment(cls, data): + if not isinstance(data, dict): + return data + deployment = data.get("deployment") + if isinstance(deployment, dict) and deployment.get("type") == "multi_node": + for key in ("num_train_gpus", "num_infer_gpus"): + deployment.pop(key, None) + return data + + @model_validator(mode="after") + def auto_setup_run_identity(self): + """Default the run name (``default_run_name``) and the run directory, W&B run name and platform run name to it + when not set explicitly.""" + if self.run.name is None: + self.run.name = self.default_run_name() + if self.run.dir is None: + self.run.dir = self.run.name + if self.monitors.wandb is not None and self.monitors.wandb.name is None: + self.monitors.wandb.name = self.run.name + if self.monitors.prime is not None and self.monitors.prime.name is None: + self.monitors.prime.name = self.run.name + return self + + @model_validator(mode="after") + def validate_deployment(self): + if self.deployment.type == "multi_node" and self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + return self + + @model_validator(mode="after") + def auto_setup_deployment(self): + if self.deployment.type == "multi_node": + if self.deployment.nodes_per_fsdp_group is not None: + if self.deployment.num_train_nodes % self.deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_train_nodes ({self.deployment.num_train_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({self.deployment.nodes_per_fsdp_group})" + ) + self.model.dp_replicate = self.deployment.num_train_nodes // self.deployment.nodes_per_fsdp_group + return self + + @model_validator(mode="after") + def auto_setup_slurm_template(self): + if self.slurm is not None and self.slurm.job_name is None: + self.slurm.job_name = self.run.name + if self.slurm is not None and self.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + if self.deployment.type == "single_node": + self.slurm.template_path = templates_dir / "single_node_sft.sbatch.j2" + else: + self.slurm.template_path = templates_dir / "multi_node_sft.sbatch.j2" + return self + + class TrainerConfig(BaseTrainerConfig): # fp32 routing matches the fp32-routed checkpoints RL trains from (e.g. GLM-5.x). auto_moe_router_dtype = "float32" From 0da07851c8883b3abf1d73ab610ac433c026c8ae Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 19:34:08 +0200 Subject: [PATCH 095/132] refactor: one resolve_resume_step for the RL, SFT and pretrain trainers and launchers The resume step (resume.dir's step, resume.step, or the latest checkpoint) was resolved by five copies of the same branch; utils/pathing.py now has resolve_resume_step next to resolve_latest_ckpt_step. --- src/prime_rl/entrypoints/rl.py | 13 +++---------- src/prime_rl/entrypoints/sft.py | 16 +++------------- src/prime_rl/trainer/pretrain/train.py | 9 ++------- src/prime_rl/trainer/rl/train.py | 11 ++--------- src/prime_rl/trainer/sft/train.py | 11 ++--------- src/prime_rl/utils/pathing.py | 14 +++++++++++++- 6 files changed, 25 insertions(+), 49 deletions(-) diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index e8d015668ed..9f05dcc20ef 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -23,7 +23,7 @@ get_launcher_log_dir, launcher_attempt_env, prepare_attempt_dirs, - resolve_latest_ckpt_step, + resolve_resume_step, validate_run_dir, write_env_server_config, write_launch_artifacts, @@ -502,15 +502,8 @@ def rl(config: RLConfig): # stale — without this, a fresh run in a dirty run dir would pick up rollouts # from a previous run and the orchestrator would see a negative async level. get_logger().info("Starting RL run") - resume_step: int | None = None - if resuming: - if config.resume.dir is not None: - resume_step = config.resume.dir_step - else: - resume_step = config.resume.step - if resume_step is None: - ckpt_base = ckpt_output_dir if ckpt_output_dir is not None else config.run_dir - resume_step = resolve_latest_ckpt_step(get_ckpt_dir(ckpt_base)) + ckpt_base = ckpt_output_dir if ckpt_output_dir is not None else config.run_dir + resume_step = resolve_resume_step(config.resume, get_ckpt_dir(ckpt_base)) if resume_step is not None: get_logger().info(f"Resuming from step {resume_step}, cleaning future rollouts and broadcasts") diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index e48963fb06a..e684d2c3b72 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -23,7 +23,7 @@ get_launcher_dir, get_launcher_log_dir, prepare_attempt_dirs, - resolve_latest_ckpt_step, + resolve_resume_step, validate_run_dir, write_env_server_config, write_launch_artifacts, @@ -60,16 +60,6 @@ def get_ckpt_base(config: SFTConfig) -> Path: return (config.ckpt.output_dir if config.ckpt else None) or config.run_dir -def resolve_resume_step(config: SFTConfig) -> int | None: - if config.resume is None: - return None - if config.resume.dir is not None: - return config.resume.dir_step - if config.resume.step is not None: - return config.resume.step - return resolve_latest_ckpt_step(get_ckpt_dir(get_ckpt_base(config))) - - def build_online_eval_monitors(monitors: TrainMonitorsConfig) -> EvalMonitorsConfig: """The online-eval process's monitors: same W&B and file config as the trainer's, with the platform monitor switched to the evaluation flavor under the same name.""" @@ -92,7 +82,7 @@ def build_online_eval_config(config: SFTConfig) -> SFTOnlineEvalConfig: weight_broadcast=config.weight_broadcast, broadcasts_dir=get_broadcast_dir(config.run_dir), max_steps=config.max_steps, - resume_step=resolve_resume_step(config), + resume_step=resolve_resume_step(config.resume, get_ckpt_dir(get_ckpt_base(config))), # Same run dir as the trainer: the dashboard reads the eval's # file-monitor artifacts from the run directory, so a separate # output dir would hide them. The train-vs-eval record collision @@ -452,7 +442,7 @@ def clean_stale_eval_artifacts(config: SFTConfig) -> None: if os.environ.get("NEVER_CLEAN"): logger.warning("NEVER_CLEAN is set - keeping stale weight broadcasts; the online-eval process may replay them") return - resume_step = resolve_resume_step(config) + resume_step = resolve_resume_step(config.resume, get_ckpt_dir(get_ckpt_base(config))) clean_future_steps(config.run_dir, resume_step if resume_step is not None else -1) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 23149c1a096..7360639ed64 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -74,7 +74,7 @@ from prime_rl.utils.cp import setup_context_parallel, setup_cp_params, shard_for_cp from prime_rl.utils.heartbeat import Heartbeat from prime_rl.utils.logger import setup_logger -from prime_rl.utils.pathing import resolve_latest_ckpt_step +from prime_rl.utils.pathing import resolve_resume_step from prime_rl.utils.process import set_proc_title from prime_rl.utils.utils import clean_exit from prime_rl.utils.weights import resolve_fqn @@ -165,12 +165,7 @@ def train(config: PretrainConfig): grad_accum_steps = total_micro_batches // micro_batches_per_step ckpt_manager = setup_ckpt_manager(config.run_dir, config.ckpt, resume=config.resume) - checkpoint_step = None - if config.resume is not None: - if config.resume.dir is not None: - checkpoint_step = config.resume.dir_step - else: - checkpoint_step = config.resume.step or resolve_latest_ckpt_step(ckpt_manager.ckpt_dir) + checkpoint_step = resolve_resume_step(config.resume, ckpt_manager.ckpt_dir) logger.info(f"Initializing model ({config.model})") # Weights come from the checkpoint when resuming and from the random init otherwise. diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 5f132da55dc..c293a4b6aad 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -73,7 +73,7 @@ from prime_rl.control.server import ControlServer from prime_rl.utils.process import set_proc_title from prime_rl.utils.worker_pool import WorkerPool -from prime_rl.utils.pathing import resolve_latest_ckpt_step, write_resolved_config +from prime_rl.utils.pathing import resolve_resume_step, write_resolved_config from prime_rl.utils.utils import clean_exit @@ -125,17 +125,10 @@ def train(config: TrainerConfig): parallel_dims = get_parallel_dims(config.model) # Check for checkpoint to resume from - checkpoint_step = None logger.info(f"Initializing checkpoint manager ({config.ckpt})") ckpt_manager = setup_ckpt_manager(config.output_dir, config.ckpt, resume=config.resume) - if config.resume is not None: - if config.resume.dir is not None: - checkpoint_step = config.resume.dir_step - else: - checkpoint_step = config.resume.step - if checkpoint_step is None: - checkpoint_step = resolve_latest_ckpt_step(ckpt_manager.ckpt_dir) + checkpoint_step = resolve_resume_step(config.resume, ckpt_manager.ckpt_dir) # Initialize the model logger.info(f"Initializing model ({config.model})") diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 4008e381854..30d3c6ae496 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -13,7 +13,7 @@ import torch from torch.profiler import profile, ProfilerActivity, record_function from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager -from prime_rl.utils.pathing import resolve_latest_ckpt_step +from prime_rl.utils.pathing import resolve_resume_step from prime_rl.configs.sft import SFTConfig from prime_rl.configs.trainer import CheckpointConfig from prime_rl.transports.weights import prune_broadcasts_beyond, setup_weight_sender @@ -131,14 +131,7 @@ def train(config: SFTConfig): logger.info(f"Initializing checkpoint manager ({config.ckpt})") ckpt_manager = setup_ckpt_manager(config.run_dir, config.ckpt, resume=config.resume) - checkpoint_step = None - if config.resume is not None: - if config.resume.dir is not None: - checkpoint_step = config.resume.dir_step - else: - checkpoint_step = config.resume.step - if checkpoint_step is None: - checkpoint_step = resolve_latest_ckpt_step(ckpt_manager.ckpt_dir) + checkpoint_step = resolve_resume_step(config.resume, ckpt_manager.ckpt_dir) # Initialize the model and tokenizer logger.info(f"Initializing model ({config.model})") diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index a02c19f6006..49c087121ba 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -11,7 +11,7 @@ from pydantic import BaseModel from prime_rl.configs.orchestrator import EnvConfig -from prime_rl.configs.shared import LogConfig +from prime_rl.configs.shared import LogConfig, ResumeConfig from prime_rl.utils.config import dump_resolved_config from prime_rl.utils.logger import get_logger @@ -335,6 +335,18 @@ def resolve_latest_ckpt_step(ckpt_dir: Path) -> int | None: return latest_step +def resolve_resume_step(resume: ResumeConfig | None, ckpt_dir: Path) -> int | None: + """The checkpoint step a run resumes from: ``resume.dir``'s step, ``resume.step``, or the latest checkpoint in + ``ckpt_dir``. None without ``resume``.""" + if resume is None: + return None + if resume.dir is not None: + return resume.dir_step + if resume.step is not None: + return resume.step + return resolve_latest_ckpt_step(ckpt_dir) + + def has_checkpoints(output_dir: Path) -> bool: """Check if the output directory contains any checkpoints.""" ckpt_dir = get_ckpt_dir(output_dir) From 5d666025b274f01bd2fcd03d66bf09d0237d10a2 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 19:34:38 +0200 Subject: [PATCH 096/132] refactor(trainer): share the gradient-accumulation step count between SFT and pretraining trainer/utils.py get_grad_accum_steps replaces the identical computation and divisibility assert in both training loops. --- src/prime_rl/trainer/pretrain/train.py | 9 +++------ src/prime_rl/trainer/sft/train.py | 9 +++------ src/prime_rl/trainer/utils.py | 12 ++++++++++++ 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 7360639ed64..a43f547fba8 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -63,6 +63,7 @@ from prime_rl.trainer.sft.data import get_dataset_state from prime_rl.trainer.scheduler import setup_scheduler from prime_rl.trainer.utils import ( + get_grad_accum_steps, get_dcp_group, GarbageCollection, clip_grad_norm_, @@ -156,13 +157,9 @@ def train(config: PretrainConfig): resolve_ep(config.model) parallel_dims = get_parallel_dims(config.model) - total_micro_batches = config.data.batch_size * config.model.cp - micro_batches_per_step = world.world_size * config.data.micro_batch_size - assert total_micro_batches % micro_batches_per_step == 0, ( - f"batch_size * cp ({total_micro_batches}) must be divisible by " - f"world_size * micro_batch_size ({micro_batches_per_step})" + grad_accum_steps = get_grad_accum_steps( + config.data.batch_size, config.data.micro_batch_size, config.model.cp, world.world_size ) - grad_accum_steps = total_micro_batches // micro_batches_per_step ckpt_manager = setup_ckpt_manager(config.run_dir, config.ckpt, resume=config.resume) checkpoint_step = resolve_resume_step(config.resume, ckpt_manager.ckpt_dir) diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 30d3c6ae496..0f60b3d4bc0 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -45,6 +45,7 @@ setup_dataset, ) from prime_rl.trainer.utils import ( + get_grad_accum_steps, GarbageCollection, MemoryProfiler, begin_backward, @@ -116,13 +117,9 @@ def train(config: SFTConfig): # Initialize parallel dimensions parallel_dims = get_parallel_dims(config.model) - total_micro_batches = config.data.batch_size * config.model.cp - micro_batches_per_step = world.world_size * config.data.micro_batch_size - assert total_micro_batches % micro_batches_per_step == 0, ( - f"batch_size * cp ({total_micro_batches}) must be divisible by " - f"world_size * micro_batch_size ({micro_batches_per_step})" + grad_accum_steps = get_grad_accum_steps( + config.data.batch_size, config.data.micro_batch_size, config.model.cp, world.world_size ) - grad_accum_steps = total_micro_batches // micro_batches_per_step if parallel_dims.cp_enabled: assert config.data.seq_len % parallel_dims.cp == 0, "Sequence length must be divisible by CP degree" diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index 3ae52354f7d..01918e51c18 100644 --- a/src/prime_rl/trainer/utils.py +++ b/src/prime_rl/trainer/utils.py @@ -363,3 +363,15 @@ def step(self): f"Finished dumping memory snapshot in {time.monotonic() - begin:.2f} seconds, load {file_path} at https://docs.pytorch.org/memory_viz to visualize the memory usage" ) self.step_num += 1 + + +def get_grad_accum_steps(batch_size: int, micro_batch_size: int, cp: int, world_size: int) -> int: + """Micro-batches each rank runs per optimizer step: ``batch_size`` samples, each split over ``cp`` ranks, in + micro-batches of ``micro_batch_size`` on ``world_size`` ranks.""" + total_micro_batches = batch_size * cp + micro_batches_per_step = world_size * micro_batch_size + assert total_micro_batches % micro_batches_per_step == 0, ( + f"batch_size * cp ({total_micro_batches}) must be divisible by " + f"world_size * micro_batch_size ({micro_batches_per_step})" + ) + return total_micro_batches // micro_batches_per_step From c1e6bf01640e5bc0071b4c35a58c912622072e1a Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 19:36:24 +0200 Subject: [PATCH 097/132] refactor(entrypoints): one prepare_run for the RL, SFT, eval and pretrain launchers utils/pathing.py prepare_run sets the run identity, validates (and with clean, unless NEVER_CLEAN is set, deletes) the run directory and creates it and the checkpoint output directory, replacing four copies. The pretrain launcher now honors NEVER_CLEAN and creates ckpt.output_dir, as the others do, and writes its config with the SFT launcher's write_config. --- src/prime_rl/entrypoints/eval.py | 14 +++----------- src/prime_rl/entrypoints/pretrain.py | 25 +++++-------------------- src/prime_rl/entrypoints/rl.py | 21 +++------------------ src/prime_rl/entrypoints/sft.py | 22 +++------------------- src/prime_rl/utils/pathing.py | 26 +++++++++++++++++++++++++- 5 files changed, 39 insertions(+), 69 deletions(-) diff --git a/src/prime_rl/entrypoints/eval.py b/src/prime_rl/entrypoints/eval.py index 29e128bbc41..0ab6651d486 100644 --- a/src/prime_rl/entrypoints/eval.py +++ b/src/prime_rl/entrypoints/eval.py @@ -11,7 +11,6 @@ import signal import sys import tomllib -import uuid from pathlib import Path from subprocess import Popen from typing import Any @@ -126,20 +125,13 @@ def main(): format_config_message, format_log_message, prepare_attempt_dirs, - validate_run_dir, + prepare_run, write_env_server_config, write_launch_artifacts, ) - # The run identity is runtime-only: $PRL_RUN_ID / $PRL_RUN_NAME are stamped on - # every episode and inherited by the env servers. - os.environ.setdefault("PRL_RUN_ID", uuid.uuid4().hex) - assert config.run.name is not None # resolved at construction - os.environ["PRL_RUN_NAME"] = config.run.name - - clean = config.clean and not os.environ.get("NEVER_CLEAN") - validate_run_dir(config.run_dir, output_dir=config.output_dir, resuming=config.resume is not None, clean=clean) - config.run_dir.mkdir(parents=True, exist_ok=True) + # The run identity is stamped on every episode and inherited by the env servers. + prepare_run(config) config_dir, log_dir = prepare_attempt_dirs(config.run_dir) os.environ["PRL_ATTEMPT_CONFIG_DIR"] = str(config_dir) os.environ["PRL_ATTEMPT_LOG_DIR"] = str(log_dir) diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py index 00d5a80ade0..bbd24036ab2 100644 --- a/src/prime_rl/entrypoints/pretrain.py +++ b/src/prime_rl/entrypoints/pretrain.py @@ -1,18 +1,16 @@ -import json import os -import uuid from pathlib import Path from prime_rl.configs.pretrain import PretrainConfig -from prime_rl.entrypoints.sft import submit_slurm_script, write_trainer_slurm_script -from prime_rl.utils.config import cli, dump_resolved_config +from prime_rl.entrypoints.sft import submit_slurm_script, write_config, write_trainer_slurm_script +from prime_rl.utils.config import cli from prime_rl.utils.logger import get_logger, setup_logger from prime_rl.utils.pathing import ( format_config_message, format_log_message, get_launcher_dir, prepare_attempt_dirs, - validate_run_dir, + prepare_run, write_launch_artifacts, ) from prime_rl.utils.process import ( @@ -42,20 +40,7 @@ def pretrain_slurm(config: PretrainConfig, config_path: Path, log_dir: Path) -> def pretrain(config: PretrainConfig): logger = setup_logger(config.log.level or "info", json_logging=config.log.json_logging) - os.environ.setdefault("PRL_RUN_ID", uuid.uuid4().hex) - assert config.run.name is not None - os.environ["PRL_RUN_NAME"] = config.run.name - - resuming = config.resume is not None - ckpt_output_dir = config.ckpt.output_dir if config.ckpt else None - validate_run_dir( - config.run_dir, - output_dir=config.output_dir, - resuming=resuming, - clean=config.clean, - ckpt_output_dir=ckpt_output_dir, - ) - config.run_dir.mkdir(parents=True, exist_ok=True) + prepare_run(config, config.ckpt.output_dir if config.ckpt else None) logger.info("Starting pretraining run") config_dir, log_dir = prepare_attempt_dirs(config.run_dir) @@ -72,7 +57,7 @@ def pretrain(config: PretrainConfig): # A single-node job runs this launcher again on the node; the trainers of a # multi-node job read the config directly. exclude |= {"slurm", "deployment"} if config.deployment.type == "multi_node" else {"slurm"} - config_path.write_text(json.dumps(dump_resolved_config(config, exclude=exclude), indent=2)) + write_config(config, config_path, exclude=exclude) logger.info(f"Configs:\n{format_config_message(config_dir, 'pretrain', [('Pretrain', config_path)])}") if config.slurm is not None: pretrain_slurm(config, config_path, log_dir) diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 9f05dcc20ef..85ada2f8486 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -2,7 +2,6 @@ import os import subprocess import sys -import uuid from pathlib import Path from urllib.parse import urlparse @@ -23,8 +22,8 @@ get_launcher_log_dir, launcher_attempt_env, prepare_attempt_dirs, + prepare_run, resolve_resume_step, - validate_run_dir, write_env_server_config, write_launch_artifacts, ) @@ -479,23 +478,9 @@ def rl_slurm(config: RLConfig): def rl(config: RLConfig): - # The run identity is runtime-only, never sub-config: $PRL_RUN_ID / $PRL_RUN_NAME are - # the vehicle for runtime info between processes, and every spawned process inherits - # them. Components launched standalone have no run identity. TODO: fetch the id from - # the Prime SDK once runs are registered there. - os.environ.setdefault("PRL_RUN_ID", uuid.uuid4().hex) - assert config.run.name is not None # resolved at construction - os.environ["PRL_RUN_NAME"] = config.run.name - - resuming = config.resume is not None - clean = config.clean and not os.environ.get("NEVER_CLEAN") + # TODO: fetch the run id from the Prime SDK once runs are registered there. ckpt_output_dir = config.ckpt.output_dir if config.ckpt else None - validate_run_dir( - config.run_dir, output_dir=config.output_dir, resuming=resuming, clean=clean, ckpt_output_dir=ckpt_output_dir - ) - config.run_dir.mkdir(parents=True, exist_ok=True) - if ckpt_output_dir is not None: - ckpt_output_dir.mkdir(parents=True, exist_ok=True) + prepare_run(config, ckpt_output_dir) # Clean stale rollouts and broadcasts. When resuming, anything past the resume # step is stale. When training from scratch, every existing step directory is diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index e684d2c3b72..6459768d22e 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -2,7 +2,6 @@ import os import subprocess import sys -import uuid from pathlib import Path from prime_rl.configs.eval import SFTOnlineEvalConfig @@ -23,8 +22,8 @@ get_launcher_dir, get_launcher_log_dir, prepare_attempt_dirs, + prepare_run, resolve_resume_step, - validate_run_dir, write_env_server_config, write_launch_artifacts, ) @@ -108,7 +107,7 @@ def sft_config_components(config: SFTConfig, config_dir: Path) -> list[tuple[str return components -def write_config(config: SFTConfig, config_path: Path, exclude: set[str] | None = None) -> None: +def write_config(config: SFTConfig | PretrainConfig, config_path: Path, exclude: set[str] | None = None) -> None: """Write resolved config to disk, excluding launcher-only fields.""" config_path.parent.mkdir(parents=True, exist_ok=True) with open(config_path, "w") as f: @@ -455,22 +454,7 @@ def sft(config: SFTConfig): "Add an [inference] block or set weight_broadcast.type = 'filesystem'." ) - # The run identity is runtime-only, never sub-config: $PRL_RUN_ID / $PRL_RUN_NAME are - # the vehicle for runtime info between processes, and every spawned process inherits - # them. Components launched standalone have no run identity. - os.environ.setdefault("PRL_RUN_ID", uuid.uuid4().hex) - assert config.run.name is not None # resolved at construction - os.environ["PRL_RUN_NAME"] = config.run.name - - resuming = config.resume is not None - clean = config.clean and not os.environ.get("NEVER_CLEAN") - ckpt_output_dir = config.ckpt.output_dir if config.ckpt else None - validate_run_dir( - config.run_dir, output_dir=config.output_dir, resuming=resuming, clean=clean, ckpt_output_dir=ckpt_output_dir - ) - config.run_dir.mkdir(parents=True, exist_ok=True) - if ckpt_output_dir is not None: - ckpt_output_dir.mkdir(parents=True, exist_ok=True) + prepare_run(config, config.ckpt.output_dir if config.ckpt else None) if config.eval is not None and not config.dry_run: clean_stale_eval_artifacts(config) diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index 49c087121ba..4fb063740f2 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -6,13 +6,14 @@ import sys import tempfile import time +import uuid from pathlib import Path from pydantic import BaseModel from prime_rl.configs.orchestrator import EnvConfig from prime_rl.configs.shared import LogConfig, ResumeConfig -from prime_rl.utils.config import dump_resolved_config +from prime_rl.utils.config import BaseConfig, dump_resolved_config from prime_rl.utils.logger import get_logger @@ -412,6 +413,29 @@ def validate_run_dir( ) +def prepare_run(config: BaseConfig, ckpt_output_dir: Path | None = None) -> None: + """Set the run identity and prepare the run directory of a launcher's ``config`` (``run``, ``run_dir``, + ``output_dir``, ``clean``, ``resume``): validate it (with ``clean``, unless ``$NEVER_CLEAN`` is set, delete it + first) and create it and ``ckpt_output_dir``. + + The run identity is runtime-only, never sub-config: ``$PRL_RUN_ID`` / ``$PRL_RUN_NAME`` carry it between + processes, and every spawned process inherits them. Components launched standalone have no run identity.""" + os.environ.setdefault("PRL_RUN_ID", uuid.uuid4().hex) + assert config.run.name is not None # resolved at construction + os.environ["PRL_RUN_NAME"] = config.run.name + clean = config.clean and not os.environ.get("NEVER_CLEAN") + validate_run_dir( + config.run_dir, + output_dir=config.output_dir, + resuming=config.resume is not None, + clean=clean, + ckpt_output_dir=ckpt_output_dir, + ) + config.run_dir.mkdir(parents=True, exist_ok=True) + if ckpt_output_dir is not None: + ckpt_output_dir.mkdir(parents=True, exist_ok=True) + + def clean_future_steps(output_dir: Path, resume_step: int) -> None: """Remove stale step artifacts past ``resume_step`` and broadcasts from it onward. From 73ae5ef0d667330cb8d93bec5a974706bc00aff8 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 19:36:43 +0200 Subject: [PATCH 098/132] refactor(monitors): register SFT and pretraining runs on the platform through one branch --- src/prime_rl/monitors/prime.py | 18 ++---------------- 1 file changed, 2 insertions(+), 16 deletions(-) diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index d7af0c8ebb0..36384532641 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -157,8 +157,8 @@ async def init(self, config: BaseConfig | None = None, output_dir: Path | None = # attached run's failure/completion marking; a clean finish() is # delivered best-effort and does not by itself mark it completed. init_kwargs = {"id": run_id} - elif isinstance(config, SFTConfig): - # SFT trains on a dataset, not on environments, and has no rollouts - the + elif isinstance(config, SFTConfig | PretrainConfig): + # SFT and pretraining train on data, not on environments, and have no rollouts - the # platform run is registered without both. init_kwargs = dict( name=self.config.name, @@ -172,20 +172,6 @@ async def init(self, config: BaseConfig | None = None, output_dir: Path | None = ), config=config.model_dump(exclude_none=True, mode="json"), ) - elif isinstance(config, PretrainConfig): - # Pretraining trains on token data from scratch: no environments, the model is its architecture. - init_kwargs = dict( - name=self.config.name, - model=config.model.name, - environments=[], - training=pr.TrainingSpec( - max_steps=config.max_steps, - batch_size=config.data.batch_size, - seq_len=config.data.seq_len, - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, - ), - config=config.model_dump(exclude_none=True, mode="json"), - ) elif config is not None: init_kwargs = dict( name=self.config.name, From 77337c4f5c63c269f0692f3385ac8933555803d6 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 19:37:08 +0200 Subject: [PATCH 099/132] refactor(prime_moe): LatentMoE projections in the MoE expert-input/output hooks PrimeMoE's latent down/up projections and norm move from inline branches in forward into the shared MoE's prepare_expert_input/prepare_expert_output hooks (as in NemotronH). Same ops in the same order. --- src/prime_rl/trainer/models/prime_moe/moe.py | 24 ++++++++++++-------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/src/prime_rl/trainer/models/prime_moe/moe.py b/src/prime_rl/trainer/models/prime_moe/moe.py index 3ed046c6cff..6010a521932 100644 --- a/src/prime_rl/trainer/models/prime_moe/moe.py +++ b/src/prime_rl/trainer/models/prime_moe/moe.py @@ -172,6 +172,19 @@ def __init__( self.latent_norm = RMSNorm(RMSNormConfig(hidden_size=latent_dim, eps=norm_eps)) self.latent_up = nn.Linear(latent_dim, dim, bias=False) + def prepare_expert_input(self, x: torch.Tensor) -> torch.Tensor: + if self.latent_down is None: + return x + x = self.latent_down(x) + return x if self.norm_latent_output else self.latent_norm(x) + + def prepare_expert_output(self, x: torch.Tensor) -> torch.Tensor: + if self.latent_up is None: + return x + if self.norm_latent_output: + x = self.latent_norm(x) + return self.latent_up(x) + def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None: # Variable-k routing dispatches one (token, expert) pair per row. token_dispatcher.top_k = 1 if self.variable_k else self.router.top_k @@ -191,11 +204,7 @@ def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) - self.tokens_per_expert, self.routing_confidence_sum, num_tokens_per_expert, routing_confidence_sum ) - expert_input = x - if self.latent_down is not None: - expert_input = self.latent_down(x) - if not self.norm_latent_output: - expert_input = self.latent_norm(expert_input) + expert_input = self.prepare_expert_input(x) if self.variable_k: # Dispatch only the (token, routed expert) pairs, without the null experts. routed = selected_experts_indices < self.experts.num_experts @@ -226,10 +235,7 @@ def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) - shared_output = self.shared_expert(x) if self.shared_expert is not None else None self.token_dispatcher.synchronize() - if self.latent_up is not None: - if self.norm_latent_output: - routed_output = self.latent_norm(routed_output) - routed_output = self.latent_up(routed_output) + routed_output = self.prepare_expert_output(routed_output) if shared_output is not None: routed_output = routed_output + shared_output return routed_output.reshape(bs, slen, dim) From c7e0e5eca887650d0399f2440e0fbe7b13532e9f Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 19:37:40 +0200 Subject: [PATCH 100/132] refactor(prime_moe): XSA and the output gate through a FlashAttention output hook FlashAttention.forward ends in project_output (o_proj of the per-head output); PrimeFlashAttention overrides only that hook for exclusive self attention and the output gate instead of copying forward. Other models run the same ops. --- src/prime_rl/trainer/models/layers/attn.py | 9 +++-- .../trainer/models/prime_moe/attention.py | 35 ++----------------- 2 files changed, 9 insertions(+), 35 deletions(-) diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py index 2e13d63b3ae..8849767c9ec 100644 --- a/src/prime_rl/trainer/models/layers/attn.py +++ b/src/prime_rl/trainer/models/layers/attn.py @@ -528,9 +528,12 @@ def forward( value_states = value_states.transpose(1, 2) out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) - attn_output = out.contiguous().view(1, out.shape[0], -1) - attn_output = self.o_proj(attn_output) - return attn_output, None + return self.project_output(out, value_states[0], hidden_states), None + + def project_output(self, out: torch.Tensor, values: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + """``o_proj`` of the per-head attention output ``out`` ([tokens, heads, dim]). Subclasses that change the + output first also get the values ([tokens, kv_heads, dim]) and the block input.""" + return self.o_proj(out.contiguous().view(1, out.shape[0], -1)) ATTN_IMPL2CLASS = { diff --git a/src/prime_rl/trainer/models/prime_moe/attention.py b/src/prime_rl/trainer/models/prime_moe/attention.py index 238596d9ed8..257237f32b7 100644 --- a/src/prime_rl/trainer/models/prime_moe/attention.py +++ b/src/prime_rl/trainer/models/prime_moe/attention.py @@ -3,7 +3,6 @@ from torch import nn from prime_rl.trainer.models.layers.attn import AttentionConfig, FlashAttention -from prime_rl.trainer.models.layers.rotary_emb import apply_rotary_pos_emb class PrimeFlashAttention(FlashAttention): @@ -20,40 +19,12 @@ def __init__(self, config: AttentionConfig, flash_attn_version: int, *, exclusiv self.exclusive = exclusive self.g_proj = nn.Linear(config.hidden_size, self.qkv_sizes[0], bias=False) if output_gate else None - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - cu_seqlens: torch.LongTensor | None = None, - max_seqlen: int | None = None, - ) -> tuple[torch.Tensor, None]: - input_shape = hidden_states.shape[:-1] - hidden_shape = (*input_shape, -1, self.head_dim) - - query_states, key_states, value_states = self.project_qkv(hidden_states) - if self.use_qk_norm and self.qk_norm_type == "per_layer": - query_states = self.q_norm(query_states) - key_states = self.k_norm(key_states) - query_states = query_states.view(hidden_shape) - key_states = key_states.view(hidden_shape) - value_states = value_states.view(hidden_shape) - if self.use_qk_norm and self.qk_norm_type == "per_head": - query_states = self.q_norm(query_states) - key_states = self.k_norm(key_states) - - if position_embeddings is not None: - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb( - query_states.transpose(1, 2), key_states.transpose(1, 2), cos, sin - ) - query_states, key_states = query_states.transpose(1, 2), key_states.transpose(1, 2) - - out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) + def project_output(self, out: torch.Tensor, values: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: if self.exclusive: # Each query head's own value is its key/value group's. - values = F.normalize(value_states[0], dim=-1).repeat_interleave(self.num_key_value_groups, dim=1) + values = F.normalize(values, dim=-1).repeat_interleave(self.num_key_value_groups, dim=1) out = out - (out * values).sum(dim=-1, keepdim=True) * values attn_output = out.contiguous().view(1, out.shape[0], -1) if self.g_proj is not None: attn_output = attn_output * torch.sigmoid(self.g_proj(hidden_states)) - return self.o_proj(attn_output), None + return self.o_proj(attn_output) From a5679651d7de49d810984b9cef1109501fdc46b2 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 20:21:43 +0200 Subject: [PATCH 101/132] feat(pretrain): evals.micro_batch_tokens sets the eval row length The tokens per packed eval row become a config field (named like data.micro_batch_tokens). It defaults to today's derived value, data.micro_batch_tokens or the longest eval context if longer, and must be at least the longest context. Longer rows can be set by hand; the docs give the measured effect (d512: 32k instead of 16k rows, -19% eval time; eval memory at most 0.6x the training peak). --- docs/pretraining.md | 2 +- .../src/prime_rl/configs/pretrain.py | 14 ++++++++++++++ src/prime_rl/trainer/pretrain/train.py | 4 +--- 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index e6de571a339..b8b1a52a66e 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -149,7 +149,7 @@ interval = 500 web = { path = "data/web-val" } ``` -**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows of `data.micro_batch_tokens` tokens (or the longest of `seq_lens`, if longer) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training source at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. +**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows of `evals.micro_batch_tokens` tokens (default: `data.micro_batch_tokens` or the longest of `seq_lens`, whichever is larger) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training source at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. Evals run without gradients, so longer rows can fit: in our tests, eval memory stayed at most 0.6x the training peak, and at d512 (8k training micro-batches), rows of 32k instead of 16k tokens cut the eval time by 19%. ```toml [evals] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 6e82013fcc2..75e10101d88 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -171,6 +171,10 @@ class EvalsConfig(BaseConfig): ``ctxcont`` sets. Metrics are named ``eval//@`` (e.g. ``bpb@4k``). Requests that are the same at several lengths are scored once.""" + micro_batch_tokens: int | None = Field(None, ge=1) + """Tokens per packed eval row (one forward pass on one GPU). Defaults to ``data.micro_batch_tokens`` or the + longest of ``seq_lens``, whichever is larger; must be at least the longest of ``seq_lens``.""" + save_arrays: bool = False """Write per-document arrays (npz) and summaries to ``/evals/step_N/``.""" @@ -249,6 +253,16 @@ def validate_cp(self): raise ValueError("[evals] does not support context parallelism yet") return self + @model_validator(mode="after") + def auto_setup_eval_micro_batch_tokens(self): + if self.evals is not None: + longest = max(self.evals.seq_lens) + if self.evals.micro_batch_tokens is None: + self.evals.micro_batch_tokens = max(self.data.micro_batch_tokens, longest) + elif self.evals.micro_batch_tokens < longest: + raise ValueError(f"evals.micro_batch_tokens must be at least the longest of evals.seq_lens ({longest})") + return self + @model_validator(mode="after") def auto_setup_model_seq_len(self): self.model.seq_len = self.data.seq_len diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index a43f547fba8..25c887cdf42 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -347,8 +347,6 @@ def run_val(step: int) -> None: for name in config.evals.sets ] logger.info(f"Prepared {len(eval_sets)} eval sets (BOS: {bos}) in {time.perf_counter() - start_time:.0f}s") - # Packed rows as long as a training micro-batch, or the longest eval context if that is longer. - eval_token_budget = max(config.data.micro_batch_tokens, *config.evals.seq_lens) @torch.no_grad() def run_suite(step: int) -> None: @@ -358,7 +356,7 @@ def run_suite(step: int) -> None: metrics = {} # Eager, also for compiled layers: the packed rows vary in length and would recompile at every forward. with scoring_head_active(model), torch.compiler.set_stance("force_eager"): - scored = score_sets(model, eval_sets, eval_token_budget, get_dcp_group()) + scored = score_sets(model, eval_sets, config.evals.micro_batch_tokens, get_dcp_group()) for eval_set, (logprobs, hits) in zip(eval_sets, scored or []): for seq_len, version in eval_set.versions.items(): index = eval_set.index[seq_len] From ae36ea48e37048b2fdb7535220c45cca49c306aa Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 20:22:30 +0200 Subject: [PATCH 102/132] feat(pretrain): quantile balancing on the router logits by default model.moe.balance.bias_update defaults to quantile and PrimeMoeConfig.router_bias_on_logits to true; the sign update stays available. Configs and model directories that relied on the old defaults must set bias_update = "sign" / "router_bias_on_logits": false explicitly. --- docs/pretraining.md | 6 +++--- .../src/prime_rl/configs/pretrain.py | 12 ++++++------ .../models/prime_moe/configuration_prime_moe.py | 4 ++-- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index b8b1a52a66e..e94e5096030 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -43,7 +43,7 @@ MoE models use `prime_moe`, pretraining's own model (`prime_rl.trainer.models.pr - Gated attention (Qwen3-Next): `attn_output_gate` multiplies the attention output elementwise by `sigmoid(g_proj(x))`, computed from the block input, before `o_proj` (after XSA). - LatentMoE: with `latent_dim`, the routed experts work in that width (projection down and RMS norm before the dispatch, projection up after the combine). The router and the shared experts stay full width. With `latent_dim = d/2`, doubling the routed experts, the null experts and `num_experts_per_tok` (and `routed_scaling_factor`, which keeps each expert's gate) keeps the expert parameters and routed FLOPs per token (Kimi K3, NVIDIA LatentMoE). `latent_norm` puts its RMS norm on the expert input (`"input"`, Marin) or on the combined expert output before the up-projection (`"output"`, Kimi K3's Stable LatentMoE). - `expert_activation = "situ_glu"`: Kimi K3's SiTU-GLU in the routed and shared experts, SwiGLU with both factors soft-capped (`4 tanh(g/4) sigmoid(g) * 25 tanh(u/25)`). -- `router_bias_on_logits` adds the selection bias to the router logits instead of the sigmoid scores (for quantile balancing). +- `router_bias_on_logits` (default true) adds the selection bias to the router logits instead of the sigmoid scores (for quantile balancing). The trainer and the HF export (`tools/convert_dcp_to_bf16.py`) read the same model directory and tokenizer; vLLM does not support `prime_moe` yet. The trainer logs the active non-embedding (routed experts counted at `top_k / num_experts`), non-embedding, embedding (input embedding and LM head) and total parameter counts at startup, and as `params/*` at step 0. @@ -104,8 +104,8 @@ The schedule scales every param group's own learning rate, so ratios between gro | Field | Meaning | |---|---| -| `bias_update` | `sign` (default): the update below. `quantile`: quantile balancing. Every step, each layer's bias is set so that every expert would have been selected by `num_experts_per_tok / num_experts` of the step's tokens (null experts included), estimated per micro-batch and averaged over the micro-batches and ranks; no rate. Best with `router_bias_on_logits`. | -| `bias_update_rate` | Every step, each layer's selection bias moves by this step size (gamma) towards underloaded experts, using the step's global expert counts. The update is mean-centred. Default 3e-3; 0 turns it off. The bias starts at 0. | +| `bias_update` | `quantile` (default): quantile balancing. Every step, each layer's bias is set so that every expert would have been selected by `num_experts_per_tok / num_experts` of the step's tokens (null experts included), estimated per micro-batch and averaged over the micro-batches and ranks; no rate. Best with `router_bias_on_logits` (on by default). `sign`: the update below. The bias starts at 0. | +| `bias_update_rate` | With `sign`: every step, each layer's selection bias moves by this step size (gamma) towards underloaded experts, using the step's global expert counts. The update is mean-centred. Default 3e-3; 0 turns it off. | | `bias_freeze` | Fraction of `max_steps` at the end during which the biases stop updating. Default 0.1. | | `seq_aux_coeff` | Weight (alpha) of the sequence-wise balance loss, computed per packed row of `data.seq_len` tokens and added to the training loss. Default 1e-3; 0 turns it off. | diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 75e10101d88..fda42b86d2a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -24,12 +24,12 @@ class MoEBalanceConfig(BaseConfig): """MoE load balancing as in DeepSeek-V3: an aux-loss-free selection bias plus a small sequence-wise balance loss.""" - bias_update: Literal["sign", "quantile"] = "sign" - """``sign``: every step, each layer's bias moves by ``bias_update_rate`` towards underloaded experts. - ``quantile``: quantile balancing (Su 2026; Kimi K3, Marin): every step, each layer's bias is set to the value - with which every expert would have been selected by ``top_k / num_experts`` of the step's tokens (null experts - included), estimated per micro-batch and averaged over the micro-batches and ranks; no rate. Marin found it works - best on the router logits (``router_bias_on_logits`` in the model config).""" + bias_update: Literal["sign", "quantile"] = "quantile" + """``quantile`` (default): quantile balancing (Su 2026; Kimi K3, Marin): every step, each layer's bias is set to + the value with which every expert would have been selected by ``top_k / num_experts`` of the step's tokens (null + experts included), estimated per micro-batch and averaged over the micro-batches and ranks; no rate. Marin found it + works best on the router logits (``router_bias_on_logits`` in the model config, on by default). ``sign``: every + step, each layer's bias moves by ``bias_update_rate`` towards underloaded experts.""" bias_update_rate: float = Field(3e-3, ge=0) """Step size of the sign update (gamma): every step, each layer's bias moves by this towards underloaded experts, diff --git a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py index e8c86efd6ed..759e2fac48f 100644 --- a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py +++ b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py @@ -51,7 +51,7 @@ class PrimeMoeConfig(PretrainedConfig): expert_activation (`str`, *optional*): Activation of the routed and shared experts: `"situ_glu"` (Kimi K3's soft-capped SwiGLU). Defaults to `hidden_act`. router_bias_on_logits (`bool`): Add the selection bias to the router logits instead of the sigmoid scores - (the space quantile balancing works best in). + (the space quantile balancing, the default, works best in). pad_token_id (`int`, *optional*): Padding token id, the input embedding's `padding_idx`. """ @@ -91,7 +91,7 @@ def __init__( latent_dim=None, latent_norm="input", expert_activation=None, - router_bias_on_logits=False, + router_bias_on_logits=True, pad_token_id=None, **kwargs, ): From 5c0bfa6129b4eac0fb833713f0ff38f76f7791d2 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Mon, 19 Oct 2026 20:43:44 +0200 Subject: [PATCH 103/132] fix(trainer): model.debug.random_init initializes the weights load_dcp_from_hf returned right after to_empty() under debug.random_init, so training started from uninitialized memory. It now initializes with models/init.py's init_weights (default InitConfig, no CompleteP multipliers), whose values do not depend on the sharding. SFT/RL behavior changes only with debug.random_init = true (uninitialized memory -> N(0, 0.02) weights, norms at 1, biases at 0). --- src/prime_rl/trainer/model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index b552abc5ed8..4c75244b42e 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -710,6 +710,7 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig): logger = get_logger() if config.debug.random_init: logger.warning("Randomly initializing model. Skipping loading weights from HF.") + init_weights(model, InitConfig()) return if not Path(config.name).exists(): From 2e3601e924335c576299981975a4e96aad460387 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 20 Oct 2026 05:07:15 +0200 Subject: [PATCH 104/132] fix(pretrain): write the last step's checkpoint before its val and evals The last step's checkpoint was written after the loop, i.e. after the final val and eval suite, so a timeout there fell back to the previous checkpoint. It is now saved in the loop like the interval checkpoints (same files), before the step's val and evals. The step's progress counters (total_tokens, total_samples) are updated before the save, so the last checkpoint holds the same values as before and interval checkpoints no longer lag one step. Resuming from the last step's checkpoint runs only its val and evals (instead of training a step past max_steps) and exits. --- src/prime_rl/trainer/pretrain/train.py | 33 ++++++++++++++++---------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 25c887cdf42..92e79003ae7 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -377,6 +377,16 @@ def run_suite(step: int) -> None: run_val(step=0) if config.evals is not None and config.evals.eval_on_start: run_suite(step=checkpoint_step or 0) + if checkpoint_step == config.max_steps: + # The last step's checkpoint is written before its val and evals: resuming from it only runs those. + logger.info(f"Resuming from the last step ({checkpoint_step}): running its val and evals only") + if config.val is not None: + run_val(checkpoint_step) + if config.evals is not None and not config.evals.eval_on_start: + run_suite(checkpoint_step) + logger.success("Pretrain trainer finished") + asyncio.run(monitors.finalize()) + return gc_handler = GarbageCollection(config.gc.interval) if config.gc else None max_memory = torch.cuda.mem_get_info()[1] / 1024**3 @@ -457,10 +467,18 @@ def run_suite(step: int) -> None: "perf/peak_memory/mean": peak_memory["sum"].item() / dist.get_world_size(dp_cp_group), } - # The last step's checkpoint is written once after the loop. + num_tokens = result.num_tokens + progress.total_tokens += num_tokens + step_tokens = source_tokens(train_dataset, next_window, next_window + config.data.batch_size) + next_window += config.data.batch_size + for name, tokens in step_tokens.items(): + source_totals[name] += tokens + progress.total_samples = next_window + + # The last step's checkpoint is written before its val and evals, so a timeout there does not lose it. save_ckpt_time = 0.0 disk_metrics = {} - if config.ckpt and config.ckpt.interval and progress.step % config.ckpt.interval == 0 and not is_last_step: + if config.ckpt and (is_last_step or (config.ckpt.interval and progress.step % config.ckpt.interval == 0)): save_ckpt_start_time = time.perf_counter() ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) ckpt_manager.maybe_clean() @@ -474,13 +492,6 @@ def run_suite(step: int) -> None: run_suite(progress.step) eval_time = time.perf_counter() - eval_start_time - num_tokens = result.num_tokens - progress.total_tokens += num_tokens - step_tokens = source_tokens(train_dataset, next_window, next_window + config.data.batch_size) - next_window += config.data.batch_size - for name, tokens in step_tokens.items(): - source_totals[name] += tokens - progress.total_samples = next_window step_time = time.perf_counter() - step_start_time - eval_time perf_counter = get_perf_counter(model, config.data.seq_len) perf_counter.count_tokens(num_tokens) @@ -536,10 +547,6 @@ def run_suite(step: int) -> None: break progress.step += 1 - if config.ckpt is not None: - logger.info(f"Saving final checkpoint at step {progress.step}") - ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) - ckpt_manager.maybe_clean() logger.success("Pretrain trainer finished") asyncio.run(monitors.finalize()) From e4f9d6abd6364251c486d057757153d5ed7d33bd Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 20 Oct 2026 10:20:08 +0200 Subject: [PATCH 105/132] =?UTF-8?q?pretrain:=20histogram=20estimator=20for?= =?UTF-8?q?=20quantile=20balancing=20(Kimi=20K3=20=C2=A72.3.3),=20default?= =?UTF-8?q?=20on?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds model.moe.balance.quantile_estimator = "histogram" (new default) | "micro_batch_mean" (the previous behavior) and quantile_bins (1024). With "histogram", every forward with grad scatter-adds each expert's required biases (threshold - score) into a per-expert histogram on the previous step's global range ([-1, 1] on the first step), the step end all-reduces the histograms once and reads each expert's top_k / num_experts quantile with in-bin interpolation: the quantile of the pooled step, independent of how tokens are split over micro-batches and ranks. Logs router/qb/bin_width and router/qb/edge_share. Evidence (pretrain-experiments experiments/moe/2026-10-19-qb-histogram): loss within noise at d384 (4 seeds) and d512 (4 seeds); expert balance better at every shape, max/mean -30% / worst layer -38% at the 1B recipe (micro-batch 8,192, 2 seeds); null share on target; +0.5% step time with the sync-free scatter_add_. Default per user decision 2026-10-20. Base: origin/feat/pretrain 2e3601e924; diff sha256 f6e44696... --- docs/pretraining.md | 4 +- .../src/prime_rl/configs/pretrain.py | 12 ++++- src/prime_rl/trainer/models/prime_moe/moe.py | 49 ++++++++++++++++++- src/prime_rl/trainer/pretrain/train.py | 33 ++++++++++++- tests/unit/train/pretrain/test_moe_balance.py | 34 ++++++++++++- 5 files changed, 126 insertions(+), 6 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index e94e5096030..de4c9f5debf 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -104,7 +104,9 @@ The schedule scales every param group's own learning rate, so ratios between gro | Field | Meaning | |---|---| -| `bias_update` | `quantile` (default): quantile balancing. Every step, each layer's bias is set so that every expert would have been selected by `num_experts_per_tok / num_experts` of the step's tokens (null experts included), estimated per micro-batch and averaged over the micro-batches and ranks; no rate. Best with `router_bias_on_logits` (on by default). `sign`: the update below. The bias starts at 0. | +| `bias_update` | `quantile` (default): quantile balancing. Every step, each layer's bias is set so that every expert would have been selected by `num_experts_per_tok / num_experts` of the step's tokens (null experts included), estimated as `quantile_estimator` sets; no rate. Best with `router_bias_on_logits` (on by default). `sign`: the update below. The bias starts at 0. | +| `quantile_estimator` | How `quantile` estimates the thresholds. `histogram` (default; Kimi K3 report, §2.3.3): one histogram of the margins per expert over the whole step, summed over the ranks with one all-reduce, read with in-bin interpolation; the quantile of the pooled step, independent of how the tokens are split. `micro_batch_mean`: each micro-batch's exact quantile, averaged over the micro-batches and ranks. | +| `quantile_bins` | Bins of the `histogram` estimator (default 1024). Its grid is the previous step's global margin range, `[-1, 1]` on the first step after a start or resume. | | `bias_update_rate` | With `sign`: every step, each layer's selection bias moves by this step size (gamma) towards underloaded experts, using the step's global expert counts. The update is mean-centred. Default 3e-3; 0 turns it off. | | `bias_freeze` | Fraction of `max_steps` at the end during which the biases stop updating. Default 0.1. | | `seq_aux_coeff` | Weight (alpha) of the sequence-wise balance loss, computed per packed row of `data.seq_len` tokens and added to the training loss. Default 1e-3; 0 turns it off. | diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index fda42b86d2a..d05aa1c86c6 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -27,10 +27,20 @@ class MoEBalanceConfig(BaseConfig): bias_update: Literal["sign", "quantile"] = "quantile" """``quantile`` (default): quantile balancing (Su 2026; Kimi K3, Marin): every step, each layer's bias is set to the value with which every expert would have been selected by ``top_k / num_experts`` of the step's tokens (null - experts included), estimated per micro-batch and averaged over the micro-batches and ranks; no rate. Marin found it + experts included), estimated as ``quantile_estimator`` sets; no rate. Marin found it works best on the router logits (``router_bias_on_logits`` in the model config, on by default). ``sign``: every step, each layer's bias moves by ``bias_update_rate`` towards underloaded experts.""" + quantile_estimator: Literal["histogram", "micro_batch_mean"] = "histogram" + """How ``quantile`` estimates the thresholds. ``histogram`` (default; Kimi K3 report, §2.3.3): per expert, a + histogram of the step's margins summed over the micro-batches and ranks (one all-reduce per step), read with in-bin + interpolation; the quantile of the pooled step, independent of how the tokens are split. Its grid is the previous + step's global margin range (``[-1, 1]`` on the first step after a start or resume). ``micro_batch_mean``: each + micro-batch's exact quantile, averaged over the micro-batches and ranks.""" + + quantile_bins: int = Field(1024, ge=2) + """Bins of the ``histogram`` estimator's grid.""" + bias_update_rate: float = Field(3e-3, ge=0) """Step size of the sign update (gamma): every step, each layer's bias moves by this towards underloaded experts, mean-centred. 0 turns the update off.""" diff --git a/src/prime_rl/trainer/models/prime_moe/moe.py b/src/prime_rl/trainer/models/prime_moe/moe.py index 6010a521932..fbbaf8f225b 100644 --- a/src/prime_rl/trainer/models/prime_moe/moe.py +++ b/src/prime_rl/trainer/models/prime_moe/moe.py @@ -33,7 +33,9 @@ class PrimeRouter(TokenChoiceTopKRouter): With ``bias_on_logits`` the selection bias is added to the logits instead of the sigmoid scores. With ``quantile_balance`` set by the training loop, every forward with grad records each expert's quantile balancing - threshold (``quantile_sum``, ``quantile_count``), from which the loop sets the next step's bias.""" + threshold (``quantile_sum``, ``quantile_count``), from which the loop sets the next step's bias. If the loop also + sets ``quantile_hist`` (``(num_experts, bins)`` counts) and ``quantile_grid`` (its ``[lo, hi]``), the forward + instead adds each expert's required biases to that histogram (``record_quantile_histogram``).""" def __init__( self, @@ -60,6 +62,9 @@ def __init__( self.quantile_balance = False self.register_buffer("quantile_sum", torch.zeros(self.num_experts, dtype=torch.float32), persistent=False) self.register_buffer("quantile_count", torch.zeros((), dtype=torch.float32), persistent=False) + self.quantile_hist: torch.Tensor | None = None + self.quantile_grid: torch.Tensor | None = None + self.quantile_range: torch.Tensor | None = None @torch.no_grad() def record_quantiles(self, selection_scores: torch.Tensor, biased_scores: torch.Tensor) -> None: @@ -72,6 +77,25 @@ def record_quantiles(self, selection_scores: torch.Tensor, biased_scores: torch. self.quantile_sum += margins.topk(rank, dim=0).values[-1] self.quantile_count += 1 + @torch.no_grad() + def record_quantile_histogram(self, selection_scores: torch.Tensor, biased_scores: torch.Tensor) -> None: + """Histogram estimator of quantile balancing (Kimi K3 report, §2.3.3 and appendix D): add each expert's + required biases ``threshold - score`` (the bias that puts the expert exactly at the token's threshold) to its + row of ``quantile_hist`` on the grid ``quantile_grid``. Values outside the grid count in the edge bins, so the + counts below any in-grid value stay exact. Also tracks the values' ``[min, max]`` in ``quantile_range``.""" + thresholds = biased_scores.topk(self.top_k + 1, dim=1).values[:, -1:] + required = thresholds - selection_scores + num_experts, bins = self.quantile_hist.shape + lo, hi = self.quantile_grid + index = ((required - lo) * (bins / (hi - lo))).floor().clamp_(0, bins - 1).long() + index += torch.arange(num_experts, device=index.device) * bins + # scatter_add_ instead of bincount: CUDA bincount reads its input's maximum to the host, one sync per call. + self.quantile_hist.view(-1).scatter_add_( + 0, index.view(-1), torch.ones_like(index.view(-1), dtype=self.quantile_hist.dtype) + ) + self.quantile_range[0] = torch.minimum(self.quantile_range[0], required.min()) + self.quantile_range[1] = torch.maximum(self.quantile_range[1], required.max()) + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Like ``TokenChoiceTopKRouter.forward`` with sigmoid scores; also returns the ``(bs*slen, num_experts)`` scores of all experts.""" @@ -86,7 +110,10 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Te biased_scores = selection_scores + self.selection_bias _, selected_experts_indices = torch.topk(biased_scores, k=self.top_k, dim=1, sorted=self.topk_sorted) if self.quantile_balance and torch.is_grad_enabled(): - self.record_quantiles(selection_scores, biased_scores) + if self.quantile_hist is None: + self.record_quantiles(selection_scores, biased_scores) + else: + self.record_quantile_histogram(selection_scores, biased_scores) top_scores = scores.gather(dim=1, index=selected_experts_indices) with torch.no_grad(): @@ -102,6 +129,24 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Te return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, scores +def histogram_quantile_biases(hist: torch.Tensor, grid: torch.Tensor, top_k: int) -> torch.Tensor: + """Each expert's ``top_k / num_experts`` quantile of the required biases counted in ``hist`` (``(num_experts, + bins)`` on ``grid = [lo, hi]``), mean-centred: the first bin whose cumulative count reaches the target ``q`` plus + linear interpolation inside it. With it, every expert would have been selected by ``q`` of the counted tokens. + ``q`` comes from each expert's own count, so a forward recomputed under activation checkpointing (which counts + every token twice) leaves the result unchanged.""" + num_experts, bins = hist.shape + cumulative = hist.cumsum(dim=1) + target = cumulative[:, -1] * top_k / num_experts + selected = (cumulative < target.ceil()[:, None]).sum(dim=1).clamp(max=bins - 1)[:, None] + inside = hist.gather(1, selected) + below = cumulative.gather(1, selected) - inside + fraction = ((target[:, None] - below) / inside.clamp(min=1)).clamp(0, 1) + lo, hi = grid + biases = (lo + (selected + fraction) * (hi - lo) / bins).squeeze(1) + return biases - biases.mean() + + def sequence_balance(scores: torch.Tensor, selected_experts: torch.Tensor, num_rows: int) -> torch.Tensor: """DeepSeek-V3's sequence-wise balance loss without its coefficient, per row: ``sum_i f_i P_i`` with ``f_i`` the row's share of expert ``i``'s selections times ``num_experts / top_k`` (1 when balanced) and diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 92e79003ae7..51658a8970f 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -32,6 +32,7 @@ from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter from prime_rl.trainer.models.init import multipliers, param_role +from prime_rl.trainer.models.prime_moe.moe import histogram_quantile_biases from prime_rl.trainer.pretrain.parametrization import param_groups from prime_rl.trainer.pretrain.bytes import token_byte_lengths from prime_rl.trainer.pretrain.data import ( @@ -110,6 +111,28 @@ def quantile_update_selection_biases(model: nn.Module, group: dist.ProcessGroup) router.quantile_count.zero_() +@torch.no_grad() +def histogram_quantile_update_selection_biases(model: nn.Module, group: dist.ProcessGroup) -> dict[str, float]: + """Quantile balancing with the histogram estimator: sum the step's histograms over the ranks, set each layer's + selection bias from them (``histogram_quantile_biases``) and use the step's global range of the required biases + as the next step's grid. Returns the grid's mean bin width and the share of values in its edge bins, which also + hold the values outside it.""" + routers = [moe.router for moe in iter_moe_blocks(model)] + hists = torch.stack([router.quantile_hist for router in routers]) + dist.all_reduce(hists, group=group) + ranges = torch.stack([torch.stack([-router.quantile_range[0], router.quantile_range[1]]) for router in routers]) + dist.all_reduce(ranges, op=dist.ReduceOp.MAX, group=group) + bin_width = clamped = 0.0 + for router, hist, (neg_lo, hi) in zip(routers, hists, ranges): + router.selection_bias.copy_(histogram_quantile_biases(hist, router.quantile_grid, router.top_k)) + bin_width += ((router.quantile_grid[1] - router.quantile_grid[0]) / hist.shape[1]).item() + clamped += ((hist[:, 0].sum() + hist[:, -1].sum()) / hist.sum()).item() + router.quantile_grid = torch.stack([-neg_lo, torch.maximum(hi, -neg_lo + 1e-6)]) + router.quantile_hist.zero_() + router.quantile_range.copy_(torch.tensor([float("inf"), float("-inf")])) + return {"router/qb/bin_width": bin_width / len(routers), "router/qb/edge_share": clamped / len(routers)} + + def count_parameters(model: nn.Module) -> dict[str, int]: """Total, embedding (input embedding and LM head), non-embedding and active non-embedding parameter counts. Active counts the routed experts of each MoE layer at ``top_k / num_experts``, the share each token uses.""" @@ -271,6 +294,12 @@ def train(config: PretrainConfig): moe.router.quantile_sum.zero_() moe.router.quantile_count.zero_() moe.router.quantile_balance = balance.bias_update == "quantile" + if moe.router.quantile_balance and balance.quantile_estimator == "histogram": + device = moe.router.selection_bias.device + moe.router.quantile_hist = torch.zeros(moe.router.num_experts, balance.quantile_bins, device=device) + # The first step's grid; later steps use the previous step's global range. + moe.router.quantile_grid = torch.tensor([-1.0, 1.0], device=device) + moe.router.quantile_range = torch.tensor([float("inf"), float("-inf")], device=device) moes_with_seq_balance = list(iter_moe_blocks(model)) if balance.seq_aux_coeff > 0 else [] for moe in moes_with_seq_balance: # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. @@ -424,7 +453,9 @@ def run_suite(step: int) -> None: updates_bias = balance.bias_update == "quantile" or balance.bias_update_rate > 0 if updates_bias and progress.step <= bias_freeze_step: previous_biases = [bias.clone() for bias in biases] - if balance.bias_update == "quantile": + if balance.bias_update == "quantile" and balance.quantile_estimator == "histogram": + moe_metrics.update(histogram_quantile_update_selection_biases(model, dp_cp_group)) + elif balance.bias_update == "quantile": quantile_update_selection_biases(model, dp_cp_group) else: update_selection_biases(model, result.tokens_per_expert, balance.bias_update_rate) diff --git a/tests/unit/train/pretrain/test_moe_balance.py b/tests/unit/train/pretrain/test_moe_balance.py index bab38a58b71..a2cdf664d84 100644 --- a/tests/unit/train/pretrain/test_moe_balance.py +++ b/tests/unit/train/pretrain/test_moe_balance.py @@ -1,7 +1,7 @@ import torch from prime_rl.trainer.models.layers.moe import MoE, MoEArgs -from prime_rl.trainer.models.prime_moe.moe import PrimeRouter, sequence_balance +from prime_rl.trainer.models.prime_moe.moe import PrimeRouter, histogram_quantile_biases, sequence_balance from prime_rl.trainer.pretrain.train import update_selection_biases @@ -48,3 +48,35 @@ def test_quantile_balancing_balances_the_recorded_batch(): with torch.no_grad(): counts = router(x)[2] assert counts.max() / counts.float().mean() < 1.15 + + +def test_histogram_quantile_balancing_balances_the_recorded_batch_however_it_is_split(): + torch.manual_seed(0) + router = PrimeRouter( + dim=16, num_experts=6, num_null_experts=2, top_k=2, route_norm=True, route_scale=1.0, bias_on_logits=True + ) + with torch.no_grad(): + router.gate.weight[0] += 0.5 # expert 0 is preferred + router.quantile_balance = True + router.quantile_grid = torch.tensor([-4.0, 4.0]) + x = torch.randn(400, 16) + hists = [] + for chunks in (1, 4): + router.quantile_hist = torch.zeros(8, 1024) + router.quantile_range = torch.tensor([float("inf"), float("-inf")]) + for part in x.chunk(chunks): + router(part) + hists.append(router.quantile_hist) + torch.testing.assert_close(hists[0], hists[1]) + # Each token counts once per expert, in the bin of its required bias. + biased = router.gate(x) + router.selection_bias + required = biased.topk(router.top_k + 1, dim=1).values[:, -1:] - router.gate(x) + bins = ((required + 4.0) * (1024 / 8.0)).floor().clamp(0, 1023).long() + expected = torch.stack([torch.bincount(bins[:, j], minlength=1024) for j in range(8)]).float() + torch.testing.assert_close(hists[0], expected) + + # One update brings every expert, null experts included, close to 400 * 2 / 8 = 100 tokens. + router.selection_bias.copy_(histogram_quantile_biases(hists[0], router.quantile_grid, router.top_k)) + with torch.no_grad(): + counts = router(x)[2] + assert counts.max() / counts.float().mean() < 1.15 From 7b4e8ae93370ae83de42dd62c766c9349c32ec20 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 20 Oct 2026 16:57:37 +0200 Subject: [PATCH 106/132] refactor(pretrain): remove CompleteP; plain N(0, 0.02) init and unscaled optimizer groups Every weight matrix and embedding starts from N(0, init.std) with no per-role, width or depth scaling; norms start at 1, biases at 0, and the architectural zero inits stay. The optimizer keeps its parameter groups (Muon for hidden matrices, router and residual outputs; AdamW for the rest; no weight decay on norms and biases), all with the configured lr, eps and weight decay. ParametrizationConfig, [parametrization], multipliers() and role_std are gone; a config that still sets [parametrization] fails validation. Without [parametrization] (the standard parametrization), init and optimizer groups are unchanged. --- examples/pretrain/prime_moe.toml | 6 +- .../src/prime_rl/configs/pretrain.py | 9 +- .../src/prime_rl/configs/trainer.py | 18 +--- src/prime_rl/trainer/model.py | 10 +-- src/prime_rl/trainer/models/init.py | 31 ++----- .../trainer/pretrain/parametrization.py | 65 +++----------- src/prime_rl/trainer/pretrain/train.py | 6 +- tests/unit/train/pretrain/test_init.py | 85 +++---------------- 8 files changed, 44 insertions(+), 186 deletions(-) diff --git a/examples/pretrain/prime_moe.toml b/examples/pretrain/prime_moe.toml index 72848dec750..29335d6fc36 100644 --- a/examples/pretrain/prime_moe.toml +++ b/examples/pretrain/prime_moe.toml @@ -1,7 +1,7 @@ # prime_moe with 64 routed experts (6 active) and 2 shared experts, d = 512 and 8 layers (moe-d512/config.json), trained from scratch on # the Dolma 3 mixture tokenized with the DeepSeek-V4 tokenizer. # 4096 steps x 256K tokens = 1.1B tokens. QK-norm and the learning rate follow the MoE QK-norm sweep (optimum 2^-7 -# with QK-norm at d = 512). [parametrization] with this base shape keeps the LR when the width grows. +# with QK-norm at d = 512). # # uv run pretrain @ examples/pretrain/prime_moe.toml @@ -34,10 +34,6 @@ type = "muon" lr = 0.0078125 max_norm = 1.0 -[parametrization] -base_hidden_size = 512 -base_num_layers = 8 - [val] interval = 512 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index d05aa1c86c6..d2ebd2956e3 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -13,7 +13,6 @@ ModelConfig, MoERuntimeConfig, MuonConfig, - ParametrizationConfig, SchedulerConfig, TrainerLaunchConfig, ) @@ -191,12 +190,12 @@ class EvalsConfig(BaseConfig): class PretrainAdamWConfig(AdamWConfig): eps: float = Field(1e-8, gt=0) - """Adam epsilon of the base shape (scaled per role under CompleteP).""" + """Adam epsilon.""" class PretrainMuonConfig(MuonConfig): eps: float = Field(1e-8, gt=0) - """Adam epsilon of the base shape for the AdamW groups (scaled per role under CompleteP).""" + """Adam epsilon of the AdamW groups.""" PretrainOptimizerConfig: TypeAlias = Annotated[PretrainAdamWConfig | PretrainMuonConfig, Field(discriminator="type")] @@ -232,10 +231,6 @@ class PretrainConfig(TrainerLaunchConfig): max_steps: int = Field(ge=1) """Training steps.""" - parametrization: ParametrizationConfig | None = None - """CompleteP relative to a base shape. If None, the standard parametrization: the same per-role parameter - groups with every multiplier at 1 (norms and biases without weight decay).""" - def default_run_name(self) -> str: return f"{Path(self.model.name).name}--{uuid.uuid4().hex[:8]}" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index b8a6cdc2074..1fd48003647 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -235,28 +235,16 @@ class MoERuntimeConfig(BaseConfig): class InitConfig(BaseConfig): - """Random weight initialization: every weight matrix and the embedding from N(0, std), scaled per role under - CompleteP (see ``[parametrization]``). Norm weights start at 1, biases at 0.""" + """Random weight initialization: every weight matrix and embedding from N(0, std). Norm weights start at 1, + biases at 0.""" std: float = Field(0.02, gt=0) - """Base standard deviation.""" + """Standard deviation.""" seed: int = 0 """Seed for the initialization. The values depend only on the seed, not on the number of GPUs.""" -class ParametrizationConfig(BaseConfig): - """CompleteP (muP extended to depth, Dey et al. 2025, alpha = 1) relative to a base shape: init std, learning - rate, Adam eps and weight decay scale per parameter role with ``m_N = hidden_size / base_hidden_size`` and - ``m_L = num_layers / base_num_layers``, so hyperparameters tuned on the base shape carry over.""" - - base_hidden_size: int = Field(ge=1) - """Hidden size of the base shape.""" - - base_num_layers: int = Field(ge=1) - """Number of layers of the base shape.""" - - class ModelConfig(BaseModelConfig): conversion_dir: Path | None = None """Directory for the auto-converted weights (written to a `prime`/`hf` subdirectory). If not set, we write into the model snapshot directory.""" diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 4c75244b42e..f4908e3c5e6 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -34,7 +34,6 @@ InitConfig, ModelConfig, MXFP8Config, - ParametrizationConfig, TokenizerConfig, ) from prime_rl.multimodal import ForwardPolicy @@ -56,7 +55,7 @@ write_back_loaded_packed_parameters, ) from prime_rl.trainer.models.glm_moe_dsa.sparse_mla_attention import Indexer -from prime_rl.trainer.models.init import init_weights, multipliers +from prime_rl.trainer.models.init import init_weights from prime_rl.trainer.models.layers.fp8_linear import replace_linear_with_fp8_blockwise_linear from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head from prime_rl.trainer.models.layers.moe import MoE, TokenChoiceTopKRouter @@ -940,12 +939,11 @@ def setup_model( parallel_dims: ParallelDims, loading_from_checkpoint_later: bool = False, init: InitConfig | None = None, - parametrization: ParametrizationConfig | None = None, compile_inside_full_checkpoint: bool = False, ) -> nn.Module: """Build the sharded model. Its weights come from ``config.name``, or later from a checkpoint - (``loading_from_checkpoint_later``); with ``init``, they are initialized from scratch (with the CompleteP - multipliers of ``parametrization``) instead. ``compile_inside_full_checkpoint`` compiles blocks inside their + (``loading_from_checkpoint_later``); with ``init``, they are initialized from scratch instead. + ``compile_inside_full_checkpoint`` compiles blocks inside their full activation checkpoints (see ``apply_compile``).""" resolve_auto_attn(config) @@ -1029,7 +1027,7 @@ def setup_model( _reset_runtime_moe_buffers(model) if init is not None: - init_weights(model, init, *multipliers(parametrization, model)) + init_weights(model, init) return model diff --git a/src/prime_rl/trainer/models/init.py b/src/prime_rl/trainer/models/init.py index 84dcb072adf..b969e99e5f7 100644 --- a/src/prime_rl/trainer/models/init.py +++ b/src/prime_rl/trainer/models/init.py @@ -7,7 +7,7 @@ from torch.distributed.tensor import DTensor from torch.distributed.tensor._utils import compute_local_shape_and_global_offset -from prime_rl.configs.trainer import InitConfig, ParametrizationConfig +from prime_rl.configs.trainer import InitConfig from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm from prime_rl.trainer.models.layers.moe import TokenChoiceTopKRouter from prime_rl.utils.weights import resolve_fqn @@ -20,8 +20,8 @@ def param_role(name: str, param: torch.Tensor) -> ParamRole: - """Classify a parameter by its canonical name. Init schemes (and later per-role learning - rates) are defined over these roles.""" + """Classify a parameter by its canonical name: the initialization and the pretraining optimizer groups are defined + over these roles.""" if "embed_tokens" in name: return "embedding" if "lm_head" in name: @@ -37,22 +37,6 @@ def param_role(name: str, param: torch.Tensor) -> ParamRole: return "hidden" -def role_std(role: ParamRole, std: float, m_n: float, m_l: float) -> float: - """Init std of a role under CompleteP with width and depth multipliers ``m_N``, ``m_L`` (both 1 in the - standard parametrization). CompleteP's 1/m_L residual-branch and 1/m_N logit multipliers are folded into the - residual output projections and the LM head, so the forward pass stays unchanged.""" - scale = {"embedding": 1.0, "hidden": m_n**-0.5, "residual_out": m_n**-0.5 / m_l, "output": 1 / m_n} - return std * scale[role] - - -def multipliers(config: ParametrizationConfig | None, model: nn.Module) -> tuple[float, float]: - """Width and depth multipliers ``(m_N, m_L)`` of the model relative to the CompleteP base shape; 1 and 1 in - the standard parametrization.""" - if config is None: - return 1.0, 1.0 - return model.config.hidden_size / config.base_hidden_size, model.config.num_hidden_layers / config.base_num_layers - - def block_seed(seed: int, name: str, block: int) -> int: digest = hashlib.sha256(f"{seed}:{name}:{block}".encode()).digest() return int.from_bytes(digest[:8], "little") & ((1 << 63) - 1) @@ -87,9 +71,10 @@ def normal_shard( @torch.no_grad() -def init_weights(model: nn.Module, config: InitConfig, m_n: float = 1.0, m_l: float = 1.0) -> None: - """Initialize every parameter and the MoE selection biases in place, with the CompleteP multipliers ``m_N``, ``m_L``. Each rank fills only - its local shard, without communication, and the values do not depend on the world size or the sharding.""" +def init_weights(model: nn.Module, config: InitConfig) -> None: + """Initialize every parameter and the MoE selection biases in place: weight matrices and embeddings from + N(0, ``config.std``), norm weights at 1, biases at 0. Each rank fills only its local shard, without communication, + and the values do not depend on the world size or the sharding.""" for name, param in model.named_parameters(): # Canonical names, so wrappers such as activation checkpointing don't change the values. name = resolve_fqn(model, name) @@ -110,7 +95,7 @@ def init_weights(model: nn.Module, config: InitConfig, m_n: float = 1.0, m_l: fl values = normal_shard( name, tuple(param.shape), tuple(local_shape), tuple(global_offset), config.seed, local.device ) - local.copy_(values.mul_(role_std(role, config.std, m_n, m_l))) + local.copy_(values.mul_(config.std)) # MoE selection biases are buffers: they start at 0, so routing starts from the router scores alone. # Gated Residual norm gains are zero-centred (scale 1 + w), so they start at 0. for module in model.modules(): diff --git a/src/prime_rl/trainer/pretrain/parametrization.py b/src/prime_rl/trainer/pretrain/parametrization.py index ef165ba8def..a27effb0bdc 100644 --- a/src/prime_rl/trainer/pretrain/parametrization.py +++ b/src/prime_rl/trainer/pretrain/parametrization.py @@ -6,48 +6,16 @@ from prime_rl.trainer.models.init import param_role from prime_rl.utils.weights import resolve_fqn +# Roles trained by Muon (``rms_norm`` learning-rate adjustment) when ``optim.type = "muon"``; everything else, and +# every role with ``optim.type = "adamw"``, is trained by AdamW. +MUON_ROLES = ("hidden", "router", "residual_out") -def param_groups( - model: nn.Module, optim: PretrainAdamWConfig | PretrainMuonConfig, m_n: float, m_l: float, ep_enabled: bool -) -> list[dict]: - """Optimizer parameter groups per role for CompleteP with alpha = 1 (Dey et al. 2025). ``optim.lr``, - ``optim.eps`` and ``optim.weight_decay`` are the values for the base shape; each role scales - them by the width and depth multipliers. With both multipliers at 1 this is the standard - parametrization, where only norms and biases differ (no weight decay). The residual output projections and the LM head - carry CompleteP's 1/m_L and 1/m_N forward multipliers folded into their weights (see - ``role_std``), which turns their rules into the ones below. - - With Muon, the hidden matrices and residual output projections are Muon groups with - ``rms_norm`` learning-rate adjustment, whose update size grows as sqrt(m_N); their learning - rate is scaled by 1/sqrt(m_N) to compensate (and weight decay by sqrt(m_N), keeping - lr * weight_decay fixed). Everything else is trained with AdamW, as with ``optim.type = "adamw"``. - - MoE experts follow the same roles per expert matrix (``rms_norm`` adjusts each expert by its - own shape). The router maps width to a fixed number of experts, so its spectral-norm target - sqrt(num_experts / d) shrinks as m_N^-1/2: with Muon its learning rate is scaled by 1/m_N; - with AdamW it follows the hidden rule. - - Gated Residual projections (``hyper_connection``) follow the hidden rule but always train with AdamW, as in - Qwen3.8-Next. - """ - # (lr, eps, weight decay) scale per role for AdamW. PyTorch AdamW and Dion decay by lr * weight_decay. - adam_scales = { - "embedding": (1.0, 1 / m_n, 1.0), - "hidden": (1 / m_n, 1 / (m_n * m_l), m_n), - "router": (1 / m_n, 1 / (m_n * m_l), m_n), - "hyper_connection": (1 / m_n, 1 / (m_n * m_l), m_n), - "residual_out": (1 / (m_n * m_l), 1 / m_n, m_n * m_l), - "output": (1 / m_n, 1.0, m_n), - "norm": (1.0, 1 / m_n, 0.0), - "bias": (1.0, 1 / m_n, 0.0), - } - # (lr, weight decay) scale per role for Muon with rms_norm adjustment. - muon_scales = { - "hidden": (m_n**-0.5, m_n**0.5), - "router": (1 / m_n, m_n), - "residual_out": (m_n**-0.5 / m_l, m_n**0.5 * m_l), - } +def param_groups(model: nn.Module, optim: PretrainAdamWConfig | PretrainMuonConfig, ep_enabled: bool) -> list[dict]: + """Optimizer parameter groups per role, all with ``optim.lr`` and ``optim.eps``, and ``optim.weight_decay`` + except for norms and biases (no weight decay). With Muon, the hidden matrices (expert matrices per expert), the + residual output projections and the MoE router are Muon groups; the embedding, the LM head, norms, biases and the + Gated Residual projections (``hyper_connection``, as in Qwen3.8-Next) train with AdamW.""" params = defaultdict(list) for name, param in model.named_parameters(): if param.requires_grad: @@ -61,24 +29,19 @@ def param_groups( groups = [] for (role, experts), role_params in params.items(): - if optim.type == "muon" and role in muon_scales: - lr, wd = muon_scales[role] + weight_decay = 0.0 if role in ("norm", "bias") else optim.weight_decay + if optim.type == "muon" and role in MUON_ROLES: group = dict( - params=role_params, - algorithm="muon", - adjust_lr="rms_norm", - lr=optim.lr * lr, - weight_decay=optim.weight_decay * wd, + params=role_params, algorithm="muon", adjust_lr="rms_norm", lr=optim.lr, weight_decay=weight_decay ) if experts and ep_enabled: group["distributed_mesh_name"] = "dp_shard_mod_ep" groups.append(group) continue - lr, eps, wd = adam_scales[role] - group = dict(params=role_params, lr=optim.lr * lr, weight_decay=optim.weight_decay * wd) + group = dict(params=role_params, lr=optim.lr, weight_decay=weight_decay) if optim.type == "muon": - group.update(algorithm="adamw", epsilon=optim.eps * eps) + group.update(algorithm="adamw", epsilon=optim.eps) else: - group["eps"] = optim.eps * eps + group["eps"] = optim.eps groups.append(group) return groups diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 51658a8970f..bd53ec49b87 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -31,7 +31,7 @@ from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter -from prime_rl.trainer.models.init import multipliers, param_role +from prime_rl.trainer.models.init import param_role from prime_rl.trainer.models.prime_moe.moe import histogram_quantile_biases from prime_rl.trainer.pretrain.parametrization import param_groups from prime_rl.trainer.pretrain.bytes import token_byte_lengths @@ -194,10 +194,8 @@ def train(config: PretrainConfig): parallel_dims, loading_from_checkpoint_later=checkpoint_step is not None, init=config.model.init if checkpoint_step is None else None, - parametrization=config.parametrization, compile_inside_full_checkpoint=True, ) - m_n, m_l = multipliers(config.parametrization, model) if parallel_dims.cp_enabled: setup_context_parallel(model, config.model, parallel_dims) num_params = count_parameters(model) @@ -216,7 +214,7 @@ def train(config: PretrainConfig): list(model.named_parameters()), parallel_dims, model=model, - param_groups=param_groups(model, config.optim, m_n, m_l, parallel_dims.ep_enabled), + param_groups=param_groups(model, config.optim, parallel_dims.ep_enabled), ) # With ckpt.skip_scheduler, the schedule is rebuilt over the remaining steps. diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index e720b18483a..72127015485 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -1,13 +1,9 @@ -import copy import itertools import torch -from torch.optim import AdamW -from transformers import LlamaConfig, LlamaForCausalLM -from prime_rl.configs.pretrain import PretrainAdamWConfig, PretrainMuonConfig -from prime_rl.configs.trainer import InitConfig -from prime_rl.trainer.models.init import init_weights, normal_shard +from prime_rl.configs.pretrain import PretrainMuonConfig +from prime_rl.trainer.models.init import normal_shard from prime_rl.trainer.models.prime_moe import PrimeMoeConfig, PrimeMoeForCausalLM from prime_rl.trainer.pretrain.parametrization import param_groups @@ -28,69 +24,7 @@ def test_normal_shard_is_independent_of_sharding(): assert not torch.equal(normal_shard("w", shape, shape, (0, 0, 0), seed=1, device=torch.device("cpu")), full) -def test_completep_folded_multipliers_match_explicit_multipliers(): - """CompleteP folds its 1/m_L residual and 1/m_N logit multipliers into the weights. Training - the folded model must match the multiplier form of CompleteP (Dey et al. 2025, Table 1).""" - torch.manual_seed(0) - config = LlamaConfig( - vocab_size=128, - hidden_size=64, - intermediate_size=128, - num_hidden_layers=4, - num_attention_heads=2, - num_key_value_heads=1, - head_dim=32, - tie_word_embeddings=False, - ) - optim = PretrainAdamWConfig(lr=1e-2, eps=1e-10, weight_decay=0.1) - m_n, m_l = 2.0, 2.0 - - folded = LlamaForCausalLM(config) - init_weights(folded, InitConfig(), m_n, m_l) - explicit = copy.deepcopy(folded) - with torch.no_grad(): - for name, param in explicit.named_parameters(): - if name.endswith(("o_proj.weight", "down_proj.weight")): - param.mul_(m_l) - elif "lm_head" in name: - param.mul_(m_n) - for layer in explicit.model.layers: - layer.self_attn.o_proj.register_forward_hook(lambda module, args, out: out / m_l) - layer.mlp.down_proj.register_forward_hook(lambda module, args, out: out / m_l) - explicit.lm_head.register_forward_pre_hook(lambda module, args: (args[0] / m_n,)) - - def explicit_scales(name: str) -> tuple[float, float, float]: - if "embed_tokens" in name or "lm_head" in name: - return 1.0, 1 / m_n, 1.0 - if "norm" in name: - return 1.0, 1 / m_n, 0.0 - return 1 / m_n, 1 / (m_n * m_l), m_n - - explicit_groups = [] - for name, param in explicit.named_parameters(): - lr, eps, wd = explicit_scales(name) - explicit_groups.append( - dict(params=[param], lr=optim.lr * lr, eps=optim.eps * eps, weight_decay=optim.weight_decay * wd) - ) - optimizers = { - "explicit": AdamW(explicit_groups, betas=(0.9, 0.95)), - "folded": AdamW(param_groups(folded, optim, m_n, m_l, ep_enabled=False), betas=(0.9, 0.95)), - } - models = {"explicit": explicit, "folded": folded} - - tokens = torch.randint(0, 128, (4, 16)) - for _ in range(5): - for key, model in models.items(): - model(input_ids=tokens, labels=tokens).loss.backward() - optimizers[key].step() - optimizers[key].zero_grad() - with torch.no_grad(): - torch.testing.assert_close( - folded(input_ids=tokens).logits, explicit(input_ids=tokens).logits, rtol=1e-4, atol=1e-4 - ) - - -def test_completep_moe_groups_cover_experts_and_scale_the_router(): +def test_moe_param_groups_cover_every_parameter(): config = PrimeMoeConfig( vocab_size=128, hidden_size=512, @@ -107,10 +41,11 @@ def test_completep_moe_groups_cover_experts_and_scale_the_router(): config._attn_implementation = "flash_attention_2" with torch.device("meta"): model = PrimeMoeForCausalLM._from_config(config) - groups = param_groups(model, PretrainMuonConfig(lr=1.0, weight_decay=0.1), m_n=2.0, m_l=1.0, ep_enabled=False) - lr = {param: group["lr"] for group in groups for param in group["params"]} - assert sum(len(group["params"]) for group in groups) == len(lr) == len(list(model.parameters())) + groups = param_groups(model, PretrainMuonConfig(lr=1.0, weight_decay=0.1), ep_enabled=False) + group_of = {param: group for group in groups for param in group["params"]} + assert sum(len(group["params"]) for group in groups) == len(group_of) == len(list(model.parameters())) moe = model.model.layers[1].mlp - # m_N = 2: Muon matrices get lr / sqrt(m_N) (residual outputs also / m_L = 1), the router lr / m_N. - assert lr[moe.experts.up_proj] == lr[moe.experts.down_proj] == 2**-0.5 - assert lr[moe.router.gate.weight] == 0.5 + assert group_of[moe.experts.up_proj]["algorithm"] == group_of[moe.router.gate.weight]["algorithm"] == "muon" + assert group_of[model.lm_head.weight]["algorithm"] == "adamw" + assert group_of[model.model.norm.weight]["weight_decay"] == 0.0 + assert all(group["lr"] == 1.0 for group in groups) From a2b8e7d6fe46822ebda167f64e35389f03ec1916 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Tue, 20 Oct 2026 16:58:03 +0200 Subject: [PATCH 107/132] docs(pretrain): drop the CompleteP section; describe the plain init and the optimizer groups --- docs/pretraining.md | 32 ++++---------------------------- 1 file changed, 4 insertions(+), 28 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index de4c9f5debf..bcd4437aa4b 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -7,7 +7,7 @@ - [Launch](#launch) - [Model](#model) - [Initialization](#initialization) -- [Parametrization](#parametrization) +- [Optimizer Groups](#optimizer-groups) - [Learning-Rate Schedule](#learning-rate-schedule) - [MoE Load Balancing](#moe-load-balancing) - [Data](#data) @@ -53,36 +53,12 @@ Weights start from a random init, unless the run resumes from a checkpoint. The | Field | Meaning | |---|---| -| `model.init.std` | Every weight matrix and the embedding start from N(0, std) (default 0.02), scaled per role under CompleteP (below). Norm weights start at 1, biases at 0. | +| `model.init.std` | Every weight matrix and embedding (including the LM head, the router, expert and Gated Residual projections) starts from N(0, std) (default 0.02), with no per-role, width or depth scaling. Norm weights start at 1, biases at 0; zero-centred norm gains and the MoE selection biases start at 0. | | `model.init.seed` | Seed. Every rank initializes only its own shard, and the values depend only on the seed and the parameter name, not on the number of GPUs. | -## Parametrization +## Optimizer Groups -The optimizer gets one parameter group per role (embedding, hidden matrices, `o_proj`/`down_proj`, LM head, norms, biases; with MoE also experts and the router). By default (standard parametrization) every role uses `optim.lr`, `optim.eps` and `optim.weight_decay`, except norms and biases, which get no weight decay. - -`[parametrization]` trains with [CompleteP](https://arxiv.org/abs/2505.01618) (muP extended to depth, alpha = 1), so that hyperparameters tuned on a small base shape carry over to wider and deeper models. `optim.lr`, `optim.eps`, `optim.weight_decay` and `model.init.std` are the base-shape values, and the base shape is set by `parametrization.base_hidden_size` and `parametrization.base_num_layers`. With `m_N = hidden_size / base_hidden_size` and `m_L = num_layers / base_num_layers`: - -| Parameters | Init std | Learning rate | Adam eps | Weight decay | -|---|---|---|---|---| -| Embedding | std | lr | eps / m_N | wd | -| Hidden matrices | std / sqrt(m_N) | lr / m_N | eps / (m_N m_L) | wd m_N | -| `o_proj`, `down_proj` | std / (sqrt(m_N) m_L) | lr / (m_N m_L) | eps / m_N | wd m_N m_L | -| LM head | std / m_N | lr / m_N | eps | wd m_N | -| Norm weights | 1 | lr | eps / m_N | 0 | - -CompleteP scales each residual branch by 1/m_L and the logits by 1/m_N. These multipliers are folded into `o_proj`, `down_proj` and the LM head, so the model, its checkpoints and the exported weights are plain Llama. - -With `optim.type = "muon"`, the hidden matrices and `o_proj`/`down_proj` are trained by Muon (`rms_norm` learning-rate adjustment), the rest by AdamW with the rules above. Muon's `rms_norm` update grows as sqrt(m_N), so its groups use: - -| Parameters | Init std | Learning rate | Weight decay | -|---|---|---|---| -| Hidden matrices | std / sqrt(m_N) | lr / sqrt(m_N) | wd sqrt(m_N) | -| `o_proj`, `down_proj` | std / (sqrt(m_N) m_L) | lr / (sqrt(m_N) m_L) | wd sqrt(m_N) m_L | -| MoE router | std / sqrt(m_N) | lr / m_N | wd m_N | - -In MoE models, the expert matrices follow the hidden and `down_proj` rules per expert. The router maps the width to a fixed number of experts, so under Muon its learning rate falls as 1/m_N (with AdamW it follows the hidden rule). - -Use QK-norm (`"use_qk_norm": true` in the model config) with CompleteP. Without it, CompleteP's attention collapses one or two LR steps above the optimum (the norm gains get no weight decay), and the cliff moves down with width. With QK-norm, CompleteP and the standard parametrization reach the same loss, and CompleteP's optimal LR stays fixed across widths (MoE QK-norm sweep). The per-role groups need AdamW or Muon and `model.offload = "none"`. The standard parametrization is CompleteP with the base shape equal to the model shape. Keep `head_dim` and the ratios `intermediate_size / hidden_size` (and `moe_intermediate_size / hidden_size`) and `num_attention_heads / num_key_value_heads` fixed across shapes. +The optimizer gets one parameter group per role (embedding, hidden matrices, `o_proj`/`down_proj`, LM head, norms, biases; with MoE also experts and the router; with Gated Residual its projections). Every group uses `optim.lr`, `optim.eps` and `optim.weight_decay`, except norms and biases, which get no weight decay. With `optim.type = "muon"` (the default), the hidden matrices (expert matrices per expert), `o_proj`/`down_proj` and the MoE router are trained by Muon (`rms_norm` learning-rate adjustment), the rest by AdamW. The groups need AdamW or Muon and `model.offload = "none"`. ## Learning-Rate Schedule From 7aee02b21418881d9ae3dfaa0b1e2f4946c05ff5 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 14:01:04 +0200 Subject: [PATCH 108/132] pretrain: stop copying data.seq_len into model.seq_len Nothing in pretraining reads model.seq_len (the trainer, data, perf counter and evals all read data.seq_len), so it keeps its default as in SFT. --- packages/prime-rl-configs/src/prime_rl/configs/pretrain.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index d2ebd2956e3..a1bafdb71be 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -267,8 +267,3 @@ def auto_setup_eval_micro_batch_tokens(self): elif self.evals.micro_batch_tokens < longest: raise ValueError(f"evals.micro_batch_tokens must be at least the longest of evals.seq_lens ({longest})") return self - - @model_validator(mode="after") - def auto_setup_model_seq_len(self): - self.model.seq_len = self.data.seq_len - return self From e30a595bb442726b4cd78270d3d2d22b4db4e3b2 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 14:03:17 +0200 Subject: [PATCH 109/132] pretrain: batch sizes in rows, as in SFT data.batch_tokens / data.micro_batch_tokens become SFT's data.batch_size (rows per optimizer step, over all GPUs) and data.micro_batch_size (rows per forward pass on one GPU): the pretrain data configs now extend SFT's BaseDataConfig and its validator. data.seq_len defaults to 8192 and evals.seq_lens to data.seq_len. examples/pretrain/prime_moe.toml trains on 8192-token rows at the same 256K tokens per step. --- .../ci/integration/pretrain-moe/resume.toml | 4 +- .../ci/integration/pretrain-moe/start.toml | 4 +- configs/debug/fake/pretrain.toml | 4 +- docs/pretraining.md | 10 ++-- examples/pretrain/prime_moe.toml | 6 +-- .../src/prime_rl/configs/pretrain.py | 49 +++++-------------- src/prime_rl/trainer/pretrain/data.py | 2 +- tests/unit/train/pretrain/test_launcher.py | 2 +- 8 files changed, 29 insertions(+), 52 deletions(-) diff --git a/configs/ci/integration/pretrain-moe/resume.toml b/configs/ci/integration/pretrain-moe/resume.toml index fba58f9aa1e..08ab3124368 100644 --- a/configs/ci/integration/pretrain-moe/resume.toml +++ b/configs/ci/integration/pretrain-moe/resume.toml @@ -22,8 +22,8 @@ name = "PrimeIntellect/Qwen3-0.6B" [data] type = "tokens" seq_len = 2048 -batch_tokens = 16384 -micro_batch_tokens = 4096 +batch_size = 8 +micro_batch_size = 2 max_epochs = 100 weights = { reverse = 1.0 } diff --git a/configs/ci/integration/pretrain-moe/start.toml b/configs/ci/integration/pretrain-moe/start.toml index 85b22407f68..8b3704a2807 100644 --- a/configs/ci/integration/pretrain-moe/start.toml +++ b/configs/ci/integration/pretrain-moe/start.toml @@ -22,8 +22,8 @@ name = "PrimeIntellect/Qwen3-0.6B" [data] type = "tokens" seq_len = 2048 -batch_tokens = 16384 -micro_batch_tokens = 4096 +batch_size = 8 +micro_batch_size = 2 max_epochs = 100 weights = { reverse = 1.0 } diff --git a/configs/debug/fake/pretrain.toml b/configs/debug/fake/pretrain.toml index 0b9de1df8db..1c778cd9a6e 100644 --- a/configs/debug/fake/pretrain.toml +++ b/configs/debug/fake/pretrain.toml @@ -10,8 +10,8 @@ name = "PrimeIntellect/Qwen3-0.6B" [data] type = "fake" seq_len = 512 -batch_tokens = 2048 -micro_batch_tokens = 512 +batch_size = 4 +micro_batch_size = 1 [optim] lr = 1e-3 diff --git a/docs/pretraining.md b/docs/pretraining.md index bcd4437aa4b..fd18f75b580 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -98,9 +98,9 @@ Each row comes from one source, drawn with the mixture `data.weights`. At startu ```toml [data] type = "tokens" -seq_len = 4096 -batch_tokens = 262144 # per optimizer step, over all GPUs -micro_batch_tokens = 65536 # per forward pass on one GPU +seq_len = 8192 # tokens per row +batch_size = 32 # rows per optimizer step, over all GPUs +micro_batch_size = 8 # rows per forward pass on one GPU max_epochs = 3 weights = { web = 0.8, code = 0.2 } @@ -109,7 +109,7 @@ web = { path = "data/web" } code = { path = "data/code" } ``` -Batch sizes are in tokens, so changing `seq_len` keeps the tokens per step and the number of steps. Resume restores the data position on the same GPU layout and `data.num_workers`; resuming on a different layout is not supported yet. `data.type = "fake"` trains on random tokens, for smoke and throughput tests. +Batch sizes are in rows, as in SFT: a step trains on `batch_size * seq_len` tokens, so changing `seq_len` changes the tokens per step. Resume restores the data position on the same GPU layout and `data.num_workers`; resuming on a different layout is not supported yet. `data.type = "fake"` trains on random tokens, for smoke and throughput tests. ## Monitoring @@ -127,7 +127,7 @@ interval = 500 web = { path = "data/web-val" } ``` -**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows of `evals.micro_batch_tokens` tokens (default: `data.micro_batch_tokens` or the longest of `seq_lens`, whichever is larger) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training source at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. Evals run without gradients, so longer rows can fit: in our tests, eval memory stayed at most 0.6x the training peak, and at d512 (8k training micro-batches), rows of 32k instead of 16k tokens cut the eval time by 19%. +**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows of `evals.micro_batch_tokens` tokens (default: the training micro-batch, `data.micro_batch_size * data.seq_len`, or the longest of `seq_lens`, whichever is larger) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` (default: `data.seq_len`) and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training source at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. Evals run without gradients, so longer rows can fit: in our tests, eval memory stayed at most 0.6x the training peak, and at d512 (8k training micro-batches), rows of 32k instead of 16k tokens cut the eval time by 19%. ```toml [evals] diff --git a/examples/pretrain/prime_moe.toml b/examples/pretrain/prime_moe.toml index 29335d6fc36..a7a3680bc26 100644 --- a/examples/pretrain/prime_moe.toml +++ b/examples/pretrain/prime_moe.toml @@ -16,9 +16,9 @@ name = "tokenizers/deepseek-v4-flash-base" # a local copy of deepseek-ai/DeepSe [data] type = "tokens" -seq_len = 4096 -batch_tokens = 262144 -micro_batch_tokens = 65536 +seq_len = 8192 +batch_size = 32 +micro_batch_size = 8 max_epochs = 3 weights = { web = 0.60, pdfs = 0.20, code = 0.12, math = 0.05, reference = 0.03 } diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index a1bafdb71be..fda4daf9784 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -5,6 +5,7 @@ from pydantic import Field, model_validator +from prime_rl.configs.sft import BaseDataConfig from prime_rl.configs.trainer import ( AdamWConfig, CheckpointConfig, @@ -77,40 +78,13 @@ class TokenSourceConfig(BaseConfig): """Tokens available from this source: the first ``max_tokens`` of its stream. Defaults to all tokens in the shards.""" -class BaseBatchConfig(BaseConfig): - seq_len: int = Field(4096, ge=1) +class BaseBatchConfig(BaseDataConfig): + seq_len: int = Field(8192, ge=1) """Tokens per row. Documents continue across rows; attention restarts at every document and every row.""" - batch_tokens: int = Field(ge=1) - """Tokens per optimizer step, over all GPUs. Fixed when ``seq_len`` changes, so the number of steps stays the same.""" - - micro_batch_tokens: int = Field(ge=1) - """Tokens per forward pass on one GPU.""" - - num_workers: int = Field(1, ge=1) - """Dataloader worker processes per GPU.""" - seed: int = 0 """Seed for choosing each row's source, the per-epoch order of every source (and the fake tokens).""" - @property - def batch_size(self) -> int: - """Rows per optimizer step.""" - return self.batch_tokens // self.seq_len - - @property - def micro_batch_size(self) -> int: - """Rows per forward pass.""" - return self.micro_batch_tokens // self.seq_len - - @model_validator(mode="after") - def validate_batch_tokens(self): - if self.batch_tokens % self.seq_len or self.micro_batch_tokens % self.seq_len: - raise ValueError("data.batch_tokens and data.micro_batch_tokens must be multiples of data.seq_len") - if self.batch_tokens % self.micro_batch_tokens: - raise ValueError("data.batch_tokens must be a multiple of data.micro_batch_tokens") - return self - class FakeDataConfig(BaseBatchConfig): """Uniformly random token ids, for smoke tests.""" @@ -175,14 +149,15 @@ class EvalsConfig(BaseConfig): sets: list[str] """Names of the sets to score.""" - seq_lens: list[int] = [2048] + seq_lens: list[int] | None = None """Context lengths to score every set at: the window length for ``ppl`` sets and the context limit for - ``ctxcont`` sets. Metrics are named ``eval//@`` (e.g. ``bpb@4k``). Requests that are the + ``ctxcont`` sets. Defaults to ``data.seq_len``. Metrics are named ``eval//@`` (e.g. ``bpb@4k``). Requests that are the same at several lengths are scored once.""" micro_batch_tokens: int | None = Field(None, ge=1) - """Tokens per packed eval row (one forward pass on one GPU). Defaults to ``data.micro_batch_tokens`` or the - longest of ``seq_lens``, whichever is larger; must be at least the longest of ``seq_lens``.""" + """Tokens per packed eval row (one forward pass on one GPU). Defaults to the training micro-batch + (``data.micro_batch_size * data.seq_len``) or the longest of ``seq_lens``, whichever is larger; must be at least + the longest of ``seq_lens``.""" save_arrays: bool = False """Write per-document arrays (npz) and summaries to ``/evals/step_N/``.""" @@ -253,17 +228,19 @@ def validate_cp(self): if self.data.seq_len % self.model.cp != 0: raise ValueError("data.seq_len must be divisible by model.cp") if self.data.micro_batch_size != 1: - raise ValueError("data.micro_batch_tokens must equal data.seq_len with context parallelism") + raise ValueError("data.micro_batch_size must be 1 with context parallelism") if self.evals is not None: raise ValueError("[evals] does not support context parallelism yet") return self @model_validator(mode="after") - def auto_setup_eval_micro_batch_tokens(self): + def auto_setup_evals(self): if self.evals is not None: + if self.evals.seq_lens is None: + self.evals.seq_lens = [self.data.seq_len] longest = max(self.evals.seq_lens) if self.evals.micro_batch_tokens is None: - self.evals.micro_batch_tokens = max(self.data.micro_batch_tokens, longest) + self.evals.micro_batch_tokens = max(self.data.micro_batch_size * self.data.seq_len, longest) elif self.evals.micro_batch_tokens < longest: raise ValueError(f"evals.micro_batch_tokens must be at least the longest of evals.seq_lens ({longest})") return self diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index ab61dd2549d..07ae88f7b90 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -225,7 +225,7 @@ def setup_train_dataset( weights = config.weights dataset = TokenDataset(sources, weights, config.seq_len, seed=config.seed, non_dp_size=non_dp_size) if config.type == "tokens": - epochs = dataset.planned_epochs(max_steps * config.batch_tokens) + epochs = dataset.planned_epochs(max_steps * config.batch_size * config.seq_len) get_logger().info(f"Planned epochs per source: {epochs}") if over := {name: epoch for name, epoch in epochs.items() if epoch > config.max_epochs}: raise ValueError(f"Planned epochs {over} exceed data.max_epochs={config.max_epochs}") diff --git a/tests/unit/train/pretrain/test_launcher.py b/tests/unit/train/pretrain/test_launcher.py index e712e4f4549..eee8ad647a2 100644 --- a/tests/unit/train/pretrain/test_launcher.py +++ b/tests/unit/train/pretrain/test_launcher.py @@ -13,7 +13,7 @@ def test_multi_node_dry_run_writes_trainer_sbatch(tmp_path, monkeypatch): "output_dir": tmp_path, "model": {"name": "configs/debug/fake/pretrain-model"}, "tokenizer": {"name": "PrimeIntellect/Qwen3-0.6B"}, - "data": {"type": "fake", "seq_len": 8, "batch_tokens": 16, "micro_batch_tokens": 8}, + "data": {"type": "fake", "seq_len": 8, "batch_size": 2, "micro_batch_size": 1}, "optim": {"lr": 1e-3}, "max_steps": 1, "deployment": {"type": "multi_node", "num_train_nodes": 2}, From ef4b34728800203dff7f4c3cdc39ea457693bb3c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 14:04:38 +0200 Subject: [PATCH 110/132] pretrain: data.document_masking On by default (attention and positions restart at every document, as before). With document_masking = false a training row is one sequence: its seq_lens are the row, so FlashAttention's cu_seqlens hold only row boundaries and the positions run over the row. Held-out val windows and the eval suite stay document-masked. --- docs/pretraining.md | 2 +- .../src/prime_rl/configs/pretrain.py | 4 ++++ src/prime_rl/trainer/pretrain/data.py | 20 +++++++++++++------ tests/unit/train/pretrain/test_data.py | 9 +++++++++ 4 files changed, 28 insertions(+), 7 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index fd18f75b580..7cdab54d71e 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -91,7 +91,7 @@ The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` ## Data -`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of a source are read as one stream. Each epoch reads every row of a source once, in a random order over the whole source and from a random start offset, both drawn from `data.seed`, the source and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole source, whatever the order of the files. Every row is a random read (about 33 ms each from a cold cache on CephFS); at the screening shapes, `data.num_workers = 2` with the default prefetch hides this (`time/data` stays at 0). Check `time/data` on other storage. Documents continue across rows, and attention restarts at every document and every row. Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. +`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of a source are read as one stream. Each epoch reads every row of a source once, in a random order over the whole source and from a random start offset, both drawn from `data.seed`, the source and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole source, whatever the order of the files. Every row is a random read (about 33 ms each from a cold cache on CephFS); at the screening shapes, `data.num_workers = 2` with the default prefetch hides this (`time/data` stays at 0). Check `time/data` on other storage. Documents continue across rows, and attention and positions restart at every document and every row; with `data.document_masking = false`, a row is one sequence (attention spans the whole row and positions run over it, as in GPT-style training without document masking). Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. Each row comes from one source, drawn with the mixture `data.weights`. At startup the trainer computes how many passes over each source the run makes and fails above `data.max_epochs`. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index fda4daf9784..b2e09e2f684 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -106,6 +106,10 @@ class TokensDataConfig(BaseBatchConfig): max_epochs: float = Field(1.0, gt=0) """Most passes over a source's available tokens that the run (``max_steps``) may make. Checked at startup.""" + document_masking: bool = True + """Attention and positions restart at every document. If False, a row is one sequence: attention spans the + whole row and positions run over it.""" + @model_validator(mode="after") def validate_weights(self): if any(weight < 0 for weight in self.weights.values()) or not math.isclose(sum(self.weights.values()), 1.0): diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index 07ae88f7b90..10ab2901bd1 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -121,7 +121,8 @@ class TokenDataset(StatefulIterableDataset): Window ``w``'s source is drawn from the mixture ``weights`` with an RNG keyed by ``(seed, w)``, and it is that source's next window (``TokenShards.window_start``). Every data rank walks the whole window sequence but reads only its own windows; the state is the window counter and the tokens read per source. Each window reads one token past - its end for the last target. With ``num_windows`` the dataset stops after that many windows.""" + its end for the last target. With ``num_windows`` the dataset stops after that many windows. Without + ``document_masking``, a window is one sequence: attention and positions run over the whole window.""" def __init__( self, @@ -131,6 +132,7 @@ def __init__( seed: int = 0, non_dp_size: int = 1, num_windows: int | None = None, + document_masking: bool = True, ): super().__init__(non_dp_size) self.sources = sources @@ -140,6 +142,7 @@ def __init__( self.seq_len = seq_len self.seed = seed self.num_windows = num_windows + self.document_masking = document_masking self.tokens = dict.fromkeys(self.names, 0) def epochs(self) -> dict[str, float]: @@ -176,10 +179,8 @@ def __iter__(self): if window % self.data_world_size != self.data_rank: continue tokens, doc_starts = self.sources[name].read(start, self.seq_len + 1) - yield { - "tokens": tokens, - "seq_lens": np.diff([0, *(s for s in doc_starts if s < self.seq_len), self.seq_len]), - } + doc_starts = [s for s in doc_starts if s < self.seq_len] if self.document_masking else [] + yield {"tokens": tokens, "seq_lens": np.diff([0, *doc_starts, self.seq_len])} def source_tokens(dataset: TokenDataset, start: int, end: int) -> dict[str, int]: @@ -223,7 +224,14 @@ def setup_train_dataset( sources = {name: TokenShards(source, seed=(config.seed, name)) for name, source in config.sources.items()} check_eos(sources, eos_token_id) weights = config.weights - dataset = TokenDataset(sources, weights, config.seq_len, seed=config.seed, non_dp_size=non_dp_size) + dataset = TokenDataset( + sources, + weights, + config.seq_len, + seed=config.seed, + non_dp_size=non_dp_size, + document_masking=config.type == "fake" or config.document_masking, + ) if config.type == "tokens": epochs = dataset.planned_epochs(max_steps * config.batch_size * config.seq_len) get_logger().info(f"Planned epochs per source: {epochs}") diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index b2fe0018202..a80fca8eaf7 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -78,6 +78,15 @@ def test_source_tokens_count_every_window_of_the_mixture(tmp_path): assert sum(source_tokens(dataset, 8, 16).values()) == 8 * 4 +def test_without_document_masking_every_window_is_one_sequence(tmp_path): + write_shard(tmp_path / "a" / "00.ds", [[1, 2, 3], [4, 5], [6, 7, 8, 9]]) + sources = {"a": TokenShards(TokenSourceConfig(path=tmp_path / "a"))} + masked = [window["seq_lens"].tolist() for window in TokenDataset(sources, {"a": 1.0}, seq_len=4, num_windows=2)] + assert masked == [[3, 1], [1, 3]] + dataset = TokenDataset(sources, {"a": 1.0}, seq_len=4, num_windows=2, document_masking=False) + assert [window["seq_lens"].tolist() for window in dataset] == [[4], [4]] + + def test_collate_rows_restarts_positions_at_documents(): rows = [ {"tokens": np.arange(5), "seq_lens": np.array([3, 1])}, From 6f2af1aa3a6d99efcd4833ea2c19e94109bc9993 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 14:06:32 +0200 Subject: [PATCH 111/132] configs(sft): pad_to_seq_len on the packed SFT data configs only Pretraining's data configs extend BaseDataConfig, and its rows are always full, so the option did nothing there. FakeDataConfig and SFTDataConfig now get it from PackedDataConfig; their fields, field order, defaults and behavior are unchanged. --- .../src/prime_rl/configs/sft.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index bfaafa87e2b..1e7fe2914f2 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -32,11 +32,6 @@ class BaseDataConfig(BaseConfig): num_workers: int = Field(1, ge=1) """Number of dataloader worker processes. Batches are prepared and pinned in the background so tokenization/packing overlaps with training.""" - pad_to_seq_len: bool = False - """Pad every packed row to ``seq_len * micro_batch_size`` instead of to a multiple of ``model.cp``. Rows then have - one static shape, which avoids recompilation on dynamic sequence lengths (torch.compile, DeepGEMM JIT kernels, ...) - at the cost of extra padding.""" - @model_validator(mode="after") def validate_batch_size(self): if self.batch_size % self.micro_batch_size != 0: @@ -46,7 +41,16 @@ def validate_batch_size(self): return self -class FakeDataConfig(BaseDataConfig): +class PackedDataConfig(BaseDataConfig): + """Samples packed into rows of ``seq_len * micro_batch_size`` tokens.""" + + pad_to_seq_len: bool = False + """Pad every packed row to ``seq_len * micro_batch_size`` instead of to a multiple of ``model.cp``. Rows then have + one static shape, which avoids recompilation on dynamic sequence lengths (torch.compile, DeepGEMM JIT kernels, ...) + at the cost of extra padding.""" + + +class FakeDataConfig(PackedDataConfig): type: Literal["fake"] = "fake" length: Literal["fixed", "variable"] = "fixed" @@ -95,7 +99,7 @@ class SFTColumnsConfig(BaseConfig): """Per-sample renderer arguments as ``renderer field = dataset column``, e.g. ``reasoning_effort = "effort"``. A row's non-null value overrides the ``[renderer]`` setting; rows and datasets without the column use it unchanged.""" -class SFTDataConfig(BaseDataConfig): +class SFTDataConfig(PackedDataConfig): type: Literal["sft"] = "sft" name: str = "PrimeIntellect/Reverse-Text-SFT" From 3db9f440807225ff3ecc596c27a4a5f8a1bb7da1 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 14:07:48 +0200 Subject: [PATCH 112/132] refactor(pretrain): move the MoE balance helpers and count_parameters out of train.py The selection-bias updates and the MoE balance setup go to pretrain/balance.py; count_parameters goes to trainer/utils.py and reuses count_active_params. --- src/prime_rl/trainer/pretrain/balance.py | 81 ++++++++++++++++++ src/prime_rl/trainer/pretrain/train.py | 100 ++--------------------- src/prime_rl/trainer/utils.py | 16 ++++ tests/unit/train/pretrain/test_init.py | 27 ++++++ 4 files changed, 132 insertions(+), 92 deletions(-) create mode 100644 src/prime_rl/trainer/pretrain/balance.py diff --git a/src/prime_rl/trainer/pretrain/balance.py b/src/prime_rl/trainer/pretrain/balance.py new file mode 100644 index 00000000000..a999294495b --- /dev/null +++ b/src/prime_rl/trainer/pretrain/balance.py @@ -0,0 +1,81 @@ +import torch +import torch.distributed as dist +from torch import Tensor, nn + +from prime_rl.configs.pretrain import MoEBalanceConfig +from prime_rl.trainer.model import is_tt_moe_model, iter_moe_blocks +from prime_rl.trainer.models.layers.moe import MoE +from prime_rl.trainer.models.prime_moe.moe import histogram_quantile_biases + + +def setup_moe_balance(model: nn.Module, balance: MoEBalanceConfig, micro_batch_size: int) -> list[MoE]: + """Prepare the MoE layers' load-balancing state for training. Returns the layers with the sequence-wise balance + loss.""" + if is_tt_moe_model(model) and model.config.model_type != "prime_moe": + raise ValueError(f"Pretraining MoE models use model_type 'prime_moe', got {model.config.model_type!r}") + for moe in iter_moe_blocks(model): + # Allocated uninitialized with the rest of the model. + if moe.routed_k_counts is not None: + moe.routed_k_counts.zero_() + moe.router.quantile_sum.zero_() + moe.router.quantile_count.zero_() + moe.router.quantile_balance = balance.bias_update == "quantile" + if moe.router.quantile_balance and balance.quantile_estimator == "histogram": + device = moe.router.selection_bias.device + moe.router.quantile_hist = torch.zeros(moe.router.num_experts, balance.quantile_bins, device=device) + # The first step's grid; later steps use the previous step's global range. + moe.router.quantile_grid = torch.tensor([-1.0, 1.0], device=device) + moe.router.quantile_range = torch.tensor([float("inf"), float("-inf")], device=device) + moes_with_seq_balance = list(iter_moe_blocks(model)) if balance.seq_aux_coeff > 0 else [] + for moe in moes_with_seq_balance: + # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. + moe.balance_rows = micro_batch_size + return moes_with_seq_balance + + +@torch.no_grad() +def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor, rate: float) -> None: + """Aux-loss-free load balancing (https://arxiv.org/abs/2408.15664): move each layer's selection bias by + ``rate`` towards underloaded experts, given the step's global ``[num_moe_layers, num_experts]`` counts. + The update is mean-centred (as in Moonlight and torchtitan), which leaves the selection unchanged but keeps + the biases from drifting together.""" + for moe, load in zip(iter_moe_blocks(model), tokens_per_expert): + if moe.router.selection_bias is not None: + delta = torch.sign(load.mean() - load) + moe.router.selection_bias += rate * (delta - delta.mean()) + + +@torch.no_grad() +def quantile_update_selection_biases(model: nn.Module, group: dist.ProcessGroup) -> None: + """Quantile balancing: set each layer's selection bias to minus the step's balancing thresholds of its experts + (``PrimeRouter.record_quantiles``), averaged over the micro-batches and ranks and mean-centred.""" + routers = [moe.router for moe in iter_moe_blocks(model)] + sums = torch.stack([torch.cat([router.quantile_sum, router.quantile_count[None]]) for router in routers]) + dist.all_reduce(sums, group=group) + for router, row in zip(routers, sums): + thresholds = row[:-1] / row[-1] + router.selection_bias.copy_(thresholds.mean() - thresholds) + router.quantile_sum.zero_() + router.quantile_count.zero_() + + +@torch.no_grad() +def histogram_quantile_update_selection_biases(model: nn.Module, group: dist.ProcessGroup) -> dict[str, float]: + """Quantile balancing with the histogram estimator: sum the step's histograms over the ranks, set each layer's + selection bias from them (``histogram_quantile_biases``) and use the step's global range of the required biases + as the next step's grid. Returns the grid's mean bin width and the share of values in its edge bins, which also + hold the values outside it.""" + routers = [moe.router for moe in iter_moe_blocks(model)] + hists = torch.stack([router.quantile_hist for router in routers]) + dist.all_reduce(hists, group=group) + ranges = torch.stack([torch.stack([-router.quantile_range[0], router.quantile_range[1]]) for router in routers]) + dist.all_reduce(ranges, op=dist.ReduceOp.MAX, group=group) + bin_width = clamped = 0.0 + for router, hist, (neg_lo, hi) in zip(routers, hists, ranges): + router.selection_bias.copy_(histogram_quantile_biases(hist, router.quantile_grid, router.top_k)) + bin_width += ((router.quantile_grid[1] - router.quantile_grid[0]) / hist.shape[1]).item() + clamped += ((hist[:, 0].sum() + hist[:, -1].sum()) / hist.sum()).item() + router.quantile_grid = torch.stack([-neg_lo, torch.maximum(hi, -neg_lo + 1e-6)]) + router.quantile_hist.zero_() + router.quantile_range.copy_(torch.tensor([float("inf"), float("-inf")])) + return {"router/qb/bin_width": bin_width / len(routers), "router/qb/edge_share": clamped / len(routers)} diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 35ea73e5c14..3d13cd7cd36 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -12,7 +12,6 @@ import numpy as np import torch import torch.distributed as dist -from torch import Tensor, nn from torch.distributed.tensor import DTensor from prime_rl import monitors @@ -31,9 +30,13 @@ from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter -from prime_rl.trainer.models.init import param_role -from prime_rl.trainer.models.prime_moe.moe import histogram_quantile_biases from prime_rl.trainer.pretrain.parametrization import param_groups +from prime_rl.trainer.pretrain.balance import ( + histogram_quantile_update_selection_biases, + quantile_update_selection_biases, + setup_moe_balance, + update_selection_biases, +) from prime_rl.trainer.pretrain.bytes import token_byte_lengths from prime_rl.trainer.pretrain.data import ( TokenShards, @@ -64,6 +67,7 @@ from prime_rl.trainer.sft.data import get_dataset_state from prime_rl.trainer.scheduler import setup_scheduler from prime_rl.trainer.utils import ( + count_parameters, get_grad_accum_steps, get_dcp_group, GarbageCollection, @@ -79,81 +83,11 @@ from prime_rl.utils.pathing import resolve_resume_step from prime_rl.utils.process import set_proc_title from prime_rl.utils.utils import clean_exit -from prime_rl.utils.weights import resolve_fqn MOE_TRACE_INTERVAL = 32 -@torch.no_grad() -def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor, rate: float) -> None: - """Aux-loss-free load balancing (https://arxiv.org/abs/2408.15664): move each layer's selection bias by - ``rate`` towards underloaded experts, given the step's global ``[num_moe_layers, num_experts]`` counts. - The update is mean-centred (as in Moonlight and torchtitan), which leaves the selection unchanged but keeps - the biases from drifting together.""" - for moe, load in zip(iter_moe_blocks(model), tokens_per_expert): - if moe.router.selection_bias is not None: - delta = torch.sign(load.mean() - load) - moe.router.selection_bias += rate * (delta - delta.mean()) - - -@torch.no_grad() -def quantile_update_selection_biases(model: nn.Module, group: dist.ProcessGroup) -> None: - """Quantile balancing: set each layer's selection bias to minus the step's balancing thresholds of its experts - (``PrimeRouter.record_quantiles``), averaged over the micro-batches and ranks and mean-centred.""" - routers = [moe.router for moe in iter_moe_blocks(model)] - sums = torch.stack([torch.cat([router.quantile_sum, router.quantile_count[None]]) for router in routers]) - dist.all_reduce(sums, group=group) - for router, row in zip(routers, sums): - thresholds = row[:-1] / row[-1] - router.selection_bias.copy_(thresholds.mean() - thresholds) - router.quantile_sum.zero_() - router.quantile_count.zero_() - - -@torch.no_grad() -def histogram_quantile_update_selection_biases(model: nn.Module, group: dist.ProcessGroup) -> dict[str, float]: - """Quantile balancing with the histogram estimator: sum the step's histograms over the ranks, set each layer's - selection bias from them (``histogram_quantile_biases``) and use the step's global range of the required biases - as the next step's grid. Returns the grid's mean bin width and the share of values in its edge bins, which also - hold the values outside it.""" - routers = [moe.router for moe in iter_moe_blocks(model)] - hists = torch.stack([router.quantile_hist for router in routers]) - dist.all_reduce(hists, group=group) - ranges = torch.stack([torch.stack([-router.quantile_range[0], router.quantile_range[1]]) for router in routers]) - dist.all_reduce(ranges, op=dist.ReduceOp.MAX, group=group) - bin_width = clamped = 0.0 - for router, hist, (neg_lo, hi) in zip(routers, hists, ranges): - router.selection_bias.copy_(histogram_quantile_biases(hist, router.quantile_grid, router.top_k)) - bin_width += ((router.quantile_grid[1] - router.quantile_grid[0]) / hist.shape[1]).item() - clamped += ((hist[:, 0].sum() + hist[:, -1].sum()) / hist.sum()).item() - router.quantile_grid = torch.stack([-neg_lo, torch.maximum(hi, -neg_lo + 1e-6)]) - router.quantile_hist.zero_() - router.quantile_range.copy_(torch.tensor([float("inf"), float("-inf")])) - return {"router/qb/bin_width": bin_width / len(routers), "router/qb/edge_share": clamped / len(routers)} - - -def count_parameters(model: nn.Module) -> dict[str, int]: - """Total, embedding (input embedding and LM head), non-embedding and active non-embedding parameter counts. - Active counts the routed experts of each MoE layer at ``top_k / num_experts``, the share each token uses.""" - total = embedding = 0 - for name, param in model.named_parameters(): - total += param.numel() - if param_role(resolve_fqn(model, name), param) in ("embedding", "output"): - embedding += param.numel() - inactive = 0 - for moe in iter_moe_blocks(model): - experts = sum(param.numel() for param in moe.experts.parameters()) - inactive += experts - experts * moe.router.top_k // moe.router.num_experts - non_embedding = total - embedding - return { - "total": total, - "embedding": embedding, - "non_embedding": non_embedding, - "active_non_embedding": non_embedding - inactive, - } - - @clean_exit def train(config: PretrainConfig): world = get_world() @@ -283,26 +217,8 @@ def train(config: PretrainConfig): cp_size = parallel_dims.cp dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() is_moe_model = is_tt_moe_model(model) - if is_moe_model and model.config.model_type != "prime_moe": - raise ValueError(f"Pretraining MoE models use model_type 'prime_moe', got {model.config.model_type!r}") balance = config.model.moe.balance - for moe in iter_moe_blocks(model): - # Allocated uninitialized with the rest of the model. - if moe.routed_k_counts is not None: - moe.routed_k_counts.zero_() - moe.router.quantile_sum.zero_() - moe.router.quantile_count.zero_() - moe.router.quantile_balance = balance.bias_update == "quantile" - if moe.router.quantile_balance and balance.quantile_estimator == "histogram": - device = moe.router.selection_bias.device - moe.router.quantile_hist = torch.zeros(moe.router.num_experts, balance.quantile_bins, device=device) - # The first step's grid; later steps use the previous step's global range. - moe.router.quantile_grid = torch.tensor([-1.0, 1.0], device=device) - moe.router.quantile_range = torch.tensor([float("inf"), float("-inf")], device=device) - moes_with_seq_balance = list(iter_moe_blocks(model)) if balance.seq_aux_coeff > 0 else [] - for moe in moes_with_seq_balance: - # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. - moe.balance_rows = config.data.micro_batch_size + moes_with_seq_balance = setup_moe_balance(model, balance, config.data.micro_batch_size) bias_freeze_step = config.max_steps - round(balance.bias_freeze * config.max_steps) # Sum and count of the per-row balance of every MoE layer over the current step's micro-batches, reduced with # the step's statistics. diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index 01918e51c18..20caa0295a0 100644 --- a/src/prime_rl/trainer/utils.py +++ b/src/prime_rl/trainer/utils.py @@ -17,6 +17,7 @@ from torch.distributed.tensor import DTensor from transformers.tokenization_utils import PreTrainedTokenizer +from prime_rl.trainer.models.base import count_active_params from prime_rl.trainer.world import get_world from prime_rl.utils.logger import format_time, get_logger from prime_rl.utils.pathing import get_ckpt_dir @@ -375,3 +376,18 @@ def get_grad_accum_steps(batch_size: int, micro_batch_size: int, cp: int, world_ f"world_size * micro_batch_size ({micro_batches_per_step})" ) return total_micro_batches // micro_batches_per_step + + +def count_parameters(model: nn.Module) -> dict[str, int]: + """Total, embedding (input embedding and LM head), non-embedding and active non-embedding parameter counts.""" + embedding_ids = {id(p) for p in (*model.get_input_embeddings().parameters(), *model.lm_head.parameters())} + params = list(model.parameters()) + non_embedding = [p for p in params if id(p) not in embedding_ids] + total = sum(p.numel() for p in params) + num_non_embedding = sum(p.numel() for p in non_embedding) + return { + "total": total, + "embedding": total - num_non_embedding, + "non_embedding": num_non_embedding, + "active_non_embedding": count_active_params(model, non_embedding), + } diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index 82e14bd6b4e..5b9373395f6 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -6,6 +6,7 @@ from prime_rl.trainer.models.init import normal_shard from prime_rl.trainer.models.prime_moe import PrimeMoeConfig, PrimeMoeForCausalLM from prime_rl.trainer.pretrain.parametrization import param_groups +from prime_rl.trainer.utils import count_parameters def test_normal_shard_is_independent_of_sharding(): @@ -24,6 +25,32 @@ def test_normal_shard_is_independent_of_sharding(): assert not torch.equal(normal_shard("w", shape, shape, (0, 0, 0), seed=1, device=torch.device("cpu")), full) +def test_count_parameters_moe(): + config = PrimeMoeConfig( + vocab_size=128, + hidden_size=512, + intermediate_size=1024, + moe_intermediate_size=128, + num_attention_heads=8, + num_key_value_heads=2, + n_routed_experts=8, + num_experts_per_tok=2, + n_shared_experts=1, + num_hidden_layers=2, + first_k_dense_replace=1, + ) + config._attn_implementation = "flash_attention_2" + with torch.device("meta"): + model = PrimeMoeForCausalLM._from_config(config) + counts = count_parameters(model) + embedding = model.model.embed_tokens.weight.numel() + model.lm_head.weight.numel() + experts = sum(param.numel() for param in model.model.layers[1].mlp.experts.parameters()) + assert counts["total"] == sum(param.numel() for param in model.parameters()) + assert counts["embedding"] == embedding + assert counts["non_embedding"] == counts["total"] - embedding + assert counts["active_non_embedding"] == counts["non_embedding"] - experts * 6 // 8 + + def test_moe_param_groups_cover_every_parameter(): config = PrimeMoeConfig( vocab_size=128, From 4aa640727b09f0abd0837bb3e4a93c7ae96d2ab0 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 14:07:54 +0200 Subject: [PATCH 113/132] refactor(pretrain): use the shared optimizer groups and drop the init roles Pretraining builds its optimizer with the default groups of SFT and RL instead of its own per-role groups, and the init classifies 1-D parameters directly, with the same values. --- docs/pretraining.md | 6 +-- .../src/prime_rl/configs/pretrain.py | 4 +- src/prime_rl/trainer/models/init.py | 34 +++----------- .../trainer/pretrain/parametrization.py | 47 ------------------- src/prime_rl/trainer/pretrain/train.py | 2 - tests/unit/train/pretrain/test_init.py | 29 ------------ 6 files changed, 12 insertions(+), 110 deletions(-) delete mode 100644 src/prime_rl/trainer/pretrain/parametrization.py diff --git a/docs/pretraining.md b/docs/pretraining.md index b1b337dfbab..7dc33e1f98e 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -53,12 +53,12 @@ Weights start from a random init, unless the run resumes from a checkpoint. The | Field | Meaning | |---|---| -| `model.init.std` | Every weight matrix and embedding (including the LM head, the router, expert and Gated Residual projections) starts from N(0, std) (default 0.02), with no per-role, width or depth scaling. Norm weights start at 1, biases at 0; zero-centred norm gains and the MoE selection biases start at 0. | +| `model.init.std` | Every weight matrix and embedding (including the LM head, the router, expert and Gated Residual projections) starts from N(0, std) (default 0.02), with no width or depth scaling. Norm weights start at 1, biases at 0; zero-centred norm gains and the MoE selection biases start at 0. | | `model.init.seed` | Seed. Every rank initializes only its own shard, and the values depend only on the seed and the parameter name, not on the number of GPUs. | -## Optimizer Groups +## Optimizer -The optimizer gets one parameter group per role (embedding, hidden matrices, `o_proj`/`down_proj`, LM head, norms, biases; with MoE also experts and the router; with Gated Residual its projections). Every group uses `optim.lr`, `optim.eps` and `optim.weight_decay`, except norms and biases, which get no weight decay. With `optim.type = "muon"` (the default), the hidden matrices (expert matrices per expert), `o_proj`/`down_proj` and the MoE router are trained by Muon (`rms_norm` learning-rate adjustment), the rest by AdamW. The groups need AdamW or Muon. `model.offload = "optimizer"` keeps the optimizer states in CPU RAM between steps, as in SFT and RL; `"full"` is not supported. +Pretraining uses the optimizer setup of SFT and RL. With `optim.type = "muon"` (the default), every weight matrix except the embedding and the LM head is trained by Muon (`rms_norm` learning-rate adjustment; expert matrices per expert), and the embedding, the LM head, norms and biases by AdamW with `optim.eps`. `optim.type = "adamw"` trains everything with AdamW. `model.offload = "optimizer"` keeps the optimizer states in CPU RAM between steps; `"full"` is not supported. ## Learning-Rate Schedule diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 5edec6477a0..f8ae9313253 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -64,8 +64,8 @@ class PretrainModelConfig(ModelConfig): moe: PretrainMoEConfig = PretrainMoEConfig() offload: Literal["none", "optimizer"] = "none" - """``optimizer`` keeps the optimizer states in CPU RAM between steps, as in SFT and RL. ``full`` is not supported: - its CPU optimizer handles only AdamW and SignSGD, without pretraining's per-role parameter groups.""" + """``optimizer`` keeps the optimizer states in CPU RAM between steps, as in SFT and RL. ``full`` (CPU optimizer + steps during backward) is not supported by the pretraining step.""" class TokenSourceConfig(BaseConfig): diff --git a/src/prime_rl/trainer/models/init.py b/src/prime_rl/trainer/models/init.py index b969e99e5f7..f2a468af576 100644 --- a/src/prime_rl/trainer/models/init.py +++ b/src/prime_rl/trainer/models/init.py @@ -1,6 +1,5 @@ import hashlib import math -from typing import Literal import torch from torch import nn @@ -12,31 +11,11 @@ from prime_rl.trainer.models.layers.moe import TokenChoiceTopKRouter from prime_rl.utils.weights import resolve_fqn -ParamRole = Literal["embedding", "output", "residual_out", "hidden", "norm", "bias"] - # Elements per random block. Blocks are seeded independently, so each rank draws only the # blocks that overlap its shard and the values do not depend on how the tensor is sharded. BLOCK_NUMEL = 1 << 22 -def param_role(name: str, param: torch.Tensor) -> ParamRole: - """Classify a parameter by its canonical name: the initialization and the pretraining optimizer groups are defined - over these roles.""" - if "embed_tokens" in name: - return "embedding" - if "lm_head" in name: - return "output" - if param.ndim == 1: - if name.endswith("bias"): - return "bias" - if "norm" in name: - return "norm" - raise ValueError(f"No init role for 1-D parameter {name!r}") - if name.endswith(("o_proj.weight", "down_proj.weight", "down_proj")): - return "residual_out" - return "hidden" - - def block_seed(seed: int, name: str, block: int) -> int: digest = hashlib.sha256(f"{seed}:{name}:{block}".encode()).digest() return int.from_bytes(digest[:8], "little") & ((1 << 63) - 1) @@ -78,13 +57,14 @@ def init_weights(model: nn.Module, config: InitConfig) -> None: for name, param in model.named_parameters(): # Canonical names, so wrappers such as activation checkpointing don't change the values. name = resolve_fqn(model, name) - role = param_role(name, param) local = param.to_local() if isinstance(param, DTensor) else param - if role == "norm": - local.fill_(1.0) - continue - if role == "bias": - local.zero_() + if param.ndim == 1: + if name.endswith("bias"): + local.zero_() + elif "norm" in name: + local.fill_(1.0) + else: + raise ValueError(f"No init for 1-D parameter {name!r}") continue if isinstance(param, DTensor): local_shape, global_offset = compute_local_shape_and_global_offset( diff --git a/src/prime_rl/trainer/pretrain/parametrization.py b/src/prime_rl/trainer/pretrain/parametrization.py deleted file mode 100644 index 66c858c98c3..00000000000 --- a/src/prime_rl/trainer/pretrain/parametrization.py +++ /dev/null @@ -1,47 +0,0 @@ -from collections import defaultdict - -from torch import nn - -from prime_rl.configs.trainer import AdamWConfig, MuonConfig -from prime_rl.trainer.models.init import param_role -from prime_rl.utils.weights import resolve_fqn - -# Roles trained by Muon (``rms_norm`` learning-rate adjustment) when ``optim.type = "muon"``; everything else, and -# every role with ``optim.type = "adamw"``, is trained by AdamW. -MUON_ROLES = ("hidden", "router", "residual_out") - - -def param_groups(model: nn.Module, optim: AdamWConfig | MuonConfig, ep_enabled: bool) -> list[dict]: - """Optimizer parameter groups per role, all with ``optim.lr`` and ``optim.eps``, and ``optim.weight_decay`` - except for norms and biases (no weight decay). With Muon, the hidden matrices (expert matrices per expert), the - residual output projections and the MoE router are Muon groups; the embedding, the LM head, norms, biases and the - Gated Residual projections (``hyper_connection``, as in Qwen3.8-Next) train with AdamW.""" - params = defaultdict(list) - for name, param in model.named_parameters(): - if param.requires_grad: - name = resolve_fqn(model, name) - role = param_role(name, param) - if role == "hidden" and ".mlp.router." in name: - role = "router" - elif role == "hidden" and "hyper_connection" in name: - role = "hyper_connection" - params[role, ".mlp.experts." in name].append(param) - - groups = [] - for (role, experts), role_params in params.items(): - weight_decay = 0.0 if role in ("norm", "bias") else optim.weight_decay - if optim.type == "muon" and role in MUON_ROLES: - group = dict( - params=role_params, algorithm="muon", adjust_lr="rms_norm", lr=optim.lr, weight_decay=weight_decay - ) - if experts and ep_enabled: - group["distributed_mesh_name"] = "dp_shard_mod_ep" - groups.append(group) - continue - group = dict(params=role_params, lr=optim.lr, weight_decay=weight_decay) - if optim.type == "muon": - group.update(algorithm="adamw", epsilon=optim.eps) - else: - group["eps"] = optim.eps - groups.append(group) - return groups diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 3d13cd7cd36..f142d41d525 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -30,7 +30,6 @@ from prime_rl.trainer.optim import setup_optimizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter -from prime_rl.trainer.pretrain.parametrization import param_groups from prime_rl.trainer.pretrain.balance import ( histogram_quantile_update_selection_biases, quantile_update_selection_biases, @@ -149,7 +148,6 @@ def train(config: PretrainConfig): parallel_dims, offload=config.model.offload, model=model, - param_groups=param_groups(model, config.optim, parallel_dims.ep_enabled), ) # With ckpt.skip_scheduler, the schedule is rebuilt over the remaining steps. diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index 5b9373395f6..6783c945cb0 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -2,10 +2,8 @@ import torch -from prime_rl.configs.trainer import MuonConfig from prime_rl.trainer.models.init import normal_shard from prime_rl.trainer.models.prime_moe import PrimeMoeConfig, PrimeMoeForCausalLM -from prime_rl.trainer.pretrain.parametrization import param_groups from prime_rl.trainer.utils import count_parameters @@ -49,30 +47,3 @@ def test_count_parameters_moe(): assert counts["embedding"] == embedding assert counts["non_embedding"] == counts["total"] - embedding assert counts["active_non_embedding"] == counts["non_embedding"] - experts * 6 // 8 - - -def test_moe_param_groups_cover_every_parameter(): - config = PrimeMoeConfig( - vocab_size=128, - hidden_size=512, - intermediate_size=1024, - moe_intermediate_size=128, - num_attention_heads=8, - num_key_value_heads=2, - n_routed_experts=8, - num_experts_per_tok=2, - n_shared_experts=1, - num_hidden_layers=2, - first_k_dense_replace=1, - ) - config._attn_implementation = "flash_attention_2" - with torch.device("meta"): - model = PrimeMoeForCausalLM._from_config(config) - groups = param_groups(model, MuonConfig(lr=1.0, weight_decay=0.1), ep_enabled=False) - group_of = {param: group for group in groups for param in group["params"]} - assert sum(len(group["params"]) for group in groups) == len(group_of) == len(list(model.parameters())) - moe = model.model.layers[1].mlp - assert group_of[moe.experts.up_proj]["algorithm"] == group_of[moe.router.gate.weight]["algorithm"] == "muon" - assert group_of[model.lm_head.weight]["algorithm"] == "adamw" - assert group_of[model.model.norm.weight]["weight_decay"] == 0.0 - assert all(group["lr"] == 1.0 for group in groups) From 99c2722a4a7d42c3508cd997b9a15234e488336d Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 14:11:14 +0200 Subject: [PATCH 114/132] pretrain: one token pool instead of weighted sources; TrainData owns the training data Mixing moves offline: data.sources / data.weights become data.paths, a list of shard folders read as one token pool (folder order, then file order), with one seeded random window order per epoch; max_epochs is checked on the pool. The per-row source draw, the per-source epoch planning and token accounting, the weight validator and the per-source progress metrics go away; progress/epochs is the pool's. TrainData (pretrain/data.py) owns the dataset, its dataloader, the resume state and the count of windows trained on: train() builds it, iterates it and checkpoints it. Its checkpoint state is the dataloader state plus the window count. --- .../ci/integration/pretrain-moe/resume.toml | 7 +- .../ci/integration/pretrain-moe/start.toml | 7 +- docs/pretraining.md | 15 +- examples/pretrain/prime_moe.toml | 17 +- .../src/prime_rl/configs/pretrain.py | 24 +- src/prime_rl/trainer/pretrain/data.py | 209 +++++++++--------- src/prime_rl/trainer/pretrain/train.py | 57 ++--- tests/integration/test_pretrain_moe.py | 2 +- tests/unit/train/pretrain/test_data.py | 49 ++-- 9 files changed, 158 insertions(+), 229 deletions(-) diff --git a/configs/ci/integration/pretrain-moe/resume.toml b/configs/ci/integration/pretrain-moe/resume.toml index 08ab3124368..952ee13a76b 100644 --- a/configs/ci/integration/pretrain-moe/resume.toml +++ b/configs/ci/integration/pretrain-moe/resume.toml @@ -1,6 +1,6 @@ # Tiny prime_moe with the base components (Gated Residual, null experts, gated attention, XSA, LatentMoE, quantile # balancing) in FP8 on 2 GPUs. The test writes the token shards (Reverse-Text-SFT, Qwen3 tokenizer) and passes their -# folder as --data.sources.reverse.path. Resumes the start run from its step-20 checkpoint. +# folder as --data.paths. Resumes the start run from its step-20 checkpoint. max_steps = 30 [model] @@ -25,10 +25,7 @@ seq_len = 2048 batch_size = 8 micro_batch_size = 2 max_epochs = 100 -weights = { reverse = 1.0 } - -[data.sources] -reverse = { path = "set-by-the-test" } +paths = ["set-by-the-test"] [optim] type = "muon" diff --git a/configs/ci/integration/pretrain-moe/start.toml b/configs/ci/integration/pretrain-moe/start.toml index 8b3704a2807..e171552c5ee 100644 --- a/configs/ci/integration/pretrain-moe/start.toml +++ b/configs/ci/integration/pretrain-moe/start.toml @@ -1,6 +1,6 @@ # Tiny prime_moe with the base components (Gated Residual, null experts, gated attention, XSA, LatentMoE, quantile # balancing) in FP8 on 2 GPUs. The test writes the token shards (Reverse-Text-SFT, Qwen3 tokenizer) and passes their -# folder as --data.sources.reverse.path. +# folder as --data.paths. max_steps = 20 [model] @@ -25,10 +25,7 @@ seq_len = 2048 batch_size = 8 micro_batch_size = 2 max_epochs = 100 -weights = { reverse = 1.0 } - -[data.sources] -reverse = { path = "set-by-the-test" } +paths = ["set-by-the-test"] [optim] type = "muon" diff --git a/docs/pretraining.md b/docs/pretraining.md index 7f5e4d4291c..42804cc098f 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -91,9 +91,9 @@ The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` ## Data -`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: per source, a folder of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of a source are read as one stream. Each epoch reads every row of a source once, in a random order over the whole source and from a random start offset, both drawn from `data.seed`, the source and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole source, whatever the order of the files. Every row is a random read (about 33 ms each from a cold cache on CephFS); at the screening shapes, `data.num_workers = 2` with the default prefetch hides this (`time/data` stays at 0). Check `time/data` on other storage. Documents continue across rows, and attention and positions restart at every document and every row; with `data.document_masking = false`, a row is one sequence (attention spans the whole row and positions run over it, as in GPT-style training without document masking). Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every source, which catches shards made with another tokenizer. +`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: folders of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of all folders in `data.paths` (in folder order, then file-name order) are read as one token pool. Each epoch reads every row of the pool once, in a random order over the whole pool and from a random start offset, both drawn from `data.seed` and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole pool, whatever the order of the files. Every row is a random read (about 33 ms each from a cold cache on CephFS); at the screening shapes, `data.num_workers = 2` with the default prefetch hides this (`time/data` stays at 0). Check `time/data` on other storage. Documents continue across rows, and attention and positions restart at every document and every row; with `data.document_masking = false`, a row is one sequence (attention spans the whole row and positions run over it, as in GPT-style training without document masking). Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every folder, which catches shards made with another tokenizer. -Each row comes from one source, drawn with the mixture `data.weights`. At startup the trainer computes how many passes over each source the run makes and fails above `data.max_epochs`. +Every folder contributes rows in proportion to its tokens, so mixtures are made offline, as folders whose token counts follow the mixture weights. At startup the trainer computes how many passes over the pool the run makes and fails above `data.max_epochs`. ```toml [data] @@ -101,19 +101,14 @@ type = "tokens" seq_len = 8192 # tokens per row batch_size = 32 # rows per optimizer step, over all GPUs micro_batch_size = 8 # rows per forward pass on one GPU -max_epochs = 3 -weights = { web = 0.8, code = 0.2 } - -[data.sources] -web = { path = "data/web" } -code = { path = "data/code" } +paths = ["data/web", "data/code"] ``` Batch sizes are in rows, as in SFT: a step trains on `batch_size * seq_len` tokens, so changing `seq_len` changes the tokens per step. Resume restores the data position on the same GPU layout and `data.num_workers`; resuming on a different layout is not supported yet. `data.type = "fake"` trains on random tokens, for smoke and throughput tests. ## Monitoring -Every step logs the loss, learning rate, gradient norms (total and per parameter group, `optim/grad_norm/`), loss and grad-norm spike flags, per-source token counts and epochs (`progress//*`), throughput, MFU, data wait and peak memory (max and mean over ranks). MoE models add the expert balance of the step's global counts (`expert_load/*`, also per layer), `router/*` and `loss/seq_balance`. The step statistics are reduced once per step. +Every step logs the loss, learning rate, gradient norms (total and per parameter group, `optim/grad_norm/`), loss and grad-norm spike flags, the rows, tokens and epochs trained on (`progress/num_samples`, `progress/num_tokens`, `progress/epochs`), throughput, MFU, data wait and peak memory (max and mean over ranks). MoE models add the expert balance of the step's global counts (`expert_load/*`, also per layer), `router/*` and `loss/seq_balance`. The step statistics are reduced once per step. ## Evals @@ -127,7 +122,7 @@ interval = 500 web = { path = "data/web-val" } ``` -**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows of `evals.micro_batch_tokens` tokens (default: the training micro-batch, `data.micro_batch_size * data.seq_len`, or the longest of `seq_lens`, whichever is larger) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` (default: `data.seq_len`) and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training source at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. Evals run without gradients, so longer rows can fit: in our tests, eval memory stayed at most 0.6x the training peak, and at d512 (8k training micro-batches), rows of 32k instead of 16k tokens cut the eval time by 19%. +**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows of `evals.micro_batch_tokens` tokens (default: the training micro-batch, `data.micro_batch_size * data.seq_len`, or the longest of `seq_lens`, whichever is larger) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` (default: `data.seq_len`) and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training folder at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. Evals run without gradients, so longer rows can fit: in our tests, eval memory stayed at most 0.6x the training peak, and at d512 (8k training micro-batches), rows of 32k instead of 16k tokens cut the eval time by 19%. ```toml [evals] diff --git a/examples/pretrain/prime_moe.toml b/examples/pretrain/prime_moe.toml index a7a3680bc26..69645498d67 100644 --- a/examples/pretrain/prime_moe.toml +++ b/examples/pretrain/prime_moe.toml @@ -19,15 +19,14 @@ type = "tokens" seq_len = 8192 batch_size = 32 micro_batch_size = 8 -max_epochs = 3 -weights = { web = 0.60, pdfs = 0.20, code = 0.12, math = 0.05, reference = 0.03 } - -[data.sources] -web = { path = "data/dolma3-50b-dsv4/web" } -pdfs = { path = "data/dolma3-50b-dsv4/pdfs" } -code = { path = "data/dolma3-50b-dsv4/code" } -math = { path = "data/dolma3-50b-dsv4/math" } -reference = { path = "data/dolma3-50b-dsv4/reference" } +# The five source folders pooled: the mixture at the token shares they were made with. +paths = [ + "data/dolma3-50b-dsv4/web", + "data/dolma3-50b-dsv4/pdfs", + "data/dolma3-50b-dsv4/code", + "data/dolma3-50b-dsv4/math", + "data/dolma3-50b-dsv4/reference", +] [optim] type = "muon" diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 04d4f6a3c26..eda1a32c3a7 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -1,4 +1,3 @@ -import math import uuid from pathlib import Path from typing import Annotated, Literal, TypeAlias @@ -82,7 +81,7 @@ class BaseBatchConfig(BaseDataConfig): """Tokens per row. Documents continue across rows; attention restarts at every document and every row.""" seed: int = 0 - """Seed for choosing each row's source, the per-epoch order of every source (and the fake tokens).""" + """Seed for the per-epoch window order of the pool (and the fake tokens).""" class FakeDataConfig(BaseBatchConfig): @@ -92,31 +91,22 @@ class FakeDataConfig(BaseBatchConfig): class TokensDataConfig(BaseBatchConfig): - """Rows of pre-tokenized datatrove shards, each row drawn from one source of a weighted mixture.""" + """Rows of pre-tokenized datatrove shards: random windows of one token pool. Mixtures are made offline, as shard + folders.""" type: Literal["tokens"] = "tokens" - sources: dict[str, TokenSourceConfig] - """Token sources by name. The name labels the source's metrics.""" - - weights: dict[str, float] - """Token share of each source. Shares are non-negative and sum to 1; unlisted sources get 0.""" + paths: list[Path] = Field(min_length=1) + """Folders of datatrove tokenized shards: ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``. + The files of all folders (in folder order, then file name order) are read as one token pool.""" max_epochs: float = Field(1.0, gt=0) - """Most passes over a source's available tokens that the run (``max_steps``) may make. Checked at startup.""" + """Most passes over the pool that the run (``max_steps``) may make. Checked at startup.""" document_masking: bool = True """Attention and positions restart at every document. If False, a row is one sequence: attention spans the whole row and positions run over it.""" - @model_validator(mode="after") - def validate_weights(self): - if any(weight < 0 for weight in self.weights.values()) or not math.isclose(sum(self.weights.values()), 1.0): - raise ValueError(f"data.weights must be non-negative and sum to 1, got {self.weights}") - if unknown := self.weights.keys() - self.sources.keys(): - raise ValueError(f"data.weights name unknown sources {sorted(unknown)}") - return self - PretrainDataConfig: TypeAlias = Annotated[FakeDataConfig | TokensDataConfig, Field(discriminator="type")] diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index 10ab2901bd1..665ce9dac5c 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -1,40 +1,44 @@ -import zlib from pathlib import Path import numpy as np import torch from torchdata.stateful_dataloader import StatefulDataLoader -from prime_rl.configs.pretrain import PretrainDataConfig, TokenSourceConfig +from prime_rl.configs.pretrain import PretrainDataConfig from prime_rl.trainer.sft.data import StatefulIterableDataset +from prime_rl.trainer.world import get_world from prime_rl.utils.logger import get_logger class TokenShards: - """Datatrove tokenized shards in one folder, the ``num_tokens`` first tokens of the files in name order, read as - one cyclic token stream. + """Datatrove tokenized shards in one or more folders, the ``num_tokens`` first tokens of their files (in folder + order, then file name order), read as one cyclic token stream. ``.ds`` holds the token ids (little-endian, 2 or 4 bytes as recorded on the first line of ``.ds.metadata``) and ``.ds.index`` the uint64 end offset of every document. File boundaries are document boundaries. - With a ``seed``, each epoch reads all ``num_tokens // seq_len`` windows of the source once, in a seeded random + With a ``seed``, each epoch reads all ``num_tokens // seq_len`` windows of the stream once, in a seeded random order and from a seeded random offset, so epochs differ in order and window boundaries, and part of an epoch is a - sample of the whole source (``window_start``). The last ``num_tokens % seq_len`` tokens of an epoch are skipped. + sample of the whole stream (``window_start``). The last ``num_tokens % seq_len`` tokens of an epoch are skipped. Without one (held-out sets), the windows are read in order.""" - def __init__(self, config: TokenSourceConfig, seed: tuple[int, str] | None = None): - self.files = sorted(config.path.glob("*.ds")) - if not self.files: - raise FileNotFoundError(f"No datatrove .ds files in {config.path}") + def __init__(self, paths: list[Path], max_tokens: int | None = None, seed: int | None = None): + self.files, self.first_files = [], [] + for path in paths: + files = sorted(path.glob("*.ds")) + if not files: + raise FileNotFoundError(f"No datatrove .ds files in {path}") + self.first_files.append(len(self.files)) + self.files += files self.dtypes = [token_dtype(file) for file in self.files] sizes = [file.stat().st_size // dtype.itemsize for file, dtype in zip(self.files, self.dtypes)] self.starts = np.cumsum([0, *sizes]) total = int(self.starts[-1]) - if config.max_tokens is not None and config.max_tokens > total: - raise ValueError(f"{config.path} has {total:,} tokens, fewer than max_tokens={config.max_tokens:,}") - self.num_tokens = total if config.max_tokens is None else config.max_tokens - self.seed = None if seed is None else [seed[0], zlib.crc32(seed[1].encode())] + if max_tokens is not None and max_tokens > total: + raise ValueError(f"{paths} have {total:,} tokens, fewer than max_tokens={max_tokens:,}") + self.num_tokens = total if max_tokens is None else max_tokens + self.seed = seed self.memmaps: dict[int, tuple[np.memmap, np.memmap]] = {} # The current epoch, its window order and its offset. self.epoch_order: tuple[int | None, np.ndarray | None, int] = (None, None, 0) @@ -80,18 +84,19 @@ def window_start(self, position: int, seq_len: int) -> int: num_windows = self.num_tokens // seq_len epoch, window = divmod(position // seq_len, num_windows) if self.epoch_order[0] != epoch: - rng = np.random.default_rng([*self.seed, epoch]) + rng = np.random.default_rng([self.seed, epoch]) order = rng.permutation(num_windows).astype(np.uint32) self.epoch_order = epoch, order, int(rng.integers(self.num_tokens)) _, order, offset = self.epoch_order return (offset + int(order[window]) * seq_len) % self.num_tokens - def first_token(self) -> int: - return int(self.open(0)[0][0]) - - def last_token_of_first_document(self) -> int: - tokens, doc_ends = self.open(0) - return int(tokens[int(doc_ends[0]) - 1]) + def first_documents(self) -> list[np.ndarray]: + """The first document of every folder.""" + documents = [] + for file_index in self.first_files: + tokens, doc_ends = self.open(file_index) + documents.append(np.asarray(tokens[: int(doc_ends[0])])) + return documents def token_dtype(file: Path) -> np.dtype: @@ -116,82 +121,46 @@ def window_start(self, position: int, seq_len: int) -> int: class TokenDataset(StatefulIterableDataset): - """``seq_len`` windows of token sources, each window from one source of a weighted mixture. + """``seq_len`` windows of a token source, window ``w`` read from ``source.window_start(w * seq_len)``. - Window ``w``'s source is drawn from the mixture ``weights`` with an RNG keyed by ``(seed, w)``, and it is that - source's next window (``TokenShards.window_start``). Every data rank walks the whole window sequence but reads only - its own windows; the state is the window counter and the tokens read per source. Each window reads one token past - its end for the last target. With ``num_windows`` the dataset stops after that many windows. Without - ``document_masking``, a window is one sequence: attention and positions run over the whole window.""" + Every data rank walks the whole window sequence but reads only its own windows; the state is the window counter. + Each window reads one token past its end for the last target. With ``num_windows`` the dataset stops after that + many windows. Without ``document_masking``, a window is one sequence: attention and positions run over the whole + window.""" def __init__( self, - sources: dict[str, TokenShards | RandomTokens], - weights: dict[str, float], + source: TokenShards | RandomTokens, seq_len: int, - seed: int = 0, non_dp_size: int = 1, num_windows: int | None = None, document_masking: bool = True, ): super().__init__(non_dp_size) - self.sources = sources - self.names = list(sources) - self.weights = np.array([weights.get(name, 0.0) for name in self.names]) - self.cdf = np.cumsum(self.weights) + self.source = source self.seq_len = seq_len - self.seed = seed self.num_windows = num_windows self.document_masking = document_masking - self.tokens = dict.fromkeys(self.names, 0) - - def epochs(self) -> dict[str, float]: - return {name: self.tokens[name] / self.sources[name].num_tokens for name in self.names} def state_dict(self) -> dict: - return {"dataset": {"step": self.step, "tokens": dict(self.tokens), "epochs": self.epochs()}} + return {"dataset": {"step": self.step}} def load_state_dict(self, state_dict: dict): self.step = state_dict["dataset"]["step"] - self.tokens.update(state_dict["dataset"]["tokens"]) - - def planned_epochs(self, num_tokens: int) -> dict[str, float]: - """Expected passes over each source's tokens when ``num_tokens`` tokens are read.""" - return { - name: num_tokens * float(weight) / self.sources[name].num_tokens - for name, weight in zip(self.names, self.weights) - } - - def source(self, window: int) -> str: - if len(self.names) == 1: - return self.names[0] - draw = np.random.default_rng([self.seed, window]).random() * self.cdf[-1] - return self.names[min(int(np.searchsorted(self.cdf, draw, side="right")), len(self.names) - 1)] def __iter__(self): self._setup_world_info() while self.num_windows is None or self.step < self.num_windows: window = self.step - name = self.source(window) - start = self.sources[name].window_start(self.tokens[name], self.seq_len) - self.tokens[name] += self.seq_len self.step += 1 if window % self.data_world_size != self.data_rank: continue - tokens, doc_starts = self.sources[name].read(start, self.seq_len + 1) + start = self.source.window_start(window * self.seq_len, self.seq_len) + tokens, doc_starts = self.source.read(start, self.seq_len + 1) doc_starts = [s for s in doc_starts if s < self.seq_len] if self.document_masking else [] yield {"tokens": tokens, "seq_lens": np.diff([0, *doc_starts, self.seq_len])} -def source_tokens(dataset: TokenDataset, start: int, end: int) -> dict[str, int]: - """Tokens each source contributes to windows ``[start, end)``. Every data rank walks the same window - sequence, so these are global counts, computed without communication.""" - tokens = dict.fromkeys(dataset.names, 0) - for window in range(start, end): - tokens[dataset.source(window)] += dataset.seq_len - return tokens - - def collate_rows(rows: list[dict]) -> dict: """Concatenate windows into one packed sample of shape ``[1, len(rows) * seq_len]``. Positions restart at every document; attention never crosses a document or a window.""" @@ -214,53 +183,85 @@ def setup_dataloader(dataset: TokenDataset, micro_batch_size: int, num_workers: ) -def setup_train_dataset( - config: PretrainDataConfig, max_steps: int, vocab_size: int, eos_token_id: int, non_dp_size: int -) -> TokenDataset: - if config.type == "fake": - sources = {"fake": RandomTokens(vocab_size, config.seed)} - weights = {"fake": 1.0} - else: - sources = {name: TokenShards(source, seed=(config.seed, name)) for name, source in config.sources.items()} - check_eos(sources, eos_token_id) - weights = config.weights - dataset = TokenDataset( - sources, - weights, - config.seq_len, - seed=config.seed, - non_dp_size=non_dp_size, - document_masking=config.type == "fake" or config.document_masking, - ) - if config.type == "tokens": - epochs = dataset.planned_epochs(max_steps * config.batch_size * config.seq_len) - get_logger().info(f"Planned epochs per source: {epochs}") - if over := {name: epoch for name, epoch in epochs.items() if epoch > config.max_epochs}: - raise ValueError(f"Planned epochs {over} exceed data.max_epochs={config.max_epochs}") - return dataset +class TrainData: + """The training data: windows of the token pool (or of fake tokens), their dataloader and the count of windows + trained on. Iterating yields this rank's micro-batches; every data rank walks the same window sequence, so the + count is global without communication. ``state_dict`` holds the dataloader position and the count, for + checkpoints.""" + + def __init__( + self, config: PretrainDataConfig, max_steps: int, vocab_size: int, eos_token_id: int, non_dp_size: int + ): + if config.type == "fake": + self.source = RandomTokens(vocab_size, config.seed) + else: + self.source = TokenShards(config.paths, seed=config.seed) + check_eos(self.source, eos_token_id) + epochs = max_steps * config.batch_size * config.seq_len / self.source.num_tokens + get_logger().info(f"Token pool: {self.source.num_tokens:,} tokens, {epochs:.3f} planned epochs") + if epochs > config.max_epochs: + raise ValueError(f"Planned epochs {epochs:.3f} exceed data.max_epochs={config.max_epochs}") + dataset = TokenDataset( + self.source, + config.seq_len, + non_dp_size=non_dp_size, + document_masking=config.type == "fake" or config.document_masking, + ) + self.loader = setup_dataloader(dataset, config.micro_batch_size, config.num_workers) + self.seq_len = config.seq_len + self.windows_per_micro_batch = config.micro_batch_size * get_world().world_size // non_dp_size + self.windows = 0 + self.iterator = None + + def __iter__(self): + return self + + def __next__(self) -> dict: + # Created at the first micro-batch, after a potential resume: iter() forks the workers with the dataset's + # current state. + if self.iterator is None: + self.iterator = iter(self.loader) + micro_batch = next(self.iterator) + self.windows += self.windows_per_micro_batch + return micro_batch + + def state_dict(self) -> dict: + return {"dataloader": self.loader.state_dict(), "windows": self.windows} + + def load_state_dict(self, state_dict: dict): + self.loader.load_state_dict(state_dict["dataloader"]) + self.windows = state_dict["windows"] + + def metrics(self) -> dict[str, float]: + return { + "progress/num_samples": self.windows, + "progress/epochs": self.windows * self.seq_len / self.source.num_tokens, + } def setup_val_dataset(source: TokenShards, seq_len: int, num_windows: int, non_dp_size: int) -> TokenDataset: """One pass over the first ``num_windows`` windows of a held-out source.""" - return TokenDataset({"val": source}, {"val": 1.0}, seq_len, non_dp_size=non_dp_size, num_windows=num_windows) + return TokenDataset(source, seq_len, non_dp_size=non_dp_size, num_windows=num_windows) -def check_eos(sources: dict[str, TokenShards], eos_token_id: int) -> None: +def check_eos(shards: TokenShards, eos_token_id: int) -> None: """Every document ends with the tokenizer's EOS. A different token there means the shards were made with - another tokenizer.""" - for name, source in sources.items(): - last = source.last_token_of_first_document() - if last != eos_token_id: + another tokenizer. Checked on the first document of every folder.""" + for file_index, document in zip(shards.first_files, shards.first_documents()): + if document[-1] != eos_token_id: raise ValueError( - f"The first document of source {name!r} ends with token {last}, not the tokenizer's EOS " - f"{eos_token_id}. Were the shards tokenized with tokenizer.name?" + f"The first document of {shards.files[file_index].parent} ends with token {document[-1]}, not the " + f"tokenizer's EOS {eos_token_id}. Were the shards tokenized with tokenizer.name?" ) -def starts_with_bos(sources: dict[str, TokenShards], bos_token_id: int | None) -> bool: - """Whether the documents of the sources start with the tokenizer's BOS, read off the first document of every - source. The eval suite formats its requests the same way.""" - starts = {name: source.first_token() == bos_token_id for name, source in sources.items()} +def starts_with_bos(shards: TokenShards, bos_token_id: int | None) -> bool: + """Whether the documents start with the tokenizer's BOS, read off the first document of every folder. The eval + suite formats its requests the same way.""" + starts = { + shards.files[i].parent.as_posix(): int(doc[0]) == bos_token_id + for i, doc in zip(shards.first_files, shards.first_documents()) + } if len(set(starts.values())) > 1: - raise ValueError(f"Only some token sources start their documents with BOS: {starts}") + raise ValueError(f"Only some token folders start their documents with BOS: {starts}") return all(starts.values()) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 35ea73e5c14..e4f9b9abdee 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -37,10 +37,9 @@ from prime_rl.trainer.pretrain.bytes import token_byte_lengths from prime_rl.trainer.pretrain.data import ( TokenShards, + TrainData, check_eos, setup_dataloader, - setup_train_dataset, - source_tokens, setup_val_dataset, starts_with_bos, ) @@ -223,15 +222,15 @@ def train(config: PretrainConfig): scheduler_steps = config.max_steps - checkpoint_step if skip_scheduler and checkpoint_step else config.max_steps scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) - train_dataset = setup_train_dataset( - config.data, config.max_steps, len(tokenizer), tokenizer.eos_token_id, non_dp_size=config.model.cp - ) - dataloader = setup_dataloader(train_dataset, config.data.micro_batch_size, config.data.num_workers) + data = TrainData(config.data, config.max_steps, len(tokenizer), tokenizer.eos_token_id, config.model.cp) val_sources = {} if config.val is not None: - val_sources = {name: TokenShards(source) for name, source in config.val.sources.items()} - check_eos(val_sources, tokenizer.eos_token_id) + val_sources = { + name: TokenShards([source.path], source.max_tokens) for name, source in config.val.sources.items() + } + for source in val_sources.values(): + check_eos(source, tokenizer.eos_token_id) byte_lengths = token_byte_lengths(tokenizer).to("cuda") # Whole micro-batches on every data-parallel rank, as every forward is a collective. windows_per_round = parallel_dims.get_mesh("dp").size() * config.data.micro_batch_size @@ -251,32 +250,14 @@ def train(config: PretrainConfig): [optimizer], None if skip_scheduler else scheduler, progress, - dataloader=dataloader if restore_dataloader else None, + dataloader=data if restore_dataloader else None, path=config.resume.dir / "trainer" if config.resume.dir is not None else None, ) progress.step += 1 if skip_scheduler: # Loading the optimizer state replaces the param groups the scheduler was built on. scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) - logger.info(f"Resuming from step {checkpoint_step} (dataset_state={get_dataset_state(dataloader)})") - - # Step s trains on the global windows [(s - 1) * batch_size, s * batch_size). The tokens per source are counted - # from the window sequence itself, which every data rank walks, so they are global without communication. - next_window = 0 - source_totals = dict.fromkeys(train_dataset.names, 0) - if restore_dataloader: - next_window = checkpoint_step * config.data.batch_size - # The loaded state counts the windows before its own position, which is near next_window. - state = next(iter(get_dataset_state(dataloader).values())) - source_totals = dict(state["tokens"]) - start, end = sorted((state["step"], next_window)) - sign = 1 if next_window > state["step"] else -1 - for name, tokens in source_tokens(train_dataset, start, end).items(): - source_totals[name] += sign * tokens - - # Create the iterator only after a potential resume: iter() forks the workers with the - # dataset's current state. - dataiter = iter(dataloader) + logger.info(f"Resuming from step {checkpoint_step} (dataset_state={get_dataset_state(data.loader)})") cp_enabled = parallel_dims.cp_enabled cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() if cp_enabled else 0 @@ -369,7 +350,7 @@ def run_val(step: int) -> None: eval_sets = [] if config.evals is not None: start_time = time.perf_counter() - bos = config.data.type == "tokens" and starts_with_bos(train_dataset.sources, tokenizer.bos_token_id) + bos = config.data.type == "tokens" and starts_with_bos(data.source, tokenizer.bos_token_id) eval_sets = [ prepare_multi_length_set(name, config.evals.data_dir, tokenizer, config.evals.seq_lens, bos) for name in config.evals.sets @@ -434,7 +415,7 @@ def run_suite(step: int) -> None: seq_balance_sums.clear() result = run_micro_batches( model, - dataiter, + data, compute_loss, grad_accum_steps, parallel_dims, @@ -499,18 +480,14 @@ def run_suite(step: int) -> None: num_tokens = result.num_tokens progress.total_tokens += num_tokens - step_tokens = source_tokens(train_dataset, next_window, next_window + config.data.batch_size) - next_window += config.data.batch_size - for name, tokens in step_tokens.items(): - source_totals[name] += tokens - progress.total_samples = next_window + progress.total_samples = data.windows # The last step's checkpoint is written before its val and evals, so a timeout there does not lose it. save_ckpt_time = 0.0 disk_metrics = {} if config.ckpt and (is_last_step or (config.ckpt.interval and progress.step % config.ckpt.interval == 0)): save_ckpt_start_time = time.perf_counter() - ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=dataloader) + ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=data) ckpt_manager.maybe_clean() save_ckpt_time = time.perf_counter() - save_ckpt_start_time disk_metrics = get_ckpt_disk_metrics(config.run_dir) @@ -543,13 +520,7 @@ def run_suite(step: int) -> None: "loss/spike": loss_spikes.update(result.loss), "optim/lr": current_lr, **{f"optim/grad_norm/{group}": norm.item() for group, norm in group_grad_norms.items()}, - **{f"progress/{name}/step_tokens": tokens for name, tokens in step_tokens.items()}, - **{f"progress/{name}/num_tokens": tokens for name, tokens in source_totals.items()}, - **{ - f"progress/{name}/epochs": tokens / train_dataset.sources[name].num_tokens - for name, tokens in source_totals.items() - }, - "progress/num_samples": progress.total_samples, + **data.metrics(), "progress/num_tokens": progress.total_tokens, "perf/throughput": throughput, "perf/throughput_per_gpu": throughput / world.world_size, diff --git a/tests/integration/test_pretrain_moe.py b/tests/integration/test_pretrain_moe.py index d94756f87fb..4d6f023d27e 100644 --- a/tests/integration/test_pretrain_moe.py +++ b/tests/integration/test_pretrain_moe.py @@ -49,7 +49,7 @@ def pretrain_command(config: str, token_dir: Path, wandb_project: str, wandb_nam "pretrain", "@", f"configs/ci/integration/pretrain-moe/{config}", - "--data.sources.reverse.path", + "--data.paths", token_dir.as_posix(), "--deployment.num-train-gpus", "2", diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index a80fca8eaf7..5e92cf1219d 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -3,8 +3,7 @@ import numpy as np -from prime_rl.configs.pretrain import TokenSourceConfig -from prime_rl.trainer.pretrain.data import TokenDataset, TokenShards, collate_rows, source_tokens +from prime_rl.trainer.pretrain.data import TokenDataset, TokenShards, collate_rows def write_shard(path: Path, documents: list[list[int]]) -> None: @@ -15,10 +14,10 @@ def write_shard(path: Path, documents: list[list[int]]) -> None: Path(f"{path}.metadata").write_text("tokenizer|4\n0\n0") -def test_token_shards_read_across_files_and_wrap(tmp_path): +def test_token_shards_pool_folders_and_wrap(tmp_path): write_shard(tmp_path / "a" / "00.ds", [[1, 2, 3], [4, 5]]) - write_shard(tmp_path / "a" / "01.ds", [[6, 7, 8, 9]]) - shards = TokenShards(TokenSourceConfig(path=tmp_path / "a")) + write_shard(tmp_path / "b" / "00.ds", [[6, 7, 8, 9]]) + shards = TokenShards([tmp_path / "a", tmp_path / "b"]) assert shards.num_tokens == 9 tokens, doc_starts = shards.read(1, 6) assert tokens.tolist() == [2, 3, 4, 5, 6, 7] @@ -31,7 +30,7 @@ def test_token_shards_read_across_files_and_wrap(tmp_path): def test_shuffled_epochs_read_every_window_once_in_a_new_order(tmp_path): write_shard(tmp_path / "a" / "00.ds", [list(range(0, 30)), list(range(30, 50))]) write_shard(tmp_path / "a" / "01.ds", [list(range(50, 72))]) - shards = TokenShards(TokenSourceConfig(path=tmp_path / "a"), seed=(0, "a")) + shards = TokenShards([tmp_path / "a"], seed=0) # 72 tokens hold 10 windows of 7; the 2-token tail of each epoch is skipped. orders = [] for epoch in range(3): @@ -41,49 +40,29 @@ def test_shuffled_epochs_read_every_window_once_in_a_new_order(tmp_path): orders.append(tuple(starts)) assert len(set(orders)) == 3 # Another reader (a dataloader worker, a resumed run) draws the same order. - again = TokenShards(TokenSourceConfig(path=tmp_path / "a"), seed=(0, "a")) + again = TokenShards([tmp_path / "a"], seed=0) assert tuple(again.window_start(position, 7) for position in range(70, 140, 7)) == orders[1] -def test_token_dataset_resumes_and_continues_each_source(tmp_path): +def test_token_dataset_resumes(tmp_path): write_shard(tmp_path / "a" / "00.ds", [list(range(100, 200))]) - write_shard(tmp_path / "b" / "00.ds", [list(range(200, 300))]) - sources = {name: TokenShards(TokenSourceConfig(path=tmp_path / name)) for name in "ab"} - - def dataset(): - return TokenDataset(sources, {"a": 0.5, "b": 0.5}, seq_len=4) - - full = [window["tokens"].tolist() for window in islice(dataset(), 8)] - for start in (100, 200): - windows = [window for window in full if start <= window[0] < start + 100] - assert windows and windows == [list(range(start + 4 * i, start + 4 * i + 5)) for i in range(len(windows))] - - first = dataset() + shards = TokenShards([tmp_path / "a"], seed=0) + full = [window["tokens"].tolist() for window in islice(TokenDataset(shards, seq_len=4), 8)] + first = TokenDataset(shards, seq_len=4) iterator = iter(first) for _ in range(3): next(iterator) - resumed = dataset() + resumed = TokenDataset(shards, seq_len=4) resumed.load_state_dict(first.state_dict()) assert [window["tokens"].tolist() for window in islice(resumed, 5)] == full[3:8] -def test_source_tokens_count_every_window_of_the_mixture(tmp_path): - for name in "ab": - write_shard(tmp_path / name / "00.ds", [list(range(100))]) - sources = {name: TokenShards(TokenSourceConfig(path=tmp_path / name)) for name in "ab"} - dataset = TokenDataset(sources, {"a": 0.3, "b": 0.7}, seq_len=4) - list(islice(dataset, 32)) - # The training loop's global per-source counts equal the tokens the dataset itself read. - assert source_tokens(dataset, 0, 32) == dataset.tokens - assert sum(source_tokens(dataset, 8, 16).values()) == 8 * 4 - - def test_without_document_masking_every_window_is_one_sequence(tmp_path): write_shard(tmp_path / "a" / "00.ds", [[1, 2, 3], [4, 5], [6, 7, 8, 9]]) - sources = {"a": TokenShards(TokenSourceConfig(path=tmp_path / "a"))} - masked = [window["seq_lens"].tolist() for window in TokenDataset(sources, {"a": 1.0}, seq_len=4, num_windows=2)] + shards = TokenShards([tmp_path / "a"]) + masked = [window["seq_lens"].tolist() for window in TokenDataset(shards, seq_len=4, num_windows=2)] assert masked == [[3, 1], [1, 3]] - dataset = TokenDataset(sources, {"a": 1.0}, seq_len=4, num_windows=2, document_masking=False) + dataset = TokenDataset(shards, seq_len=4, num_windows=2, document_masking=False) assert [window["seq_lens"].tolist() for window in dataset] == [[4], [4]] From 60ca836fcd0d88ead3fa6b5ab124d0f3abda8976 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 15:38:43 +0200 Subject: [PATCH 115/132] fix(pretrain): train the Gated Residual projections with AdamW Main's default Muon groups put the hyper_connection projections on Muon, which trains them worse (d384 with Gated Residual: +0.029 val loss over 4 seeds). Pretraining moves them to the AdamW group. --- docs/pretraining.md | 2 +- src/prime_rl/trainer/pretrain/optim.py | 18 ++++++++++++++++++ src/prime_rl/trainer/pretrain/train.py | 2 ++ tests/unit/train/pretrain/test_init.py | 23 +++++++++++++++++++++-- 4 files changed, 42 insertions(+), 3 deletions(-) create mode 100644 src/prime_rl/trainer/pretrain/optim.py diff --git a/docs/pretraining.md b/docs/pretraining.md index 7dc33e1f98e..2ed314a7ee4 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -58,7 +58,7 @@ Weights start from a random init, unless the run resumes from a checkpoint. The ## Optimizer -Pretraining uses the optimizer setup of SFT and RL. With `optim.type = "muon"` (the default), every weight matrix except the embedding and the LM head is trained by Muon (`rms_norm` learning-rate adjustment; expert matrices per expert), and the embedding, the LM head, norms and biases by AdamW with `optim.eps`. `optim.type = "adamw"` trains everything with AdamW. `model.offload = "optimizer"` keeps the optimizer states in CPU RAM between steps; `"full"` is not supported. +Pretraining uses the optimizer setup of SFT and RL. With `optim.type = "muon"` (the default), every weight matrix except the embedding and the LM head is trained by Muon (`rms_norm` learning-rate adjustment; expert matrices per expert), and the embedding, the LM head, norms, biases and the Gated Residual projections (`hyper_connection`) by AdamW with `optim.eps`. `optim.type = "adamw"` trains everything with AdamW. `model.offload = "optimizer"` keeps the optimizer states in CPU RAM between steps; `"full"` is not supported. ## Learning-Rate Schedule diff --git a/src/prime_rl/trainer/pretrain/optim.py b/src/prime_rl/trainer/pretrain/optim.py new file mode 100644 index 00000000000..9d8a686b5d9 --- /dev/null +++ b/src/prime_rl/trainer/pretrain/optim.py @@ -0,0 +1,18 @@ +from torch import nn + +from prime_rl.configs.trainer import AdamWConfig, MuonConfig +from prime_rl.trainer.optim import _default_muon_param_groups +from prime_rl.trainer.parallel_dims import ParallelDims + + +def param_groups(model: nn.Module, optim: AdamWConfig | MuonConfig, parallel_dims: ParallelDims) -> list[dict] | None: + """The default Muon groups, with the Gated Residual projections (``hyper_connection``) moved to the AdamW group, + which trains them better. None (the default groups) without Muon or Gated Residual.""" + named_params = list(model.named_parameters()) + hyper_connection = [param for name, param in named_params if "hyper_connection" in name and param.requires_grad] + if optim.type != "muon" or not hyper_connection: + return None + rest = [(name, param) for name, param in named_params if "hyper_connection" not in name] + groups = _default_muon_param_groups(optim, rest, parallel_dims, optim.lr) + next(group for group in groups if group["algorithm"] == "adamw")["params"].extend(hyper_connection) + return groups diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index f142d41d525..7b11fd06cab 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -37,6 +37,7 @@ update_selection_biases, ) from prime_rl.trainer.pretrain.bytes import token_byte_lengths +from prime_rl.trainer.pretrain.optim import param_groups from prime_rl.trainer.pretrain.data import ( TokenShards, check_eos, @@ -148,6 +149,7 @@ def train(config: PretrainConfig): parallel_dims, offload=config.model.offload, model=model, + param_groups=param_groups(model, config.optim, parallel_dims), ) # With ckpt.skip_scheduler, the schedule is rebuilt over the remaining steps. diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py index 6783c945cb0..77d90922070 100644 --- a/tests/unit/train/pretrain/test_init.py +++ b/tests/unit/train/pretrain/test_init.py @@ -1,9 +1,12 @@ import itertools +from types import SimpleNamespace import torch +from prime_rl.configs.trainer import MuonConfig from prime_rl.trainer.models.init import normal_shard from prime_rl.trainer.models.prime_moe import PrimeMoeConfig, PrimeMoeForCausalLM +from prime_rl.trainer.pretrain.optim import param_groups from prime_rl.trainer.utils import count_parameters @@ -23,7 +26,7 @@ def test_normal_shard_is_independent_of_sharding(): assert not torch.equal(normal_shard("w", shape, shape, (0, 0, 0), seed=1, device=torch.device("cpu")), full) -def test_count_parameters_moe(): +def small_moe(**kwargs) -> PrimeMoeForCausalLM: config = PrimeMoeConfig( vocab_size=128, hidden_size=512, @@ -36,10 +39,15 @@ def test_count_parameters_moe(): n_shared_experts=1, num_hidden_layers=2, first_k_dense_replace=1, + **kwargs, ) config._attn_implementation = "flash_attention_2" with torch.device("meta"): - model = PrimeMoeForCausalLM._from_config(config) + return PrimeMoeForCausalLM._from_config(config) + + +def test_count_parameters_moe(): + model = small_moe() counts = count_parameters(model) embedding = model.model.embed_tokens.weight.numel() + model.lm_head.weight.numel() experts = sum(param.numel() for param in model.model.layers[1].mlp.experts.parameters()) @@ -47,3 +55,14 @@ def test_count_parameters_moe(): assert counts["embedding"] == embedding assert counts["non_embedding"] == counts["total"] - embedding assert counts["active_non_embedding"] == counts["non_embedding"] - experts * 6 // 8 + + +def test_gated_residual_projections_train_with_adamw(): + model = small_moe(hc_count=4, hc_lowrank=32) + groups = param_groups(model, MuonConfig(lr=1.0), SimpleNamespace(ep_enabled=False)) + algorithm = {param: group["algorithm"] for group in groups for param in group["params"]} + assert sum(len(group["params"]) for group in groups) == len(algorithm) == len(list(model.parameters())) + hyper_connection = [param for name, param in model.named_parameters() if "hyper_connection" in name] + assert hyper_connection and all(algorithm[param] == "adamw" for param in hyper_connection) + assert algorithm[model.model.layers[0].self_attn.q_proj.weight] == "muon" + assert param_groups(small_moe(), MuonConfig(lr=1.0), SimpleNamespace(ep_enabled=False)) is None From e7071d342094c1b4683767ba5626f7b5c2ccf012 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 16:01:08 +0200 Subject: [PATCH 116/132] docs(pretrain): why the Gated Residual projections train with AdamW --- src/prime_rl/trainer/pretrain/optim.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/trainer/pretrain/optim.py b/src/prime_rl/trainer/pretrain/optim.py index 9d8a686b5d9..53802a87dff 100644 --- a/src/prime_rl/trainer/pretrain/optim.py +++ b/src/prime_rl/trainer/pretrain/optim.py @@ -6,8 +6,12 @@ def param_groups(model: nn.Module, optim: AdamWConfig | MuonConfig, parallel_dims: ParallelDims) -> list[dict] | None: - """The default Muon groups, with the Gated Residual projections (``hyper_connection``) moved to the AdamW group, - which trains them better. None (the default groups) without Muon or Gated Residual.""" + """The default Muon groups, with the Gated Residual projections (``hyper_connection``) moved to the AdamW group. + None (the default groups) without Muon or Gated Residual. + + These projections produce sigmoid gate logits (the write projection only one scalar per stream), not hidden + activations. Muon's orthogonalized update moves every direction by the same amount, which makes the gates noisy: + on Muon they cost +0.028 held-out loss at d384 (4 seeds).""" named_params = list(model.named_parameters()) hyper_connection = [param for name, param in named_params if "hyper_connection" in name and param.requires_grad] if optim.type != "muon" or not hyper_connection: From 1872c860bb015a39a92aadbff55639de13a62858 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 16:03:06 +0200 Subject: [PATCH 117/132] fix(pretrain): accept topk_ids in the eval scoring head Since the main merge, the LM-head wrapper passes topk_ids to the head's forward; the scoring head that replaces it for the eval suite did not take it, so every run with [evals] crashed at its first eval. --- src/prime_rl/trainer/pretrain/suite.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/prime_rl/trainer/pretrain/suite.py b/src/prime_rl/trainer/pretrain/suite.py index b173419fbeb..3a3b0abc2c9 100644 --- a/src/prime_rl/trainer/pretrain/suite.py +++ b/src/prime_rl/trainer/pretrain/suite.py @@ -64,7 +64,7 @@ def head_logprobs(hidden, weight, targets, chunk: int = 2048): return torch.cat(logprobs), torch.cat(hits) -def scoring_head(self, hidden_states, labels, temperature=None, sampling_mask=None): +def scoring_head(self, hidden_states, labels, temperature=None, sampling_mask=None, topk_ids=None): """Replaces the forward of prime-rl's LM head: log-prob and greedy hit of every label that is not IGNORE_INDEX, in fp32 without TF32, returned in the ``logprobs`` / ``entropy`` slots of its output dict.""" hidden = hidden_states.reshape(-1, hidden_states.shape[-1]) From a7997c718c646520d294ca785811c0dd3351ee70 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 16:04:01 +0200 Subject: [PATCH 118/132] pretrain: held-out token shards are eval sets; [val] goes [evals] scores both kinds of sets on one interval: evals.sets lists pretrain-evals set names (in evals.data_dir, as before) and held-out token folders ({ name, path }). A token folder is scored like a ppl set through the suite's scoring path: every document as stored, in windows of each of seq_lens, every token after its first scored; bpb counts the bytes of the scored tokens. Its metrics are eval//loss@, bpb@, ... ValConfig, TokenSourceConfig, run_val and the val windows go, and with them the seq-balance guard for no-grad forwards and the unseeded window order of TokenShards. The W&B overview and the dashboard drop their val/*/bpb panels. --- docs/pretraining.md | 18 ++--- examples/pretrain/prime_moe.toml | 12 +-- .../src/prime_rl/configs/pretrain.py | 57 ++++++------- src/prime_rl/dashboard/static/app.js | 4 +- src/prime_rl/monitors/wandb/overview.py | 2 +- src/prime_rl/trainer/pretrain/data.py | 45 ++++------- src/prime_rl/trainer/pretrain/suite.py | 42 +++++++++- src/prime_rl/trainer/pretrain/train.py | 80 ++++--------------- tests/unit/train/pretrain/test_data.py | 11 +-- tests/unit/train/pretrain/test_suite.py | 21 +++++ 10 files changed, 140 insertions(+), 152 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index d0c89941c94..536f49f7249 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -91,7 +91,7 @@ The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` ## Data -`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: folders of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of all folders in `data.paths` (in folder order, then file-name order) are read as one token pool. Each epoch reads every row of the pool once, in a random order over the whole pool and from a random start offset, both drawn from `data.seed` and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole pool, whatever the order of the files. Every row is a random read (about 33 ms each from a cold cache on CephFS); at the screening shapes, `data.num_workers = 2` with the default prefetch hides this (`time/data` stays at 0). Check `time/data` on other storage. Documents continue across rows, and attention and positions restart at every document and every row; with `data.document_masking = false`, a row is one sequence (attention spans the whole row and positions run over it, as in GPT-style training without document masking). Held-out sources are read in file-name order. Every document must end with the tokenizer's EOS; the trainer checks the first document of every folder, which catches shards made with another tokenizer. +`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: folders of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of all folders in `data.paths` (in folder order, then file-name order) are read as one token pool. Each epoch reads every row of the pool once, in a random order over the whole pool and from a random start offset, both drawn from `data.seed` and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole pool, whatever the order of the files. Every row is a random read (about 33 ms each from a cold cache on CephFS); at the screening shapes, `data.num_workers = 2` with the default prefetch hides this (`time/data` stays at 0). Check `time/data` on other storage. Documents continue across rows, and attention and positions restart at every document and every row; with `data.document_masking = false`, a row is one sequence (attention spans the whole row and positions run over it, as in GPT-style training without document masking). Every document must end with the tokenizer's EOS; the trainer checks the first document of every folder, which catches shards made with another tokenizer. Every folder contributes rows in proportion to its tokens, so mixtures are made offline, as folders whose token counts follow the mixture weights. At startup the trainer computes how many passes over the pool the run makes and fails above `data.max_epochs`. @@ -112,26 +112,22 @@ Every step logs the loss, learning rate, gradient norms (total and per parameter ## Evals -`[val]` computes the held-out loss (nats per token), perplexity and bits per UTF-8 byte (comparable across tokenizers) of every source in `val.sources`, through the same loss code as training, so it also works with context parallelism. Each eval reads the same leading windows of every source: as many as fill whole micro-batches on every data-parallel rank. Evals run after the optimizer update, so the metrics at step N describe the same weights as the step-N checkpoint. With `eval_on_start` (default), they also run before the first step (logged as step 0), and they always run at the last step. +`[evals]`: every `interval` steps, at the last step and (with `eval_on_start`, default) on the weights the run starts from, the trainer scores its eval sets with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Evals run after the optimizer update, so the metrics at step N describe the same weights as the step-N checkpoint. A set is either -```toml -[val] -interval = 500 - -[val.sources] -web = { path = "data/web-val" } -``` +- the name of a log-likelihood set of the pretrain-evals offline suite (`ppl`, `ctxcont` or `pairs`, in `//{test.parquet,meta.json}`), or +- a held-out folder of datatrove token shards (`{ name = ..., path = ... }`), scored like a `ppl` set: every document of the folder as stored, in windows of each of `seq_lens`, every token after its first scored (loss in nats per token and bits per UTF-8 byte, comparable across tokenizers). Each eval scores the whole folder, so keep held-out folders small. -**Eval suite (`[evals]`).** Every `interval` steps, at the last step and (with `eval_on_start`) on the weights the run starts from, the trainer scores the eval suite: log-likelihood sets (`ppl`, `ctxcont` and `pairs`, in `//{test.parquet,meta.json}`, the format of the pretrain-evals offline suite) with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Requests are packed into document-masked rows of `evals.micro_batch_tokens` tokens (default: the training micro-batch, `data.micro_batch_size * data.seq_len`, or the longest of `seq_lens`, whichever is larger) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` (default: `data.seq_len`) and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests start with BOS only if the training documents do (read off the first document of every training folder at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. Evals run without gradients, so longer rows can fit: in our tests, eval memory stayed at most 0.6x the training peak, and at d512 (8k training micro-batches), rows of 32k instead of 16k tokens cut the eval time by 19%. +Requests are packed into document-masked rows of `evals.micro_batch_tokens` tokens (default: the training micro-batch, `data.micro_batch_size * data.seq_len`, or the longest of `seq_lens`, whichever is larger) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` (default: `data.seq_len`) and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests of the pretrain-evals sets start with BOS only if the training documents do (read off the first document of every training folder at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. Evals run without gradients, so longer rows can fit: in our tests, eval memory stayed at most 0.6x the training peak, and at d512 (8k training micro-batches), rows of 32k instead of 16k tokens cut the eval time by 19%. ```toml [evals] interval = 1000 data_dir = "/shared/evals-data" -sets = ["fwedu-holdout", "lambada", "math"] +sets = ["fwedu-holdout", "lambada", "math", { name = "web-val", path = "data/web-val" }] seq_lens = [2048, 4096, 16384] ``` + ## Multi-Node `uv run pretrain` takes the same `[deployment]` and `[slurm]` blocks as `uv run sft` (see [Scaling § SLURM](scaling.md#slurm)) and submits the same bundled templates. The launcher writes the model dir on the submitting node, then submits one job in which every node runs `torchrun` on the pretrain trainer: diff --git a/examples/pretrain/prime_moe.toml b/examples/pretrain/prime_moe.toml index 69645498d67..08c763f21fb 100644 --- a/examples/pretrain/prime_moe.toml +++ b/examples/pretrain/prime_moe.toml @@ -33,13 +33,13 @@ type = "muon" lr = 0.0078125 max_norm = 1.0 -[val] +[evals] interval = 512 - -[val.sources] -web = { path = "data/dolma3-50b-dsv4/web-val" } -code = { path = "data/dolma3-50b-dsv4/code-val" } -math = { path = "data/dolma3-50b-dsv4/math-val" } +sets = [ + { name = "web-val", path = "data/dolma3-50b-dsv4/web-val" }, + { name = "code-val", path = "data/dolma3-50b-dsv4/code-val" }, + { name = "math-val", path = "data/dolma3-50b-dsv4/math-val" }, +] [ckpt] interval = 1024 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py index 25f86c82904..2106d89041b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py @@ -68,14 +68,6 @@ class PretrainModelConfig(ModelConfig): steps during backward) is not supported by the pretraining step.""" -class TokenSourceConfig(BaseConfig): - path: Path - """Folder of datatrove tokenized shards: ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``. The files are read in name order as one token stream.""" - - max_tokens: int | None = Field(None, ge=1) - """Tokens available from this source: the first ``max_tokens`` of its stream. Defaults to all tokens in the shards.""" - - class BaseBatchConfig(BaseDataConfig): seq_len: int = Field(8192, ge=1) """Tokens per row. Documents continue across rows; attention restarts at every document and every row.""" @@ -111,41 +103,39 @@ class TokensDataConfig(BaseBatchConfig): PretrainDataConfig: TypeAlias = Annotated[FakeDataConfig | TokensDataConfig, Field(discriminator="type")] -class ValConfig(BaseConfig): - """Held-out loss, perplexity and bits per byte, computed like the training loss on held-out token shards.""" - - interval: int = Field(500, ge=1) - """Evaluate every N steps, after that step's optimizer update.""" +class TokenEvalSetConfig(BaseConfig): + """A held-out folder of datatrove tokenized shards, scored like a ``ppl`` set: every document on its own (in + windows of each of ``seq_lens``), all documents of the folder.""" - eval_on_start: bool = True - """Evaluate the initial weights before the first step (logged as step 0).""" + name: str + """Names the metrics, ``eval//@``.""" - sources: dict[str, TokenSourceConfig] - """Held-out sources by name, logged as ``val//loss``, ``/ppl`` and ``/bpb``. Each eval reads the same - windows: as many as fill whole micro-batches on every data-parallel rank.""" + path: Path + """Folder of ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``, as for training.""" class EvalsConfig(BaseConfig): - """The eval suite, scored inside the trainer: forward passes without gradients on the training model, every - ``interval`` steps. The sets are log-likelihood evals in the format of the pretrain-evals offline suite.""" + """The evals, scored inside the trainer: forward passes without gradients on the training model, every + ``interval`` steps. The sets are log-likelihood evals in the format of the pretrain-evals offline suite, and + held-out token shards.""" interval: int = Field(ge=1) - """Score the suite every N steps (after that step's optimizer update), and at the last step.""" + """Score the sets every N steps (after that step's optimizer update), and at the last step.""" eval_on_start: bool = True """Also score the weights the run starts from (the initial weights, logged as step 0, or the checkpoint it resumes from, logged at that step).""" - data_dir: Path - """Eval sets, one directory per set with ``test.parquet`` and ``meta.json``.""" + data_dir: Path | None = None + """Directory of the pretrain-evals sets, one directory per set with ``test.parquet`` and ``meta.json``.""" - sets: list[str] - """Names of the sets to score.""" + sets: list[str | TokenEvalSetConfig] = Field(min_length=1) + """The sets to score: the name of a pretrain-evals set in ``data_dir``, or a held-out token folder.""" seq_lens: list[int] | None = None - """Context lengths to score every set at: the window length for ``ppl`` sets and the context limit for - ``ctxcont`` sets. Defaults to ``data.seq_len``. Metrics are named ``eval//@`` (e.g. ``bpb@4k``). Requests that are the - same at several lengths are scored once.""" + """Context lengths to score every set at: the window length for ``ppl`` and token sets and the context limit for + ``ctxcont`` sets. Defaults to ``data.seq_len``. Metrics are named ``eval//@`` (e.g. + ``bpb@4k``). Requests that are the same at several lengths are scored once.""" micro_batch_tokens: int | None = Field(None, ge=1) """Tokens per packed eval row (one forward pass on one GPU). Defaults to the training micro-batch @@ -155,6 +145,12 @@ class EvalsConfig(BaseConfig): save_arrays: bool = False """Write per-document arrays (npz) and summaries to ``/evals/step_N/``.""" + @model_validator(mode="after") + def validate_data_dir(self): + if self.data_dir is None and any(isinstance(eval_set, str) for eval_set in self.sets): + raise ValueError("evals.data_dir is required for pretrain-evals sets") + return self + PretrainOptimizerConfig: TypeAlias = Annotated[AdamWConfig | MuonConfig, Field(discriminator="type")] @@ -170,11 +166,8 @@ class PretrainConfig(TrainerLaunchConfig): data: PretrainDataConfig - val: ValConfig | None = None - """Held-out loss on token shards (``val/*``). If None, it never runs.""" - evals: EvalsConfig | None = None - """The eval suite, scored inside the trainer. If None, it never runs.""" + """The evals, scored inside the trainer. If None, they never run.""" optim: PretrainOptimizerConfig = MuonConfig() """Optimizer. ``optim.lr`` is required.""" diff --git a/src/prime_rl/dashboard/static/app.js b/src/prime_rl/dashboard/static/app.js index 38b42970e8a..228bfa36a7d 100644 --- a/src/prime_rl/dashboard/static/app.js +++ b/src/prime_rl/dashboard/static/app.js @@ -446,8 +446,8 @@ const SFT_TRAIN_METRICS = ["loss/mean", "loss/perplexity", "val/loss", "val/perp // max_vio only exists on MoE models: a regex panel stays away on dense ones const SFT_STABILITY_PANELS = [{ metric: "optim/grad_norm" }, { metric: "optim/lr" }, { metric: "loss/nan_count" }, { regex: "max_vio/(mean|max)" }]; const SFT_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memory", "time/forward_backward"]; -// Pretrain flavor (overview.py PRETRAIN_*): held-out bpb from [val] (val/) and the eval suite (eval/) -const PRETRAIN_TRAIN_PANELS = [{ metric: "loss/mean" }, { metric: "loss/spike" }, { regex: "val/.*/bpb" }, { regex: "eval/.*/bpb" }]; +// Pretrain flavor (overview.py PRETRAIN_*): bpb of the evals (eval/) +const PRETRAIN_TRAIN_PANELS = [{ metric: "loss/mean" }, { metric: "loss/spike" }, { regex: "eval/.*/bpb" }]; const PRETRAIN_STABILITY_PANELS = [{ metric: "optim/grad_norm" }, { metric: "optim/grad_norm_spike" }, { metric: "optim/lr" }, { metric: "loss/nan_count" }, { regex: "optim/grad_norm/.*" }]; const PRETRAIN_MOE_METRICS = ["expert_load/max_mean/mean", "expert_load/max_mean/max", "expert_load/micro_max_mean/mean", "expert_load/micro_max_mean/max", "expert_load/entropy/min", "expert_load/dead", "router/confidence", "router/bias/absmax", "loss/seq_balance"]; const PRETRAIN_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memory/max", "time/step", "time/data/max", "time/eval_suite"]; diff --git a/src/prime_rl/monitors/wandb/overview.py b/src/prime_rl/monitors/wandb/overview.py index 6a763f3d5d8..8c6c52a9210 100644 --- a/src/prime_rl/monitors/wandb/overview.py +++ b/src/prime_rl/monitors/wandb/overview.py @@ -93,7 +93,7 @@ # Pretrain flavor: loss and held-out evals, training stability, MoE balance, performance. PRETRAIN_TRAIN_METRICS = ["loss/mean", "loss/spike"] -PRETRAIN_TRAIN_REGEXES = ["val/.*/bpb", "eval/.*/bpb"] +PRETRAIN_TRAIN_REGEXES = ["eval/.*/bpb"] PRETRAIN_STABILITY_METRICS = ["optim/grad_norm", "optim/grad_norm_spike", "optim/lr", "loss/nan_count"] PRETRAIN_STABILITY_REGEXES = ["optim/grad_norm/.*"] PRETRAIN_MOE_METRICS = [ diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py index 665ce9dac5c..d94751f94cb 100644 --- a/src/prime_rl/trainer/pretrain/data.py +++ b/src/prime_rl/trainer/pretrain/data.py @@ -18,12 +18,12 @@ class TokenShards: ``.ds.metadata``) and ``.ds.index`` the uint64 end offset of every document. File boundaries are document boundaries. - With a ``seed``, each epoch reads all ``num_tokens // seq_len`` windows of the stream once, in a seeded random - order and from a seeded random offset, so epochs differ in order and window boundaries, and part of an epoch is a - sample of the whole stream (``window_start``). The last ``num_tokens % seq_len`` tokens of an epoch are skipped. - Without one (held-out sets), the windows are read in order.""" + Each epoch reads all ``num_tokens // seq_len`` windows of the stream once, in a random order and from a random + offset drawn from ``seed`` and the epoch, so epochs differ in order and window boundaries, and part of an epoch is + a sample of the whole stream (``window_start``). The last ``num_tokens % seq_len`` tokens of an epoch are + skipped.""" - def __init__(self, paths: list[Path], max_tokens: int | None = None, seed: int | None = None): + def __init__(self, paths: list[Path], seed: int = 0): self.files, self.first_files = [], [] for path in paths: files = sorted(path.glob("*.ds")) @@ -34,10 +34,7 @@ def __init__(self, paths: list[Path], max_tokens: int | None = None, seed: int | self.dtypes = [token_dtype(file) for file in self.files] sizes = [file.stat().st_size // dtype.itemsize for file, dtype in zip(self.files, self.dtypes)] self.starts = np.cumsum([0, *sizes]) - total = int(self.starts[-1]) - if max_tokens is not None and max_tokens > total: - raise ValueError(f"{paths} have {total:,} tokens, fewer than max_tokens={max_tokens:,}") - self.num_tokens = total if max_tokens is None else max_tokens + self.num_tokens = int(self.starts[-1]) self.seed = seed self.memmaps: dict[int, tuple[np.memmap, np.memmap]] = {} # The current epoch, its window order and its offset. @@ -79,8 +76,6 @@ def read(self, start: int, length: int) -> tuple[np.ndarray, list[int]]: def window_start(self, position: int, seq_len: int) -> int: """Stream offset of the window to read when ``position`` tokens have been read.""" - if self.seed is None: - return position % self.num_tokens num_windows = self.num_tokens // seq_len epoch, window = divmod(position // seq_len, num_windows) if self.epoch_order[0] != epoch: @@ -124,22 +119,19 @@ class TokenDataset(StatefulIterableDataset): """``seq_len`` windows of a token source, window ``w`` read from ``source.window_start(w * seq_len)``. Every data rank walks the whole window sequence but reads only its own windows; the state is the window counter. - Each window reads one token past its end for the last target. With ``num_windows`` the dataset stops after that - many windows. Without ``document_masking``, a window is one sequence: attention and positions run over the whole - window.""" + Each window reads one token past its end for the last target. Without ``document_masking``, a window is one + sequence: attention and positions run over the whole window.""" def __init__( self, source: TokenShards | RandomTokens, seq_len: int, non_dp_size: int = 1, - num_windows: int | None = None, document_masking: bool = True, ): super().__init__(non_dp_size) self.source = source self.seq_len = seq_len - self.num_windows = num_windows self.document_masking = document_masking def state_dict(self) -> dict: @@ -150,7 +142,7 @@ def load_state_dict(self, state_dict: dict): def __iter__(self): self._setup_world_info() - while self.num_windows is None or self.step < self.num_windows: + while True: window = self.step self.step += 1 if window % self.data_world_size != self.data_rank: @@ -177,12 +169,6 @@ def collate_rows(rows: list[dict]) -> dict: } -def setup_dataloader(dataset: TokenDataset, micro_batch_size: int, num_workers: int) -> StatefulDataLoader: - return StatefulDataLoader( - dataset, batch_size=micro_batch_size, collate_fn=collate_rows, num_workers=num_workers, pin_memory=True - ) - - class TrainData: """The training data: windows of the token pool (or of fake tokens), their dataloader and the count of windows trained on. Iterating yields this rank's micro-batches; every data rank walks the same window sequence, so the @@ -207,7 +193,13 @@ def __init__( non_dp_size=non_dp_size, document_masking=config.type == "fake" or config.document_masking, ) - self.loader = setup_dataloader(dataset, config.micro_batch_size, config.num_workers) + self.loader = StatefulDataLoader( + dataset, + batch_size=config.micro_batch_size, + collate_fn=collate_rows, + num_workers=config.num_workers, + pin_memory=True, + ) self.seq_len = config.seq_len self.windows_per_micro_batch = config.micro_batch_size * get_world().world_size // non_dp_size self.windows = 0 @@ -239,11 +231,6 @@ def metrics(self) -> dict[str, float]: } -def setup_val_dataset(source: TokenShards, seq_len: int, num_windows: int, non_dp_size: int) -> TokenDataset: - """One pass over the first ``num_windows`` windows of a held-out source.""" - return TokenDataset(source, seq_len, non_dp_size=non_dp_size, num_windows=num_windows) - - def check_eos(shards: TokenShards, eos_token_id: int) -> None: """Every document ends with the tokenizer's EOS. A different token there means the shards were made with another tokenizer. Checked on the first document of every folder.""" diff --git a/src/prime_rl/trainer/pretrain/suite.py b/src/prime_rl/trainer/pretrain/suite.py index b173419fbeb..860c41e4892 100644 --- a/src/prime_rl/trainer/pretrain/suite.py +++ b/src/prime_rl/trainer/pretrain/suite.py @@ -12,6 +12,9 @@ continuation span is found by character offsets. - ``pairs``: ``good``/``bad`` sentences, each scored as [BOS +] sentence from its second token. +Held-out token shards (``TokenEvalSetConfig``) are scored as a ``ppl`` set (``prepare_token_set``): every document +as stored, cut into windows the same way; bpb counts the bytes of the scored tokens. + A set is tokenized once (`prepare_set`); every round scores its requests (`score_requests`, the model's LM head replaced by `scoring_head`) and summarizes them (`summarize`).""" @@ -31,7 +34,9 @@ from torch import nn from transformers.tokenization_utils import PreTrainedTokenizer +from prime_rl.configs.pretrain import TokenEvalSetConfig from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX +from prime_rl.trainer.pretrain.data import TokenShards, check_eos from prime_rl.trainer.world import get_world POSITION_BINS = (0, 128, 512, 2048, 4096, 16384) @@ -48,7 +53,7 @@ class EvalSet: table: Any requests: list[Request] context: dict - sha256: str + sha256: str | None def head_logprobs(hidden, weight, targets, chunk: int = 2048): @@ -228,6 +233,24 @@ def prepare_set(name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_l return EvalSet(name, kind, table, requests, context, meta["parquet_sha256"]) +def prepare_token_set(name: str, shards: TokenShards, seq_len: int, byte_lengths: np.ndarray) -> EvalSet: + """Every document of held-out token shards as a ``ppl`` set: cut into non-overlapping windows of ``seq_len`` + scored tokens (later windows start cold), every token after the first scored. ``byte_lengths``: the UTF-8 bytes + of every token id, for bpb.""" + requests, owner, num_bytes = [], [], [] + for file_index in range(len(shards.files)): + tokens, doc_ends = shards.open(file_index) + begin = 0 + for end in doc_ends.tolist(): + ids = tokens[begin:end] + for window in range(0, len(ids) - 1, seq_len): + requests.append((ids[window : window + seq_len + 1].tolist(), 1)) + owner.append(len(num_bytes)) + num_bytes.append(int(byte_lengths[ids[1:]].sum())) + begin = end + return EvalSet(name, "ppl", None, requests, {"owner": owner, "bytes": np.array(num_bytes)}, None) + + @dataclass class MultiLengthSet: """One eval set at several context lengths (``ppl`` window / ``ctxcont`` context limit). A request that is the @@ -242,9 +265,22 @@ class MultiLengthSet: def prepare_multi_length_set( - name: str, data_dir: Path, tokenizer: PreTrainedTokenizer, seq_lens: list[int], bos: bool + eval_set: str | TokenEvalSetConfig, + data_dir: Path | None, + tokenizer: PreTrainedTokenizer, + seq_lens: list[int], + bos: bool, + byte_lengths: np.ndarray | None = None, ) -> MultiLengthSet: - versions = {seq_len: prepare_set(name, data_dir, tokenizer, seq_len, bos) for seq_len in seq_lens} + """A pretrain-evals set (by name, in ``data_dir``) or held-out token shards at every length of ``seq_lens``.""" + if isinstance(eval_set, str): + name = eval_set + versions = {seq_len: prepare_set(name, data_dir, tokenizer, seq_len, bos) for seq_len in seq_lens} + else: + name = eval_set.name + shards = TokenShards([eval_set.path]) + check_eos(shards, tokenizer.eos_token_id) + versions = {seq_len: prepare_token_set(name, shards, seq_len, byte_lengths) for seq_len in seq_lens} distinct: dict[tuple[tuple[int, ...], int], int] = {} index = { seq_len: [distinct.setdefault((tuple(ids), start), len(distinct)) for ids, start in version.requests] diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 77f38ab58f1..0775d197cd6 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -38,11 +38,7 @@ ) from prime_rl.trainer.pretrain.bytes import token_byte_lengths from prime_rl.trainer.pretrain.data import ( - TokenShards, TrainData, - check_eos, - setup_dataloader, - setup_val_dataset, starts_with_bos, ) from prime_rl.trainer.pretrain.suite import ( @@ -156,23 +152,6 @@ def train(config: PretrainConfig): data = TrainData(config.data, config.max_steps, len(tokenizer), tokenizer.eos_token_id, config.model.cp) - val_sources = {} - if config.val is not None: - val_sources = { - name: TokenShards([source.path], source.max_tokens) for name, source in config.val.sources.items() - } - for source in val_sources.values(): - check_eos(source, tokenizer.eos_token_id) - byte_lengths = token_byte_lengths(tokenizer).to("cuda") - # Whole micro-batches on every data-parallel rank, as every forward is a collective. - windows_per_round = parallel_dims.get_mesh("dp").size() * config.data.micro_batch_size - val_windows = {} - for name, source in val_sources.items(): - val_windows[name] = (source.num_tokens - 1) // config.data.seq_len // windows_per_round * windows_per_round - if val_windows[name] == 0: - raise ValueError(f"Val source {name!r} has fewer than {windows_per_round} windows of data.seq_len") - logger.info(f"Val source {name!r}: {val_windows[name] * config.data.seq_len:,} tokens") - progress = Progress() restore_dataloader = checkpoint_step is not None and not (config.ckpt or CheckpointConfig()).skip_dataloader if checkpoint_step is not None: @@ -204,7 +183,7 @@ def train(config: PretrainConfig): seq_balance_sums: dict[str, torch.Tensor] = {} def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Forward pass returning (loss_sum, token_count) over unmasked tokens and, in training, the MoE + """Forward pass returning (loss_sum, token_count) over unmasked tokens and the MoE sequence-wise balance loss summed over tokens (each token carries its row's loss, so after the division by the token count it is the mean over rows).""" input_ids = micro_batch["input_ids"].to("cuda", non_blocking=True) @@ -222,14 +201,13 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.T model, input_ids, position_ids, seq_lens=seq_lens, labels=labels, seq_lens_are_pre_shard=cp_enabled ) aux_loss_sum = torch.zeros((), device="cuda") - if torch.is_grad_enabled(): - row_len = input_ids.shape[1] // config.data.micro_batch_size - for moe in moes_with_seq_balance: - aux_loss_sum += balance.seq_aux_coeff * moe.seq_balance.sum() * row_len - seq_balance_sums["balance"] = seq_balance_sums.get("balance", 0) + moe.seq_balance.detach().sum() - seq_balance_sums["rows"] = seq_balance_sums.get("rows", 0) + moe.seq_balance.new_tensor( - moe.seq_balance.numel() - ) + row_len = input_ids.shape[1] // config.data.micro_batch_size + for moe in moes_with_seq_balance: + aux_loss_sum += balance.seq_aux_coeff * moe.seq_balance.sum() * row_len + seq_balance_sums["balance"] = seq_balance_sums.get("balance", 0) + moe.seq_balance.detach().sum() + seq_balance_sums["rows"] = seq_balance_sums.get("rows", 0) + moe.seq_balance.new_tensor( + moe.seq_balance.numel() + ) return out["loss"], loss_mask.sum(dtype=torch.int64), aux_loss_sum def finish_no_grad_forwards() -> None: @@ -239,35 +217,17 @@ def finish_no_grad_forwards() -> None: if is_moe_model: take_routing_counts(model) - @torch.no_grad() - def run_val(step: int) -> None: - """Held-out loss, perplexity and bits per byte on the weights after ``step`` optimizer updates.""" - metrics = {} - for name, source in val_sources.items(): - val_dataset = setup_val_dataset(source, config.data.seq_len, val_windows[name], config.model.cp) - totals = torch.zeros(3, dtype=torch.float64, device="cuda") # loss sum, tokens, bytes - for micro_batch in setup_dataloader(val_dataset, config.data.micro_batch_size, num_workers=1): - loss_sum, token_count, _ = compute_loss(micro_batch) - # The bytes are counted on the unsharded targets, so every CP rank counts them once. - num_bytes = byte_lengths[micro_batch["target_ids"].to("cuda")].sum() / cp_size - totals += torch.stack([loss_sum.double(), token_count.double(), num_bytes.double()]) - dist.all_reduce(totals, op=dist.ReduceOp.SUM, group=dp_cp_group) - loss_sum, num_tokens, num_bytes = totals.tolist() - loss = loss_sum / num_tokens - metrics[f"val/{name}/loss"] = loss - metrics[f"val/{name}/ppl"] = math.exp(min(loss, 20)) - metrics[f"val/{name}/bpb"] = loss_sum / (math.log(2) * num_bytes) - finish_no_grad_forwards() - logger.success(f"Val | Step {step} | " + " | ".join(f"{k} {v:.4f}" for k, v in sorted(metrics.items()))) - asyncio.run(monitors.log({**metrics, "step": step}, step=step)) - eval_sets = [] if config.evals is not None: start_time = time.perf_counter() bos = config.data.type == "tokens" and starts_with_bos(data.source, tokenizer.bos_token_id) + has_token_sets = any(not isinstance(eval_set, str) for eval_set in config.evals.sets) + byte_lengths = token_byte_lengths(tokenizer).numpy() if has_token_sets else None eval_sets = [ - prepare_multi_length_set(name, config.evals.data_dir, tokenizer, config.evals.seq_lens, bos) - for name in config.evals.sets + prepare_multi_length_set( + eval_set, config.evals.data_dir, tokenizer, config.evals.seq_lens, bos, byte_lengths + ) + for eval_set in config.evals.sets ] logger.info(f"Prepared {len(eval_sets)} eval sets (BOS: {bos}) in {time.perf_counter() - start_time:.0f}s") @@ -296,15 +256,11 @@ def run_suite(step: int) -> None: logger.success(f"Eval suite | Step {step} | {metrics['time/eval_suite']:.0f}s | {len(eval_sets)} sets") asyncio.run(monitors.log({**metrics, "step": step}, step=step)) - if checkpoint_step is None and config.val is not None and config.val.eval_on_start: - run_val(step=0) if config.evals is not None and config.evals.eval_on_start: run_suite(step=checkpoint_step or 0) if checkpoint_step == config.max_steps: - # The last step's checkpoint is written before its val and evals: resuming from it only runs those. - logger.info(f"Resuming from the last step ({checkpoint_step}): running its val and evals only") - if config.val is not None: - run_val(checkpoint_step) + # The last step's checkpoint is written before its evals: resuming from it only runs those. + logger.info(f"Resuming from the last step ({checkpoint_step}): running its evals only") if config.evals is not None and not config.evals.eval_on_start: run_suite(checkpoint_step) logger.success("Pretrain trainer finished") @@ -396,7 +352,7 @@ def run_suite(step: int) -> None: progress.total_tokens += num_tokens progress.total_samples = data.windows - # The last step's checkpoint is written before its val and evals, so a timeout there does not lose it. + # The last step's checkpoint is written before its evals, so a timeout there does not lose it. save_ckpt_time = 0.0 disk_metrics = {} if config.ckpt and (is_last_step or (config.ckpt.interval and progress.step % config.ckpt.interval == 0)): @@ -407,8 +363,6 @@ def run_suite(step: int) -> None: disk_metrics = get_ckpt_disk_metrics(config.run_dir) eval_start_time = time.perf_counter() - if config.val is not None and (is_last_step or progress.step % config.val.interval == 0): - run_val(progress.step) if config.evals is not None and (is_last_step or progress.step % config.evals.interval == 0): run_suite(progress.step) eval_time = time.perf_counter() - eval_start_time diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py index 5e92cf1219d..20aeb21586c 100644 --- a/tests/unit/train/pretrain/test_data.py +++ b/tests/unit/train/pretrain/test_data.py @@ -58,12 +58,13 @@ def test_token_dataset_resumes(tmp_path): def test_without_document_masking_every_window_is_one_sequence(tmp_path): - write_shard(tmp_path / "a" / "00.ds", [[1, 2, 3], [4, 5], [6, 7, 8, 9]]) + write_shard(tmp_path / "a" / "00.ds", [[1, 2], [3, 4], [5, 6], [7, 8]]) shards = TokenShards([tmp_path / "a"]) - masked = [window["seq_lens"].tolist() for window in TokenDataset(shards, seq_len=4, num_windows=2)] - assert masked == [[3, 1], [1, 3]] - dataset = TokenDataset(shards, seq_len=4, num_windows=2, document_masking=False) - assert [window["seq_lens"].tolist() for window in dataset] == [[4], [4]] + # Every window of 4 holds a document start. + masked = [window["seq_lens"] for window in islice(TokenDataset(shards, seq_len=4), 8)] + assert all(len(seq_lens) > 1 for seq_lens in masked) + unmasked = islice(TokenDataset(shards, seq_len=4, document_masking=False), 8) + assert [window["seq_lens"].tolist() for window in unmasked] == [[4]] * 8 def test_collate_rows_restarts_positions_at_documents(): diff --git a/tests/unit/train/pretrain/test_suite.py b/tests/unit/train/pretrain/test_suite.py index dd2ad8cf61e..b82ff94d5d4 100644 --- a/tests/unit/train/pretrain/test_suite.py +++ b/tests/unit/train/pretrain/test_suite.py @@ -10,11 +10,13 @@ from torch import nn from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX +from prime_rl.trainer.pretrain.data import TokenShards from prime_rl.trainer.pretrain.suite import ( head_logprobs, pack_requests, prepare_multi_length_set, prepare_set, + prepare_token_set, score_requests, score_sets, set_metrics, @@ -72,6 +74,25 @@ def test_multi_length_set_scores_shared_requests_once(tmp_path): assert multi.index[8][0] == multi.index[32][0] +def test_token_set_scores_documents_like_a_ppl_set_without_bos(tmp_path): + texts = ["short", "a much longer document"] + (tmp_path / "docs").mkdir() + pq.write_table(pa.table({"text": texts}), tmp_path / "docs" / "test.parquet") + (tmp_path / "docs" / "meta.json").write_text(json.dumps({"kind": "ppl", "parquet_sha256": "x"})) + documents = [[ord(char) for char in text] + [CharTokenizer.eos_token_id] for text in texts] + (tmp_path / "shards").mkdir() + np.array([token for document in documents for token in document], dtype=" Date: Wed, 21 Oct 2026 18:44:07 +0200 Subject: [PATCH 119/132] docs(pretrain): the 5M-token held-out set in the examples, its cost, and why old val numbers don't compare examples/pretrain/prime_moe.toml scores data/dolma3-50b-dsv4/mix-val-5M instead of three whole source val folders. docs/pretraining.md shows it, gives the measured cost at d512 (about 1 s per 200K tokens per length on one H200; 5M tokens at [2048, 8192] in 40 s against 280 s for the suite) and says why token-set losses sit 0.014-0.027 nats below the former [val]. --- docs/pretraining.md | 4 ++-- examples/pretrain/prime_moe.toml | 7 ++----- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/docs/pretraining.md b/docs/pretraining.md index d89dc34bc71..37aad5a4f25 100644 --- a/docs/pretraining.md +++ b/docs/pretraining.md @@ -115,7 +115,7 @@ Every step logs the loss, learning rate, gradient norms (total and per parameter `[evals]`: every `interval` steps, at the last step and (with `eval_on_start`, default) on the weights the run starts from, the trainer scores its eval sets with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Evals run after the optimizer update, so the metrics at step N describe the same weights as the step-N checkpoint. A set is either - the name of a log-likelihood set of the pretrain-evals offline suite (`ppl`, `ctxcont` or `pairs`, in `//{test.parquet,meta.json}`), or -- a held-out folder of datatrove token shards (`{ name = ..., path = ... }`), scored like a `ppl` set: every document of the folder as stored, in windows of each of `seq_lens`, every token after its first scored (loss in nats per token and bits per UTF-8 byte, comparable across tokenizers). Each eval scores the whole folder, so keep held-out folders small. +- a held-out folder of datatrove token shards (`{ name = ..., path = ... }`), scored like a `ppl` set: every document of the folder as stored, in windows of each of `seq_lens`, every token after its first scored (loss in nats per token and bits per UTF-8 byte, comparable across tokenizers). Each eval scores the whole folder at every length, so keep held-out folders small: at d512 on one H200, a token set costs about 1 s per 200K tokens per length (5M tokens at `seq_lens = [2048, 8192]`: 40 s, against 280 s for the 58-set suite). Because every document is scored from its start, these losses are lower than those of the former `[val]` (by 0.014-0.027 nats at d512), which started rows mid-document and predicted each document's first token across the document boundary: don't compare old `val/*` numbers with `eval/*` token-set numbers. Requests are packed into document-masked rows of `evals.micro_batch_tokens` tokens (default: the training micro-batch, `data.micro_batch_size * data.seq_len`, or the longest of `seq_lens`, whichever is larger) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` (default: `data.seq_len`) and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests of the pretrain-evals sets start with BOS only if the training documents do (read off the first document of every training folder at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. Evals run without gradients, so longer rows can fit: in our tests, eval memory stayed at most 0.6x the training peak, and at d512 (8k training micro-batches), rows of 32k instead of 16k tokens cut the eval time by 19%. @@ -123,7 +123,7 @@ Requests are packed into document-masked rows of `evals.micro_batch_tokens` toke [evals] interval = 1000 data_dir = "/shared/evals-data" -sets = ["fwedu-holdout", "lambada", "math", { name = "web-val", path = "data/web-val" }] +sets = ["fwedu-holdout", "lambada", "math", { name = "heldout", path = "data/dolma3-50b-dsv4/mix-val-5M" }] seq_lens = [2048, 4096, 16384] ``` diff --git a/examples/pretrain/prime_moe.toml b/examples/pretrain/prime_moe.toml index 08c763f21fb..45cf067cb98 100644 --- a/examples/pretrain/prime_moe.toml +++ b/examples/pretrain/prime_moe.toml @@ -35,11 +35,8 @@ max_norm = 1.0 [evals] interval = 512 -sets = [ - { name = "web-val", path = "data/dolma3-50b-dsv4/web-val" }, - { name = "code-val", path = "data/dolma3-50b-dsv4/code-val" }, - { name = "math-val", path = "data/dolma3-50b-dsv4/math-val" }, -] +# 5M held-out tokens of the five sources, at their training shares. +sets = [{ name = "heldout", path = "data/dolma3-50b-dsv4/mix-val-5M" }] [ckpt] interval = 1024 From 2dec0e898de3516676e75ecf7254e9d24a813f62 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Wed, 21 Oct 2026 21:58:00 +0200 Subject: [PATCH 120/132] refactor(config): rename model.debug.random_init to model.debug.init The flag now initializes the weights (N(0, 0.02), norms 1, biases 0) instead of leaving them uninitialized, so it is named for what it does. --- packages/prime-rl-configs/src/prime_rl/configs/trainer.py | 4 ++-- src/prime_rl/entrypoints/rl.py | 2 +- src/prime_rl/entrypoints/sft.py | 2 +- src/prime_rl/trainer/model.py | 6 +++--- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index d071ed371db..864b80a55be 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -103,8 +103,8 @@ class DebugModelConfig(BaseConfig): num_layers: int | None = None """Override the number of transformer layers (truncates the model).""" - random_init: bool = False - """Randomly initialize the model instead of loading weights.""" + init: bool = False + """Initialize the weights (N(0, 0.02), norms 1, biases 0) instead of loading them.""" force_balanced_routing: bool = False """Replace MoE token-choice routing with a round-robin assignment so every expert sees an equal share. Intended for fake-data smoke tests where untrained routing would otherwise OOM under severe imbalance. Gating scores are still gathered from the override indices so the forward pass stays consistent.""" diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 035c7d0963a..444338c0b41 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -507,7 +507,7 @@ def rl(config: RLConfig): if not config.dry_run: from prime_rl.trainer.model import pre_download_model - pre_download_model(config.trainer.model.name, skip_weights=config.trainer.model.debug.random_init) + pre_download_model(config.trainer.model.name, skip_weights=config.trainer.model.debug.init) if config.slurm is not None: rl_slurm(config) diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 6459768d22e..367748d881b 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -463,7 +463,7 @@ def sft(config: SFTConfig): from prime_rl.trainer.model import pre_download_model from prime_rl.trainer.sft.data import pre_download_data - pre_download_model(config.model.name, skip_weights=config.model.debug.random_init) + pre_download_model(config.model.name, skip_weights=config.model.debug.init) pre_download_data(config.data, config.env_vars) if config.val is not None: pre_download_data(config.val.data, config.env_vars) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 2385c5e8d83..66a9e5fdc85 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -79,7 +79,7 @@ def pre_download_model(model_name: str, *, skip_weights: bool = False) -> None: """Pre-download model from HuggingFace Hub so all nodes have cached weights before training. - With ``skip_weights`` (random-init debug runs), only config and tokenizer files are fetched. + With ``skip_weights`` (``model.debug.init`` runs), only config and tokenizer files are fetched. """ if Path(model_name).exists(): get_logger().info(f"Model {model_name} found at local path, skipping download") @@ -650,8 +650,8 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig): model.init_buffers_post_meta() logger = get_logger() - if config.debug.random_init: - logger.warning("Randomly initializing model. Skipping loading weights from HF.") + if config.debug.init: + logger.warning("Initializing the model weights. Skipping loading weights from HF.") init_weights(model, InitConfig()) return From 99b83346f24fc0f156b9ffcee8f6304008116b39 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 22 Oct 2026 09:58:29 +0200 Subject: [PATCH 121/132] Revert "refactor(config): rename model.debug.random_init to model.debug.init" This reverts commit 2dec0e898de3516676e75ecf7254e9d24a813f62. --- packages/prime-rl-configs/src/prime_rl/configs/trainer.py | 4 ++-- src/prime_rl/entrypoints/rl.py | 2 +- src/prime_rl/entrypoints/sft.py | 2 +- src/prime_rl/trainer/model.py | 6 +++--- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 864b80a55be..d071ed371db 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -103,8 +103,8 @@ class DebugModelConfig(BaseConfig): num_layers: int | None = None """Override the number of transformer layers (truncates the model).""" - init: bool = False - """Initialize the weights (N(0, 0.02), norms 1, biases 0) instead of loading them.""" + random_init: bool = False + """Randomly initialize the model instead of loading weights.""" force_balanced_routing: bool = False """Replace MoE token-choice routing with a round-robin assignment so every expert sees an equal share. Intended for fake-data smoke tests where untrained routing would otherwise OOM under severe imbalance. Gating scores are still gathered from the override indices so the forward pass stays consistent.""" diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 444338c0b41..035c7d0963a 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -507,7 +507,7 @@ def rl(config: RLConfig): if not config.dry_run: from prime_rl.trainer.model import pre_download_model - pre_download_model(config.trainer.model.name, skip_weights=config.trainer.model.debug.init) + pre_download_model(config.trainer.model.name, skip_weights=config.trainer.model.debug.random_init) if config.slurm is not None: rl_slurm(config) diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 367748d881b..6459768d22e 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -463,7 +463,7 @@ def sft(config: SFTConfig): from prime_rl.trainer.model import pre_download_model from prime_rl.trainer.sft.data import pre_download_data - pre_download_model(config.model.name, skip_weights=config.model.debug.init) + pre_download_model(config.model.name, skip_weights=config.model.debug.random_init) pre_download_data(config.data, config.env_vars) if config.val is not None: pre_download_data(config.val.data, config.env_vars) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 66a9e5fdc85..2385c5e8d83 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -79,7 +79,7 @@ def pre_download_model(model_name: str, *, skip_weights: bool = False) -> None: """Pre-download model from HuggingFace Hub so all nodes have cached weights before training. - With ``skip_weights`` (``model.debug.init`` runs), only config and tokenizer files are fetched. + With ``skip_weights`` (random-init debug runs), only config and tokenizer files are fetched. """ if Path(model_name).exists(): get_logger().info(f"Model {model_name} found at local path, skipping download") @@ -650,8 +650,8 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig): model.init_buffers_post_meta() logger = get_logger() - if config.debug.init: - logger.warning("Initializing the model weights. Skipping loading weights from HF.") + if config.debug.random_init: + logger.warning("Randomly initializing model. Skipping loading weights from HF.") init_weights(model, InitConfig()) return From b988f08ae136ff81d7f64b69da68348a319a889a Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 22 Oct 2026 22:52:18 +0200 Subject: [PATCH 122/132] revert(model): remove compile_inside_full_checkpoint --- src/prime_rl/trainer/model.py | 19 ++++--------------- src/prime_rl/trainer/pretrain/train.py | 1 - 2 files changed, 4 insertions(+), 16 deletions(-) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 2385c5e8d83..3893c0521f6 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -770,20 +770,13 @@ def mark_dynamic_int_args() -> None: torch.compiler.config.dynamic_sources = ",".join(sources) -def apply_compile(model: nn.Module, compile_config: CompileConfig, inside_checkpoint: bool = False): - """Compile every transformer block. With ``inside_checkpoint``, compile the block inside its activation - checkpoint wrapper instead of the checkpointed block: a graph break inside a compiled checkpoint (the EP - all-to-all, FSDP hooks) makes dynamo run the whole block eagerly, while inside an eager checkpoint it only - splits the graph.""" +def apply_compile(model: nn.Module, compile_config: CompileConfig): torch._dynamo.config.capture_scalar_outputs = True mark_dynamic_int_args() language_model = get_language_model(model) for layer_id in range(len(language_model.layers)): - layer = language_model.layers[layer_id] - if inside_checkpoint: - layer = getattr(layer, "_checkpoint_wrapped_module", layer) # Doing it in-place avoids mangled fqn which can break checkpoint loading - layer.compile(fullgraph=compile_config.fullgraph, mode=compile_config.mode) + language_model.layers[layer_id].compile(fullgraph=compile_config.fullgraph, mode=compile_config.mode) get_logger().info( f"Compiled {len(language_model.layers)} layers (fullgraph={compile_config.fullgraph}, mode={compile_config.mode})" ) @@ -883,12 +876,9 @@ def setup_model( parallel_dims: ParallelDims, loading_from_checkpoint_later: bool = False, init: InitConfig | None = None, - compile_inside_full_checkpoint: bool = False, ) -> nn.Module: """Build the sharded model. Its weights come from ``config.name``, or later from a checkpoint - (``loading_from_checkpoint_later``); with ``init``, they are initialized from scratch instead. - ``compile_inside_full_checkpoint`` compiles blocks inside their - full activation checkpoints (see ``apply_compile``).""" + (``loading_from_checkpoint_later``); with ``init``, they are initialized from scratch instead.""" resolve_auto_attn(config) if config.attn == "flash_attention_3" and not is_flash_attn_3_available(): @@ -948,8 +938,7 @@ def setup_model( if config.ac is not None: apply_ac(model, config.ac) if config.compile is not None: - inside_checkpoint = compile_inside_full_checkpoint and config.ac is not None and config.ac.mode == "full" - apply_compile(model, config.compile, inside_checkpoint=inside_checkpoint) + apply_compile(model, config.compile) setup_fsdp(model, config, parallel_dims) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py index 30ce79b0814..a2f2cc7996e 100644 --- a/src/prime_rl/trainer/pretrain/train.py +++ b/src/prime_rl/trainer/pretrain/train.py @@ -123,7 +123,6 @@ def train(config: PretrainConfig): parallel_dims, loading_from_checkpoint_later=checkpoint_step is not None, init=config.model.init if checkpoint_step is None else None, - compile_inside_full_checkpoint=True, ) if parallel_dims.cp_enabled: setup_context_parallel(model, config.model, parallel_dims) From 44cc5efcc1ab11e997815380ca4aedcd3fd4407c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 22 Oct 2026 23:04:21 +0200 Subject: [PATCH 123/132] fix(model): model.debug.random_init behaves as on main; the init zeroes unknown 1-D parameters --- src/prime_rl/trainer/model.py | 1 - src/prime_rl/trainer/models/init.py | 13 ++++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 3893c0521f6..eb008e49acb 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -652,7 +652,6 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig): logger = get_logger() if config.debug.random_init: logger.warning("Randomly initializing model. Skipping loading weights from HF.") - init_weights(model, InitConfig()) return if not Path(config.name).exists(): diff --git a/src/prime_rl/trainer/models/init.py b/src/prime_rl/trainer/models/init.py index f2a468af576..962f2def55a 100644 --- a/src/prime_rl/trainer/models/init.py +++ b/src/prime_rl/trainer/models/init.py @@ -1,5 +1,6 @@ import hashlib import math +import re import torch from torch import nn @@ -9,6 +10,7 @@ from prime_rl.configs.trainer import InitConfig from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm from prime_rl.trainer.models.layers.moe import TokenChoiceTopKRouter +from prime_rl.utils.logger import get_logger from prime_rl.utils.weights import resolve_fqn # Elements per random block. Blocks are seeded independently, so each rank draws only the @@ -52,8 +54,9 @@ def normal_shard( @torch.no_grad() def init_weights(model: nn.Module, config: InitConfig) -> None: """Initialize every parameter and the MoE selection biases in place: weight matrices and embeddings from - N(0, ``config.std``), norm weights at 1, biases at 0. Each rank fills only its local shard, without communication, - and the values do not depend on the world size or the sharding.""" + N(0, ``config.std``), norm weights at 1, biases and other 1-D parameters at 0. Each rank fills only its local + shard, without communication, and the values do not depend on the world size or the sharding.""" + zeroed = [] for name, param in model.named_parameters(): # Canonical names, so wrappers such as activation checkpointing don't change the values. name = resolve_fqn(model, name) @@ -64,7 +67,8 @@ def init_weights(model: nn.Module, config: InitConfig) -> None: elif "norm" in name: local.fill_(1.0) else: - raise ValueError(f"No init for 1-D parameter {name!r}") + local.zero_() + zeroed.append(name) continue if isinstance(param, DTensor): local_shape, global_offset = compute_local_shape_and_global_offset( @@ -76,6 +80,9 @@ def init_weights(model: nn.Module, config: InitConfig) -> None: name, tuple(param.shape), tuple(local_shape), tuple(global_offset), config.seed, local.device ) local.copy_(values.mul_(config.std)) + if zeroed: + names = sorted({re.sub(r"\.\d+\.", ".*.", name) for name in zeroed}) + get_logger().warning(f"No dedicated init for the 1-D parameters {names}: they start at 0") # MoE selection biases are buffers: they start at 0, so routing starts from the router scores alone. # Gated Residual norm gains are zero-centred (scale 1 + w), so they start at 0. for module in model.modules(): From d9295a78a4054839fdd6b7d07a1dc2fc963774d8 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 22 Oct 2026 23:05:21 +0200 Subject: [PATCH 124/132] refactor(pretrain): move prime_moe, its shared-code changes, MoE load balancing and the fractional scheduler out of this PR They move to the stacked branches feat/pretrain-prime-moe and feat/pretrain-moe-balance, and the scheduler change to feat/scheduler-fractional. ckpt.py and moe_runtime.py are back to main. --- .../pretrain-moe/model/config.json | 39 --- .../ci/integration/pretrain-moe/resume.toml | 40 --- .../ci/integration/pretrain-moe/start.toml | 37 --- examples/pretrain/moe-d512/config.json | 36 --- examples/pretrain/prime_moe.toml | 42 --- .../src/prime_rl/configs/trainer.py | 28 +- src/prime_rl/trainer/ckpt.py | 4 +- src/prime_rl/trainer/models/__init__.py | 3 - src/prime_rl/trainer/models/init.py | 4 - src/prime_rl/trainer/models/layers/attn.py | 9 +- .../trainer/models/prime_moe/__init__.py | 4 - .../trainer/models/prime_moe/attention.py | 30 -- .../prime_moe/configuration_prime_moe.py | 138 --------- .../models/prime_moe/modeling_prime_moe.py | 268 ---------------- src/prime_rl/trainer/models/prime_moe/moe.py | 289 ------------------ .../models/qwen3_8_flash_next/__init__.py | 2 +- .../hyper_connection.py | 0 .../modeling_qwen3_8_flash_next.py | 2 +- .../qwen3_8_flash_next/position_learning.py | 2 +- src/prime_rl/trainer/moe_runtime.py | 3 +- src/prime_rl/trainer/optim/__init__.py | 33 +- src/prime_rl/trainer/pretrain/balance.py | 81 ----- src/prime_rl/trainer/pretrain/optim.py | 22 -- src/prime_rl/trainer/scheduler.py | 37 +-- tests/integration/test_pretrain_moe.py | 106 ------- tests/unit/train/models/test_init.py | 21 ++ tests/unit/train/pretrain/test_init.py | 68 ----- tests/unit/train/pretrain/test_moe_balance.py | 82 ----- tests/unit/train/test_scheduler.py | 28 -- 29 files changed, 51 insertions(+), 1407 deletions(-) delete mode 100644 configs/ci/integration/pretrain-moe/model/config.json delete mode 100644 configs/ci/integration/pretrain-moe/resume.toml delete mode 100644 configs/ci/integration/pretrain-moe/start.toml delete mode 100644 examples/pretrain/moe-d512/config.json delete mode 100644 examples/pretrain/prime_moe.toml delete mode 100644 src/prime_rl/trainer/models/prime_moe/__init__.py delete mode 100644 src/prime_rl/trainer/models/prime_moe/attention.py delete mode 100644 src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py delete mode 100644 src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py delete mode 100644 src/prime_rl/trainer/models/prime_moe/moe.py rename src/prime_rl/trainer/models/{layers => qwen3_8_flash_next}/hyper_connection.py (100%) delete mode 100644 src/prime_rl/trainer/pretrain/balance.py delete mode 100644 src/prime_rl/trainer/pretrain/optim.py delete mode 100644 tests/integration/test_pretrain_moe.py create mode 100644 tests/unit/train/models/test_init.py delete mode 100644 tests/unit/train/pretrain/test_init.py delete mode 100644 tests/unit/train/pretrain/test_moe_balance.py delete mode 100644 tests/unit/train/test_scheduler.py diff --git a/configs/ci/integration/pretrain-moe/model/config.json b/configs/ci/integration/pretrain-moe/model/config.json deleted file mode 100644 index 178c5251779..00000000000 --- a/configs/ci/integration/pretrain-moe/model/config.json +++ /dev/null @@ -1,39 +0,0 @@ -{ - "model_type": "prime_moe", - "vocab_size": 151680, - "bos_token_id": null, - "eos_token_id": 151645, - "pad_token_id": 151643, - "hidden_size": 256, - "num_hidden_layers": 4, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 64, - "intermediate_size": 1024, - "hidden_act": "silu", - "max_position_embeddings": 4096, - "rms_norm_eps": 1e-06, - "partial_rotary_factor": 1.0, - "rope_parameters": { - "rope_type": "default", - "rope_theta": 10000.0, - "partial_rotary_factor": 1.0 - }, - "tie_word_embeddings": false, - "attention_bias": false, - "use_qk_norm": true, - "first_k_dense_replace": 1, - "n_routed_experts": 16, - "n_null_experts": 8, - "num_experts_per_tok": 6, - "moe_intermediate_size": 128, - "n_shared_experts": 1, - "norm_topk_prob": true, - "routed_scaling_factor": 2.5, - "hc_count": 2, - "hc_lowrank": 32, - "attn_output_gate": true, - "exclusive_attention": true, - "latent_dim": 128, - "router_bias_on_logits": true -} diff --git a/configs/ci/integration/pretrain-moe/resume.toml b/configs/ci/integration/pretrain-moe/resume.toml deleted file mode 100644 index 952ee13a76b..00000000000 --- a/configs/ci/integration/pretrain-moe/resume.toml +++ /dev/null @@ -1,40 +0,0 @@ -# Tiny prime_moe with the base components (Gated Residual, null experts, gated attention, XSA, LatentMoE, quantile -# balancing) in FP8 on 2 GPUs. The test writes the token shards (Reverse-Text-SFT, Qwen3 tokenizer) and passes their -# folder as --data.paths. Resumes the start run from its step-20 checkpoint. -max_steps = 30 - -[model] -name = "configs/ci/integration/pretrain-moe/model" -ac = "None" - -[model.quantization] -type = "fp8" - -[model.moe.compute] -type = "deepgemm_fp8" - -[model.moe.balance] -bias_update = "quantile" - -[tokenizer] -name = "PrimeIntellect/Qwen3-0.6B" - -[data] -type = "tokens" -seq_len = 2048 -batch_size = 8 -micro_batch_size = 2 -max_epochs = 100 -paths = ["set-by-the-test"] - -[optim] -type = "muon" -lr = 0.004 - -[scheduler] -type = "constant" - -[ckpt] - -[resume] -step = 20 diff --git a/configs/ci/integration/pretrain-moe/start.toml b/configs/ci/integration/pretrain-moe/start.toml deleted file mode 100644 index e171552c5ee..00000000000 --- a/configs/ci/integration/pretrain-moe/start.toml +++ /dev/null @@ -1,37 +0,0 @@ -# Tiny prime_moe with the base components (Gated Residual, null experts, gated attention, XSA, LatentMoE, quantile -# balancing) in FP8 on 2 GPUs. The test writes the token shards (Reverse-Text-SFT, Qwen3 tokenizer) and passes their -# folder as --data.paths. -max_steps = 20 - -[model] -name = "configs/ci/integration/pretrain-moe/model" -ac = "None" - -[model.quantization] -type = "fp8" - -[model.moe.compute] -type = "deepgemm_fp8" - -[model.moe.balance] -bias_update = "quantile" - -[tokenizer] -name = "PrimeIntellect/Qwen3-0.6B" - -[data] -type = "tokens" -seq_len = 2048 -batch_size = 8 -micro_batch_size = 2 -max_epochs = 100 -paths = ["set-by-the-test"] - -[optim] -type = "muon" -lr = 0.004 - -[scheduler] -type = "constant" - -[ckpt] diff --git a/examples/pretrain/moe-d512/config.json b/examples/pretrain/moe-d512/config.json deleted file mode 100644 index 84fe7c66888..00000000000 --- a/examples/pretrain/moe-d512/config.json +++ /dev/null @@ -1,36 +0,0 @@ -{ - "attention_bias": false, - "attention_dropout": 0.0, - "bos_token_id": 0, - "eos_token_id": 1, - "first_k_dense_replace": 1, - "head_dim": 64, - "hidden_act": "silu", - "hidden_size": 512, - "initializer_range": 0.02, - "intermediate_size": 2048, - "max_position_embeddings": 131072, - "model_type": "prime_moe", - "moe_intermediate_size": 256, - "n_routed_experts": 64, - "n_shared_experts": 2, - "norm_topk_prob": true, - "num_attention_heads": 8, - "num_experts_per_tok": 6, - "num_hidden_layers": 8, - "num_key_value_heads": 2, - "pad_token_id": 1, - "partial_rotary_factor": 1.0, - "rms_norm_eps": 1e-05, - "rope_parameters": { - "partial_rotary_factor": 1.0, - "rope_theta": 500000.0, - "rope_type": "default" - }, - "routed_scaling_factor": 2.5, - "tie_word_embeddings": false, - "transformers_version": "5.6.2", - "use_cache": true, - "use_qk_norm": true, - "vocab_size": 129280 -} diff --git a/examples/pretrain/prime_moe.toml b/examples/pretrain/prime_moe.toml deleted file mode 100644 index 45cf067cb98..00000000000 --- a/examples/pretrain/prime_moe.toml +++ /dev/null @@ -1,42 +0,0 @@ -# prime_moe with 64 routed experts (6 active) and 2 shared experts, d = 512 and 8 layers (moe-d512/config.json), trained from scratch on -# the Dolma 3 mixture tokenized with the DeepSeek-V4 tokenizer. -# 4096 steps x 256K tokens = 1.1B tokens. QK-norm and the learning rate follow the MoE QK-norm sweep (optimum 2^-7 -# with QK-norm at d = 512). -# -# uv run pretrain @ examples/pretrain/prime_moe.toml - -max_steps = 4096 - -[model] -name = "examples/pretrain/moe-d512" # the architecture: config.json only, no weights -ac = "None" - -[tokenizer] -name = "tokenizers/deepseek-v4-flash-base" # a local copy of deepseek-ai/DeepSeek-V4-Flash-Base without config.json - -[data] -type = "tokens" -seq_len = 8192 -batch_size = 32 -micro_batch_size = 8 -# The five source folders pooled: the mixture at the token shares they were made with. -paths = [ - "data/dolma3-50b-dsv4/web", - "data/dolma3-50b-dsv4/pdfs", - "data/dolma3-50b-dsv4/code", - "data/dolma3-50b-dsv4/math", - "data/dolma3-50b-dsv4/reference", -] - -[optim] -type = "muon" -lr = 0.0078125 -max_norm = 1.0 - -[evals] -interval = 512 -# 5M held-out tokens of the five sources, at their training shares. -sets = [{ name = "heldout", path = "data/dolma3-50b-dsv4/mix-val-5M" }] - -[ckpt] -interval = 1024 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index d071ed371db..96e0969d74c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -1,4 +1,3 @@ -import math import re import uuid import warnings @@ -393,25 +392,14 @@ class ConstantSchedulerConfig(BaseConfig): type: Literal["constant"] = "constant" -StepsOrFraction: TypeAlias = Annotated[int, Field(ge=0)] | Annotated[float, Field(gt=0, lt=1)] - - -def resolve_steps(value: int | float, max_steps: int | None) -> int: - """Steps as given, or a fraction of ``max_steps`` rounded down.""" - return math.floor(value * max_steps) if isinstance(value, float) else value - - class LinearSchedulerConfig(BaseConfig): type: Literal["linear"] = "linear" - warmup_steps: StepsOrFraction = 10 - """Warmup steps for the learning rate scheduler, or a fraction of ``max_steps``.""" - - decay_steps: StepsOrFraction = 10 - """Steps to decay the learning rate during the final portion of training, or a fraction of ``max_steps``.""" + warmup_steps: int = Field(10, ge=0) + """Warmup steps for the learning rate scheduler.""" - decay_shape: Literal["linear", "cosine", "1-sqrt"] = "linear" - """Shape of the decay from the peak to ``min_lr``.""" + decay_steps: int = Field(10, ge=0) + """Steps to decay the learning rate during the final portion of training.""" min_lr: float = Field(0.0, ge=0) """Minimum learning rate to converge to.""" @@ -442,15 +430,11 @@ def validate_scheduler(scheduler: SchedulerConfig, max_steps: int | None) -> Non if scheduler.decay_steps > 0: if max_steps is None: raise ValueError("Must specify max_steps when using a linear scheduler with decay_steps > 0") - warmup_steps = resolve_steps(scheduler.warmup_steps, max_steps) - decay_steps = resolve_steps(scheduler.decay_steps, max_steps) - if warmup_steps + decay_steps > max_steps: + if scheduler.warmup_steps + scheduler.decay_steps > max_steps: raise ValueError( - f"warmup_steps ({warmup_steps}) + decay_steps ({decay_steps}) " + f"warmup_steps ({scheduler.warmup_steps}) + decay_steps ({scheduler.decay_steps}) " f"must not exceed max_steps ({max_steps})" ) - elif isinstance(scheduler.warmup_steps, float) and max_steps is None: - raise ValueError("Must specify max_steps when warmup_steps is a fraction") if isinstance(scheduler, CosineSchedulerConfig): if max_steps is None: raise ValueError("Must specify max_steps when using a cosine scheduler") diff --git a/src/prime_rl/trainer/ckpt.py b/src/prime_rl/trainer/ckpt.py index 7205ca9f640..125abc581ba 100644 --- a/src/prime_rl/trainer/ckpt.py +++ b/src/prime_rl/trainer/ckpt.py @@ -178,7 +178,7 @@ def save_to_path( path: Path, model: nn.Module, optimizers: list[OptimizerLike], - scheduler: LRScheduler | None, + scheduler: LRScheduler, progress: Progress, dataloader: StatefulDataLoader | None = None, ): @@ -270,7 +270,7 @@ def save( step: int, model: nn.Module, optimizers: list[OptimizerLike], - scheduler: LRScheduler | None, + scheduler: LRScheduler, progress: Progress, dataloader: StatefulDataLoader | None = None, ) -> None: diff --git a/src/prime_rl/trainer/models/__init__.py b/src/prime_rl/trainer/models/__init__.py index b004af3638b..af9a0684f14 100644 --- a/src/prime_rl/trainer/models/__init__.py +++ b/src/prime_rl/trainer/models/__init__.py @@ -20,7 +20,6 @@ from prime_rl.trainer.models.llama import LlamaForCausalLM from prime_rl.trainer.models.minimax_m2 import MiniMaxM2Config, MiniMaxM2ForCausalLM from prime_rl.trainer.models.nemotron_h import NemotronHConfig, NemotronHForCausalLM -from prime_rl.trainer.models.prime_moe import PrimeMoeConfig, PrimeMoeForCausalLM from prime_rl.trainer.models.qwen3 import Qwen3ForCausalLM from prime_rl.trainer.models.qwen3_5 import ( Qwen3_5Config, @@ -46,7 +45,6 @@ AutoConfig.register("laguna", LagunaConfig, exist_ok=True) AutoConfig.register("minimax_m2", MiniMaxM2Config, exist_ok=True) AutoConfig.register("nemotron_h", NemotronHConfig, exist_ok=True) -AutoConfig.register("prime_moe", PrimeMoeConfig, exist_ok=True) AutoConfig.register("qwen3_moe", Qwen3MoeConfig, exist_ok=True) AutoConfig.register("qwen3_5", Qwen3_5Config, exist_ok=True) AutoConfig.register("qwen3_5_text", Qwen3_5TextConfig, exist_ok=True) @@ -70,7 +68,6 @@ (LagunaConfig, LagunaForCausalLM), (MiniMaxM2Config, MiniMaxM2ForCausalLM), (NemotronHConfig, NemotronHForCausalLM), - (PrimeMoeConfig, PrimeMoeForCausalLM), (Qwen3MoeConfig, Qwen3MoeForCausalLM), (Qwen3_5TextConfig, Qwen3_5ForCausalLM), (Qwen3_5MoeTextConfig, Qwen3_5ForCausalLM), diff --git a/src/prime_rl/trainer/models/init.py b/src/prime_rl/trainer/models/init.py index 962f2def55a..b7ed7ff8313 100644 --- a/src/prime_rl/trainer/models/init.py +++ b/src/prime_rl/trainer/models/init.py @@ -8,7 +8,6 @@ from torch.distributed.tensor._utils import compute_local_shape_and_global_offset from prime_rl.configs.trainer import InitConfig -from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm from prime_rl.trainer.models.layers.moe import TokenChoiceTopKRouter from prime_rl.utils.logger import get_logger from prime_rl.utils.weights import resolve_fqn @@ -84,9 +83,6 @@ def init_weights(model: nn.Module, config: InitConfig) -> None: names = sorted({re.sub(r"\.\d+\.", ".*.", name) for name in zeroed}) get_logger().warning(f"No dedicated init for the 1-D parameters {names}: they start at 0") # MoE selection biases are buffers: they start at 0, so routing starts from the router scores alone. - # Gated Residual norm gains are zero-centred (scale 1 + w), so they start at 0. for module in model.modules(): if isinstance(module, TokenChoiceTopKRouter) and module.selection_bias is not None: module.selection_bias.zero_() - if isinstance(module, ExpandedRMSNorm): - module.weight.zero_() diff --git a/src/prime_rl/trainer/models/layers/attn.py b/src/prime_rl/trainer/models/layers/attn.py index a2dee2ca304..7a8cae675f7 100644 --- a/src/prime_rl/trainer/models/layers/attn.py +++ b/src/prime_rl/trainer/models/layers/attn.py @@ -561,12 +561,9 @@ def forward( value_states = value_states.transpose(1, 2) out = self._compute_attention(query_states[0], key_states[0], value_states[0], cu_seqlens, max_seqlen) - return self.project_output(out, value_states[0], hidden_states), None - - def project_output(self, out: torch.Tensor, values: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: - """``o_proj`` of the per-head attention output ``out`` ([tokens, heads, dim]). Subclasses that change the - output first also get the values ([tokens, kv_heads, dim]) and the block input.""" - return self.o_proj(out.contiguous().view(1, out.shape[0], -1)) + attn_output = out.contiguous().view(1, out.shape[0], -1) + attn_output = self.o_proj(attn_output) + return attn_output, None ATTN_IMPL2CLASS = { diff --git a/src/prime_rl/trainer/models/prime_moe/__init__.py b/src/prime_rl/trainer/models/prime_moe/__init__.py deleted file mode 100644 index 00162be6fbb..00000000000 --- a/src/prime_rl/trainer/models/prime_moe/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from prime_rl.trainer.models.prime_moe.configuration_prime_moe import PrimeMoeConfig -from prime_rl.trainer.models.prime_moe.modeling_prime_moe import PrimeMoeForCausalLM, PrimeMoeModel - -__all__ = ["PrimeMoeConfig", "PrimeMoeForCausalLM", "PrimeMoeModel"] diff --git a/src/prime_rl/trainer/models/prime_moe/attention.py b/src/prime_rl/trainer/models/prime_moe/attention.py deleted file mode 100644 index 257237f32b7..00000000000 --- a/src/prime_rl/trainer/models/prime_moe/attention.py +++ /dev/null @@ -1,30 +0,0 @@ -import torch -import torch.nn.functional as F -from torch import nn - -from prime_rl.trainer.models.layers.attn import AttentionConfig, FlashAttention - - -class PrimeFlashAttention(FlashAttention): - """FlashAttention with two optional changes to the attention output, in this order: - - - ``exclusive`` (XSA; Zhai 2026, "Exclusive Self Attention"): remove from each head's output its component along the - token's own (normalized) value vector, so attention carries only context the residual path does not already have. - - ``output_gate`` (Qiu et al. 2025, "Gated Attention for Large Language Models", their best variant; Qwen3-Next): - multiply the output elementwise by ``sigmoid(g_proj(x))``, computed from the block input. - """ - - def __init__(self, config: AttentionConfig, flash_attn_version: int, *, exclusive: bool, output_gate: bool): - super().__init__(config, flash_attn_version=flash_attn_version) - self.exclusive = exclusive - self.g_proj = nn.Linear(config.hidden_size, self.qkv_sizes[0], bias=False) if output_gate else None - - def project_output(self, out: torch.Tensor, values: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: - if self.exclusive: - # Each query head's own value is its key/value group's. - values = F.normalize(values, dim=-1).repeat_interleave(self.num_key_value_groups, dim=1) - out = out - (out * values).sum(dim=-1, keepdim=True) * values - attn_output = out.contiguous().view(1, out.shape[0], -1) - if self.g_proj is not None: - attn_output = attn_output * torch.sigmoid(self.g_proj(hidden_states)) - return self.o_proj(attn_output) diff --git a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py deleted file mode 100644 index 759e2fac48f..00000000000 --- a/src/prime_rl/trainer/models/prime_moe/configuration_prime_moe.py +++ /dev/null @@ -1,138 +0,0 @@ -from transformers.configuration_utils import PretrainedConfig - - -class PrimeMoeConfig(PretrainedConfig): - r""" - Configuration of the pretraining MoE (`PrimeMoeForCausalLM`): grouped-query attention (optional QK norm, gated - attention, exclusive self attention), MoE layers with sigmoid-scored token-choice routed experts, a selection bias, - shared experts and optional null experts or LatentMoE, after leading dense layers, and optional Gated Residual - connections. Fields that also exist in the Hugging Face MoE configs keep their names there. - - Args: - vocab_size (`int`): Vocabulary size. - hidden_size (`int`): Width of the residual stream. - intermediate_size (`int`): Hidden size of the dense MLP layers. - num_hidden_layers (`int`): Number of decoder layers. - num_attention_heads (`int`): Number of query heads. - num_key_value_heads (`int`): Number of key/value heads (grouped-query attention). - head_dim (`int`, *optional*): Dimension of each attention head. Defaults to `hidden_size // num_attention_heads`. - partial_rotary_factor (`float`): Share of each head's dimensions that RoPE rotates. - hidden_act (`str`): Activation of the gated MLPs. - max_position_embeddings (`int`): Maximum sequence length. - rms_norm_eps (`float`): Epsilon of the RMS norms. - tie_word_embeddings (`bool`): Whether the LM head shares the input embedding. - rope_theta (`float`): RoPE base period. - rope_scaling (`dict`, *optional*): RoPE scaling configuration. - attention_bias (`bool`): Whether the attention projections have biases. - moe_intermediate_size (`int`): Hidden size of each routed and shared expert. - num_experts_per_tok (`int`): Experts each token selects (`top_k`), null experts included. - n_shared_experts (`int`): Number of shared experts; their hidden sizes add up. - shared_expert_intermediate_size (`int`, *optional*): Hidden size of the shared experts together. Defaults to - `moe_intermediate_size * n_shared_experts`. - n_routed_experts (`int`): Number of routed experts. - routed_scaling_factor (`float`): Scale of the routed experts' output. - first_k_dense_replace (`int`): Number of leading dense layers. - norm_topk_prob (`bool`): Whether the selected experts' scores are normalized to sum to 1. - use_qk_norm (`bool`): Whether queries and keys are RMS-normalized. - hc_count (`int`): Number of residual streams. Above 1, Gated Residual connections (Qwen3.8-Next) replace the - pre-norms and the final norm. - hc_lowrank (`int`, *optional*): Rank of the Gated Residual read gates; required when `hc_count` is above 1. - n_null_experts (`int`): Zero-output experts the router can select next to the routed experts (AdaMoE); a - token that selects them uses fewer routed experts. - attn_output_gate (`bool`): Elementwise sigmoid gate on the attention output, computed from the block input - (gated attention, Qwen3-Next). - exclusive_attention (`bool`): Exclusive self attention (XSA): remove from each head's output its component - along the token's own value vector. - latent_dim (`int`, *optional*): LatentMoE: the routed experts work in this width (`hidden_size -> latent_dim` - projection and RMS norm before them, `latent_dim -> hidden_size` projection after the combine), with hidden - size `moe_intermediate_size`. The router and the shared experts stay full width. - latent_norm (`str`): Where LatentMoE's RMS norm sits: `"input"` on the projected expert input (Marin) or - `"output"` on the combined expert output before the up-projection (Kimi K3's Stable LatentMoE). - expert_activation (`str`, *optional*): Activation of the routed and shared experts: `"situ_glu"` (Kimi K3's - soft-capped SwiGLU). Defaults to `hidden_act`. - router_bias_on_logits (`bool`): Add the selection bias to the router logits instead of the sigmoid scores - (the space quantile balancing, the default, works best in). - pad_token_id (`int`, *optional*): Padding token id, the input embedding's `padding_idx`. - """ - - model_type = "prime_moe" - - def __init__( - self, - vocab_size=151552, - hidden_size=4096, - intermediate_size=10944, - num_hidden_layers=46, - num_attention_heads=96, - partial_rotary_factor=0.5, - num_key_value_heads=8, - head_dim=None, - hidden_act="silu", - max_position_embeddings=131072, - rms_norm_eps=1e-5, - tie_word_embeddings=False, - rope_theta=10000.0, - rope_scaling=None, - attention_bias=False, - moe_intermediate_size=1408, - num_experts_per_tok=8, - n_shared_experts=1, - n_routed_experts=128, - routed_scaling_factor=1.0, - shared_expert_intermediate_size=None, - first_k_dense_replace=1, - norm_topk_prob=True, - use_qk_norm=False, - hc_count=1, - hc_lowrank=None, - n_null_experts=0, - attn_output_gate=False, - exclusive_attention=False, - latent_dim=None, - latent_norm="input", - expert_activation=None, - router_bias_on_logits=True, - pad_token_id=None, - **kwargs, - ): - self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size - self.num_hidden_layers = num_hidden_layers - self.num_attention_heads = num_attention_heads - self.partial_rotary_factor = partial_rotary_factor - self.num_key_value_heads = num_key_value_heads - self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads - self.hidden_act = hidden_act - self.rms_norm_eps = rms_norm_eps - self.rope_theta = rope_theta - self.rope_scaling = rope_scaling - self.attention_bias = attention_bias - self.pad_token_id = pad_token_id - - self.moe_intermediate_size = moe_intermediate_size - self.num_experts_per_tok = num_experts_per_tok - self.n_shared_experts = n_shared_experts - self.shared_expert_intermediate_size = ( - shared_expert_intermediate_size - if shared_expert_intermediate_size is not None - else moe_intermediate_size * n_shared_experts - ) - self.n_routed_experts = n_routed_experts - self.routed_scaling_factor = routed_scaling_factor - self.first_k_dense_replace = first_k_dense_replace - self.norm_topk_prob = norm_topk_prob - self.use_qk_norm = use_qk_norm - - self.hc_count = hc_count - self.hc_lowrank = hc_lowrank - self.n_null_experts = n_null_experts - self.attn_output_gate = attn_output_gate - self.exclusive_attention = exclusive_attention - self.latent_dim = latent_dim - self.latent_norm = latent_norm - self.expert_activation = expert_activation - self.router_bias_on_logits = router_bias_on_logits - - super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) diff --git a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py b/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py deleted file mode 100644 index b78993bba44..00000000000 --- a/src/prime_rl/trainer/models/prime_moe/modeling_prime_moe.py +++ /dev/null @@ -1,268 +0,0 @@ -"""The pretraining MoE (`prime_moe`), pretraining's own model, so architecture work leaves the shared MoE models -unchanged. Module names follow `glm4_moe`, so checkpoints keep that state-dict layout and its conversion.""" - -from typing import Optional - -import torch -from torch import Tensor, nn -from transformers.modeling_layers import GradientCheckpointingLayer -from transformers.modeling_outputs import BaseModelOutputWithPast -from transformers.processing_utils import Unpack -from transformers.utils import TransformersKwargs - -from prime_rl.trainer.models.base import PreTrainedModelPrimeRL -from prime_rl.trainer.models.glm4_moe.converting_glm4_moe import conversion_chain -from prime_rl.trainer.models.layers.attn import ATTN_IMPL2CLASS, AttentionConfig -from prime_rl.trainer.models.layers.hyper_connection import HyperConnection -from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput -from prime_rl.trainer.models.layers.mlp import FeedForward -from prime_rl.trainer.models.layers.moe import GroupedExperts -from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig -from prime_rl.trainer.models.layers.rotary_emb import RotaryEmbedding, RotaryEmbeddingConfig -from prime_rl.trainer.models.prime_moe.attention import PrimeFlashAttention -from prime_rl.trainer.models.prime_moe.configuration_prime_moe import PrimeMoeConfig -from prime_rl.trainer.models.prime_moe.moe import PrimeMoE, PrimeRouter, SituGlu -from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens - - -class PrimeMoeDecoderLayer(GradientCheckpointingLayer): - def __init__(self, config: PrimeMoeConfig, layer_idx: int): - super().__init__() - attn_config = AttentionConfig( - hidden_size=config.hidden_size, - head_dim=config.head_dim, - num_attention_heads=config.num_attention_heads, - num_key_value_heads=config.num_key_value_heads, - is_causal=True, - attention_bias=config.attention_bias, - use_qk_norm=config.use_qk_norm, - rms_norm_eps=config.rms_norm_eps, - ) - if config.attn_output_gate or config.exclusive_attention: - self.self_attn = PrimeFlashAttention( - attn_config, - flash_attn_version=int(config._attn_implementation.removeprefix("flash_attention_")), - exclusive=config.exclusive_attention, - output_gate=config.attn_output_gate, - ) - else: - self.self_attn = ATTN_IMPL2CLASS[config._attn_implementation](attn_config) - - if layer_idx >= config.first_k_dense_replace: - expert_activation = SituGlu() if config.expert_activation == "situ_glu" else config.hidden_act - shared_expert = None - if config.shared_expert_intermediate_size > 0: - shared_expert = FeedForward( - dim=config.hidden_size, - hidden_dim=config.shared_expert_intermediate_size, - expert_type="gated", - activation=expert_activation, - ) - self.mlp = PrimeMoE( - router=PrimeRouter( - dim=config.hidden_size, - num_experts=config.n_routed_experts, - num_null_experts=config.n_null_experts, - top_k=config.num_experts_per_tok, - route_norm=config.norm_topk_prob, - route_scale=config.routed_scaling_factor, - bias_on_logits=config.router_bias_on_logits, - ), - experts=GroupedExperts( - dim=config.latent_dim or config.hidden_size, - hidden_dim=config.moe_intermediate_size, - num_experts=config.n_routed_experts, - expert_type="gated", - activation=expert_activation, - ), - shared_expert=shared_expert, - dim=config.hidden_size, - latent_dim=config.latent_dim, - latent_norm=config.latent_norm, - norm_eps=config.rms_norm_eps, - ) - else: - self.mlp = FeedForward( - dim=config.hidden_size, hidden_dim=config.intermediate_size, activation=config.hidden_act - ) - - self.gated_residual = config.hc_count > 1 - if self.gated_residual: - connection_args = { - "hidden_size": config.hidden_size, - "stream_count": config.hc_count, - "low_rank": config.hc_lowrank, - "norm_eps": config.rms_norm_eps, - } - self.attn_hyper_connection = HyperConnection(**connection_args) - self.mlp_hyper_connection = HyperConnection(**connection_args) - else: - self.input_layernorm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) - self.post_attention_layernorm = RMSNorm( - RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps) - ) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - cu_seqlens: torch.LongTensor, - max_seqlen: int, - ) -> torch.Tensor: - if self.gated_residual: - block_input, residual_state = self.attn_hyper_connection.mix(hidden_states) - block_output, _ = self.self_attn( - hidden_states=block_input, - position_embeddings=position_embeddings, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - ) - hidden_states = self.attn_hyper_connection.combine(block_output, residual_state) - block_input, residual_state = self.mlp_hyper_connection.mix(hidden_states) - block_output = self.mlp(block_input) - return self.mlp_hyper_connection.combine(block_output, residual_state) - - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states, _ = self.self_attn( - hidden_states=hidden_states, - position_embeddings=position_embeddings, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - ) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - return residual + hidden_states - - -class PrimeMoePreTrainedModel(PreTrainedModelPrimeRL): - config: PrimeMoeConfig - base_model_prefix = "model" - supports_gradient_checkpointing = True - _no_split_modules = ["PrimeMoeDecoderLayer"] - _supports_flash_attn = True - _supports_sdpa = False - _supports_flex_attn = True - _can_compile_fullgraph = False - _supports_attention_backend = True - - @classmethod - def keep_in_fp32_for_weight_transfer(cls, name: str) -> bool: - return name.endswith("mlp.router.selection_bias") - - @classmethod - def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - return any("mlp.experts.1.up_proj" in name or "mlp.experts.gate_up_proj" in name for name in state_dict) - - @classmethod - def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: - return any("mlp.experts.gate_proj" in name for name in state_dict) - - @classmethod - def conversion_chain(cls, config): - # The expert state-dict layout of glm4_moe, so its conversion applies. - return conversion_chain(config) - - -class PrimeMoeModel(PrimeMoePreTrainedModel): - def __init__(self, config: PrimeMoeConfig): - super().__init__(config) - self.padding_idx = config.pad_token_id - self.vocab_size = config.vocab_size - - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) - self.layers = nn.ModuleList( - [PrimeMoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] - ) - if config.hc_count > 1: - self.hyper_connection_mixer = HyperConnection( - hidden_size=config.hidden_size, - stream_count=config.hc_count, - low_rank=config.hc_lowrank, - norm_eps=config.rms_norm_eps, - with_residual_injection=False, - ) - else: - self.norm = RMSNorm(RMSNormConfig(hidden_size=config.hidden_size, eps=config.rms_norm_eps)) - - if isinstance(config.rope_scaling, dict): - rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) - else: - rope_type = "default" - self.rotary_emb = RotaryEmbedding( - RotaryEmbeddingConfig( - max_position_embeddings=config.max_position_embeddings, rope_type=rope_type, model_config=config - ) - ) - self.gradient_checkpointing = False - self.post_init() - - def forward( - self, - input_ids: torch.LongTensor, - position_ids: torch.LongTensor, - *, - seq_lens: torch.LongTensor, - seq_lens_are_pre_shard: bool = False, - ) -> BaseModelOutputWithPast: - inputs_embeds = self.embed_tokens(input_ids) - - cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens( - seq_lens.to(device=inputs_embeds.device), - total_tokens=None if seq_lens_are_pre_shard else inputs_embeds.shape[1], - ) - torch._dynamo.mark_dynamic(cu_seqlens, 0) - - position_embeddings = self.rotary_emb(inputs_embeds, position_ids) - # Gated Residual: every residual stream starts from the token embedding. - hidden_states = inputs_embeds.repeat(1, 1, self.config.hc_count) - for decoder_layer in self.layers: - hidden_states = decoder_layer( - hidden_states, position_embeddings=position_embeddings, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen - ) - - if self.config.hc_count > 1: - hidden_states, _ = self.hyper_connection_mixer(hidden_states) - else: - hidden_states = self.norm(hidden_states) - return BaseModelOutputWithPast(last_hidden_state=hidden_states) - - -class PrimeMoeForCausalLM(PrimeMoePreTrainedModel): - def __init__(self, config: PrimeMoeConfig): - super().__init__(config) - self.model = PrimeMoeModel(config) - self.vocab_size = config.vocab_size - self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - self.post_init() - - def forward( - self, - input_ids: torch.LongTensor, - position_ids: Optional[torch.LongTensor] = None, - labels: Optional[torch.LongTensor] = None, - temperature: Optional[torch.Tensor] = None, - *, - seq_lens: torch.LongTensor, - seq_lens_are_pre_shard: bool = False, - **kwargs: Unpack[TransformersKwargs], - ) -> PrimeLmOutput: - """``labels`` and ``temperature`` go to prime-rl's wrapped LM head, which then returns per-token logprobs - and entropy; without them it returns logits.""" - if position_ids is None: - position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).unsqueeze(0) - outputs = self.model( - input_ids=input_ids, - position_ids=position_ids, - seq_lens=seq_lens, - seq_lens_are_pre_shard=seq_lens_are_pre_shard, - ) - return self.lm_head(outputs.last_hidden_state, labels, temperature=temperature) - - def init_buffers_post_meta(self): - rotary_emb = self.model.rotary_emb - inv_freq, rotary_emb.attention_scaling = rotary_emb.rope_init_fn(rotary_emb.config, rotary_emb.inv_freq.device) - rotary_emb.inv_freq.copy_(inv_freq) diff --git a/src/prime_rl/trainer/models/prime_moe/moe.py b/src/prime_rl/trainer/models/prime_moe/moe.py deleted file mode 100644 index 2a8d79b8560..00000000000 --- a/src/prime_rl/trainer/models/prime_moe/moe.py +++ /dev/null @@ -1,289 +0,0 @@ -from typing import Literal - -import torch -from torch import nn - -from prime_rl.trainer.distributed.token_dispatcher import TokenDispatcher -from prime_rl.trainer.models.layers.activations import Activation -from prime_rl.trainer.models.layers.mlp import FeedForward -from prime_rl.trainer.models.layers.moe import ( - GroupedExperts, - MoE, - TokenChoiceTopKRouter, - record_moe_routing_statistics, -) -from prime_rl.trainer.models.layers.norms import RMSNorm, RMSNormConfig - - -class SituGlu(Activation): - """Sigmoid Tanh Unit GLU (Kimi K3 report, eq. 12): SwiGLU with both factors soft-capped, - ``(4 tanh(gate / 4) * sigmoid(gate)) * (25 tanh(up / 25))``; close to SwiGLU near 0 and bounded by 100.""" - - @staticmethod - def apply(gate: torch.Tensor | None, up: torch.Tensor) -> torch.Tensor: - assert gate is not None - return 4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate) * (25.0 * torch.tanh(up / 25.0)) - - -class PrimeRouter(TokenChoiceTopKRouter): - """Sigmoid-scored top-k router with a selection bias and ``num_null_experts`` zero-output experts next to the - routed ones (AdaMoE). Null experts are scored, selected and balanced like the others, so a token that selects - them uses fewer routed experts. Returns the scores of all experts too, for the sequence-wise balance loss. - - With ``bias_on_logits`` the selection bias is added to the logits instead of the sigmoid scores. With - ``quantile_balance`` set by the training loop, every forward with grad records each expert's quantile balancing - threshold (``quantile_sum``, ``quantile_count``), from which the loop sets the next step's bias. If the loop also - sets ``quantile_hist`` (``(num_experts, bins)`` counts) and ``quantile_grid`` (its ``[lo, hi]``), the forward - instead adds each expert's required biases to that histogram (``record_quantile_histogram``).""" - - def __init__( - self, - *, - dim: int, - num_experts: int, - num_null_experts: int, - top_k: int, - route_norm: bool, - route_scale: float, - bias_on_logits: bool, - ): - super().__init__( - dim=dim, - num_experts=num_experts + num_null_experts, - top_k=top_k, - score_func="sigmoid", - route_norm=route_norm, - route_scale=route_scale, - selection_bias=True, - ) - self.num_null_experts = num_null_experts - self.bias_on_logits = bias_on_logits - self.quantile_balance = False - self.register_buffer("quantile_sum", torch.zeros(self.num_experts, dtype=torch.float32), persistent=False) - self.register_buffer("quantile_count", torch.zeros((), dtype=torch.float32), persistent=False) - self.quantile_hist: torch.Tensor | None = None - self.quantile_grid: torch.Tensor | None = None - self.quantile_range: torch.Tensor | None = None - - @torch.no_grad() - def record_quantiles(self, selection_scores: torch.Tensor, biased_scores: torch.Tensor) -> None: - """Quantile balancing (Su 2026; Kimi K3, Marin): each token's threshold is its largest unselected biased score, - and each expert's new bias is minus the margin over that threshold that ``top_k / num_experts`` of the tokens - exceed, so that with it every expert would have been selected by exactly that share of these tokens.""" - thresholds = biased_scores.topk(self.top_k + 1, dim=1).values[:, -1:] - margins = selection_scores - thresholds - rank = max(1, round(margins.shape[0] * self.top_k / self.num_experts)) - self.quantile_sum += margins.topk(rank, dim=0).values[-1] - self.quantile_count += 1 - - @torch.no_grad() - def record_quantile_histogram(self, selection_scores: torch.Tensor, biased_scores: torch.Tensor) -> None: - """Histogram estimator of quantile balancing (Kimi K3 report, §2.3.3 and appendix D): add each expert's - required biases ``threshold - score`` (the bias that puts the expert exactly at the token's threshold) to its - row of ``quantile_hist`` on the grid ``quantile_grid``. Values outside the grid count in the edge bins, so the - counts below any in-grid value stay exact. Also tracks the values' ``[min, max]`` in ``quantile_range``.""" - thresholds = biased_scores.topk(self.top_k + 1, dim=1).values[:, -1:] - required = thresholds - selection_scores - num_experts, bins = self.quantile_hist.shape - lo, hi = self.quantile_grid - index = ((required - lo) * (bins / (hi - lo))).floor().clamp_(0, bins - 1).long() - index += torch.arange(num_experts, device=index.device) * bins - # scatter_add_ instead of bincount: CUDA bincount reads its input's maximum to the host, one sync per call. - self.quantile_hist.view(-1).scatter_add_( - 0, index.view(-1), torch.ones_like(index.view(-1), dtype=self.quantile_hist.dtype) - ) - self.quantile_range[0] = torch.minimum(self.quantile_range[0], required.min()) - self.quantile_range[1] = torch.maximum(self.quantile_range[1], required.max()) - - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Like ``TokenChoiceTopKRouter.forward`` with sigmoid scores; also returns the ``(bs*slen, num_experts)`` - scores of all experts.""" - logits = self.gate(x) - scores = torch.sigmoid(logits) - # The selection bias only affects which experts are selected, not their scores. - selection_scores = logits.float() if self.bias_on_logits else scores - biased_scores = selection_scores + self.selection_bias - _, selected_experts_indices = torch.topk(biased_scores, k=self.top_k, dim=1, sorted=self.topk_sorted) - if self.quantile_balance and torch.is_grad_enabled(): - if self.quantile_hist is None: - self.record_quantiles(selection_scores, biased_scores) - else: - self.record_quantile_histogram(selection_scores, biased_scores) - top_scores = scores.gather(dim=1, index=selected_experts_indices) - - with torch.no_grad(): - routing_confidence_sum = (top_scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)).sum() - - if self.route_norm: - top_scores = top_scores / (top_scores.sum(dim=-1, keepdim=True) + 1e-20) - top_scores = top_scores * self.route_scale - - num_tokens_per_expert = torch.histc( - selected_experts_indices.reshape(-1).float(), bins=self.num_experts, min=0, max=self.num_experts - ).to(torch.int64) - return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, scores - - -def histogram_quantile_biases(hist: torch.Tensor, grid: torch.Tensor, top_k: int) -> torch.Tensor: - """Each expert's ``top_k / num_experts`` quantile of the required biases counted in ``hist`` (``(num_experts, - bins)`` on ``grid = [lo, hi]``), mean-centred: the first bin whose cumulative count reaches the target ``q`` plus - linear interpolation inside it. With it, every expert would have been selected by ``q`` of the counted tokens. - ``q`` comes from each expert's own count, so a forward recomputed under activation checkpointing (which counts - every token twice) leaves the result unchanged.""" - num_experts, bins = hist.shape - cumulative = hist.cumsum(dim=1) - target = cumulative[:, -1] * top_k / num_experts - selected = (cumulative < target.ceil()[:, None]).sum(dim=1).clamp(max=bins - 1)[:, None] - inside = hist.gather(1, selected) - below = cumulative.gather(1, selected) - inside - fraction = ((target[:, None] - below) / inside.clamp(min=1)).clamp(0, 1) - lo, hi = grid - biases = (lo + (selected + fraction) * (hi - lo) / bins).squeeze(1) - return biases - biases.mean() - - -def sequence_balance(scores: torch.Tensor, selected_experts: torch.Tensor, num_rows: int) -> torch.Tensor: - """DeepSeek-V3's sequence-wise balance loss without its coefficient, per row: ``sum_i f_i P_i`` with - ``f_i`` the row's share of expert ``i``'s selections times ``num_experts / top_k`` (1 when balanced) and - ``P_i`` its mean score normalized over all experts. 1 for a balanced row, up to ``num_experts / top_k``. - - Args: - scores: ``(num_rows * slen, num_experts)`` scores of all experts. - selected_experts: ``(num_rows * slen, top_k)`` selected experts. - """ - num_experts = scores.shape[1] - selections = selected_experts.view(num_rows, -1) - counts = torch.zeros(num_rows, num_experts, device=scores.device).scatter_add_( - 1, selections, torch.ones_like(selections, dtype=torch.float32) - ) - f = counts * (num_experts / selections.shape[1]) - p = (scores / scores.sum(dim=-1, keepdim=True)).view(num_rows, -1, num_experts).mean(dim=1) - return (f * p).sum(dim=-1) - - -class PrimeMoE(MoE): - """The shared token-choice MoE runtime with a ``PrimeRouter``. With null experts, tokens use different numbers of - routed experts: only the (token, routed expert) pairs are dispatched, one per row. Also computes the sequence-wise - balance for its loss (``balance_rows``) and counts the routed experts per token (``routed_k_counts``). - - With ``latent_dim`` (LatentMoE), the routed experts work in that width: the token is projected down before the - dispatch and the combined expert output is projected back up, with an RMS norm on the expert input - (``latent_norm = "input"``, Marin) or on the combined output before the up-projection (``"output"``, Kimi K3's - Stable LatentMoE). The router and the shared experts read the full-width token.""" - - def __init__( - self, - *, - router: PrimeRouter, - experts: GroupedExperts, - shared_expert: FeedForward | None, - dim: int, - latent_dim: int | None, - latent_norm: Literal["input", "output"], - norm_eps: float, - ) -> None: - super().__init__( - router=router, - experts=experts, - shared_expert=shared_expert, - score_before_experts=False, - load_balance_coeff=1e-3, - ) - self.variable_k = router.num_null_experts > 0 - self.set_token_dispatcher(self.token_dispatcher) - # Usage of every expert, null experts included. - self.register_buffer( - "tokens_per_expert", torch.zeros(router.num_experts, dtype=torch.float32), persistent=False - ) - # Training tokens per number of routed experts used, for variable-k routing. - self.register_buffer( - "routed_k_counts", - torch.zeros(router.top_k + 1, dtype=torch.float32) if self.variable_k else None, - persistent=False, - ) - # Sequence-wise balance (DeepSeek-V3 eq. 17-20): with balance_rows set by the training loop, which knows - # how many packed rows a sample holds, every forward with grad computes the per-row balance for its loss. - self.balance_rows: int | None = None - self.seq_balance: torch.Tensor | None = None - self.latent_down = self.latent_norm = self.latent_up = None - self.norm_latent_output = latent_norm == "output" - if latent_dim is not None: - self.latent_down = nn.Linear(dim, latent_dim, bias=False) - self.latent_norm = RMSNorm(RMSNormConfig(hidden_size=latent_dim, eps=norm_eps)) - self.latent_up = nn.Linear(latent_dim, dim, bias=False) - - def prepare_expert_input(self, x: torch.Tensor) -> torch.Tensor: - if self.latent_down is None: - return x - x = self.latent_down(x) - return x if self.norm_latent_output else self.latent_norm(x) - - def prepare_expert_output(self, x: torch.Tensor) -> torch.Tensor: - if self.latent_up is None: - return x - if self.norm_latent_output: - x = self.latent_norm(x) - return self.latent_up(x) - - def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None: - # Variable-k routing dispatches one (token, expert) pair per row. - token_dispatcher.top_k = 1 if self.variable_k else self.router.top_k - self.token_dispatcher = token_dispatcher - - def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) -> torch.Tensor: - assert routed_experts is None, "PrimeMoE does not support router replay" - bs, slen, dim = x.shape - x = x.view(-1, dim) - top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, scores = self.router(x) - - if self.balance_rows is not None and torch.is_grad_enabled(): - self.seq_balance = sequence_balance(scores, selected_experts_indices, self.balance_rows) - - with torch.no_grad(): - record_moe_routing_statistics( - self.tokens_per_expert, self.routing_confidence_sum, num_tokens_per_expert, routing_confidence_sum - ) - - expert_input = self.prepare_expert_input(x) - if self.variable_k: - # Dispatch only the (token, routed expert) pairs, without the null experts. - routed = selected_experts_indices < self.experts.num_experts - if torch.is_grad_enabled(): - with torch.no_grad(): - self.routed_k_counts += torch.bincount(routed.sum(dim=1), minlength=self.routed_k_counts.numel()) - token_indices, slots = routed.nonzero(as_tuple=True) - pair_output = self.token_dispatcher.run( - expert_input[token_indices], - top_scores[token_indices, slots].unsqueeze(1), - selected_experts_indices[token_indices, slots].unsqueeze(1), - self.experts, - score_before_experts=self.score_before_experts, - ) - routed_output = ( - torch.zeros(expert_input.shape, dtype=torch.float32, device=x.device) - .index_add(0, token_indices, pair_output.float()) - .to(x.dtype) - ) - else: - routed_output = self.token_dispatcher.run( - expert_input, - top_scores, - selected_experts_indices, - self.experts, - score_before_experts=self.score_before_experts, - ) - - shared_output = self.shared_expert(x) if self.shared_expert is not None else None - self.token_dispatcher.synchronize() - routed_output = self.prepare_expert_output(routed_output) - if shared_output is not None: - routed_output = routed_output + shared_output - return routed_output.reshape(bs, slen, dim) - - def init_weights(self, init_std: float, buffer_device: torch.device): - super().init_weights(init_std, buffer_device) - with torch.device(buffer_device): - self.tokens_per_expert = torch.zeros(self.router.num_experts, dtype=torch.float32) - self.router.selection_bias = torch.zeros(self.router.num_experts, dtype=torch.float32) - if self.routed_k_counts is not None: - self.routed_k_counts = torch.zeros_like(self.routed_k_counts) diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/__init__.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/__init__.py index 0f89c639c65..03f2798043c 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/__init__.py +++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/__init__.py @@ -1,9 +1,9 @@ -from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm, HyperConnection from prime_rl.trainer.models.qwen3_8_flash_next.attention import IndexedGatedAttention from prime_rl.trainer.models.qwen3_8_flash_next.configuration_qwen3_8_flash_next import ( Qwen3_8FlashNextConfig, Qwen3_8FlashNextTextConfig, ) +from prime_rl.trainer.models.qwen3_8_flash_next.hyper_connection import ExpandedRMSNorm, HyperConnection from prime_rl.trainer.models.qwen3_8_flash_next.modeling_qwen3_8_flash_next import ( Qwen3_8FlashNextDecoderLayer, Qwen3_8FlashNextForCausalLM, diff --git a/src/prime_rl/trainer/models/layers/hyper_connection.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/hyper_connection.py similarity index 100% rename from src/prime_rl/trainer/models/layers/hyper_connection.py rename to src/prime_rl/trainer/models/qwen3_8_flash_next/hyper_connection.py diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py index 43706b84f3e..2dad01e9031 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py +++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py @@ -4,7 +4,6 @@ from prime_rl.trainer.models.base import PreTrainedModelPrimeRL from prime_rl.trainer.models.layers.gated_delta_net import GatedDeltaNet -from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm, HyperConnection from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear from prime_rl.trainer.models.layers.moe import MoE from prime_rl.trainer.models.layers.norms import ZeroCenteredRMSNorm @@ -18,6 +17,7 @@ is_hf_state_dict, is_prime_state_dict, ) +from prime_rl.trainer.models.qwen3_8_flash_next.hyper_connection import ExpandedRMSNorm, HyperConnection from prime_rl.trainer.models.qwen3_8_flash_next.moe import SigmoidOutputGatedMoE from prime_rl.trainer.models.qwen3_8_flash_next.ngram_embedding import NGramEmbedding from prime_rl.trainer.models.qwen3_8_flash_next.position_learning import PositionLearningEnhancement diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/position_learning.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/position_learning.py index df7ea6aeb34..f6e70f5db4a 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/position_learning.py +++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/position_learning.py @@ -4,7 +4,7 @@ import torch.nn.functional as F from torch import nn -from prime_rl.trainer.models.layers.hyper_connection import ExpandedRMSNorm +from prime_rl.trainer.models.qwen3_8_flash_next.hyper_connection import ExpandedRMSNorm from prime_rl.trainer.models.qwen3_8_flash_next.ngram_embedding import NGramEmbedding from prime_rl.utils.cp import CPContext, gather_for_cp diff --git a/src/prime_rl/trainer/moe_runtime.py b/src/prime_rl/trainer/moe_runtime.py index 5308749255e..d0f01e8a4aa 100644 --- a/src/prime_rl/trainer/moe_runtime.py +++ b/src/prime_rl/trainer/moe_runtime.py @@ -86,8 +86,7 @@ def _resolve_mega_moe_compute( def configure_moe_runtime(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims) -> None: moe_layers = [module for module in model.modules() if isinstance(module, MoE)] if not moe_layers: - default = MoERuntimeConfig() - if (config.moe.compute, config.moe.dispatch) != (default.compute, default.dispatch): + if config.moe != MoERuntimeConfig(): raise ValueError("A non-default model.moe runtime was configured, but the model has no custom MoE layers.") return diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index 7feac0bca15..f2551412536 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/src/prime_rl/trainer/optim/__init__.py @@ -54,10 +54,7 @@ def setup_optimizer( offload: OffloadMode = "none", model: nn.Module | None = None, full_offload_dtype_policy: dict[int, tuple[torch.dtype, torch.dtype]] | None = None, - param_groups: list[dict] | None = None, ) -> tuple[OptimizerLike, GradientOffloadManager | None]: - """``param_groups`` replaces the default parameter groups (for Muon, each group needs an - ``algorithm`` key). It is not supported with full offload.""" cpu_offload = offload == "optimizer" full_offload = offload == "full" if full_offload and config.type not in ("adamw", "sign_sgd"): @@ -82,7 +79,6 @@ def setup_optimizer( parallel_dims, fused_adamw=config.type == "adamw" and not cpu_offload, model=model, - param_groups=param_groups, ) if full_offload: @@ -109,7 +105,6 @@ def _create_optimizer( lr: float | None = None, fused_adamw: bool = False, model: nn.Module | None = None, - param_groups: list[dict] | None = None, ) -> Optimizer: """Create optimizer. If lr is None, uses config.lr.""" if lr is None: @@ -122,7 +117,7 @@ def _create_optimizer( match config.type: case "adamw": return AdamW( - params=param_groups or trainable_params, + params=trainable_params, lr=lr, weight_decay=config.weight_decay, betas=(config.betas1, config.betas2), @@ -130,7 +125,7 @@ def _create_optimizer( fused=fused_adamw, ) case "muon": - return _create_muon_optimizer(config, named_params, parallel_dims, model, lr, param_groups) + return _create_muon_optimizer(config, named_params, parallel_dims, model, lr) case "sign_sgd": return SignSGD( params=trainable_params, @@ -139,12 +134,13 @@ def _create_optimizer( ) -def _default_muon_param_groups( +def _create_muon_optimizer( config: OptimizerConfig, named_params: list[tuple[str, nn.Parameter]], parallel_dims: ParallelDims, - lr: float, -) -> list[dict]: + model: nn.Module | None, + lr: float | None = None, +) -> Optimizer: def muon_enabled(n, p): if p.ndim < 2: return False @@ -173,6 +169,8 @@ def muon_enabled(n, p): muon_params.append(p) elif p.requires_grad: adamw_params.append(p) + else: + pass param_groups = [] @@ -205,19 +203,6 @@ def muon_enabled(n, p): ) param_groups.append(dict(params=adamw_params, algorithm="adamw", lr=lr, weight_decay=config.weight_decay)) - return param_groups - - -def _create_muon_optimizer( - config: OptimizerConfig, - named_params: list[tuple[str, nn.Parameter]], - parallel_dims: ParallelDims, - model: nn.Module | None, - lr: float | None = None, - param_groups: list[dict] | None = None, -) -> Optimizer: - if param_groups is None: - param_groups = _default_muon_param_groups(config, named_params, parallel_dims, lr) if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled: distributed_mesh = parallel_dims.get_mesh("dp_shard_cp") @@ -252,6 +237,6 @@ def _create_muon_optimizer( # Keep both warm-ups after Muon construction and before its first step. The # main and expert groups establish independent NCCL peer connections. _warmup_muon_mesh(distributed_mesh) - if any(group.get("distributed_mesh_name") == "dp_shard_mod_ep" for group in param_groups): + if expert_params and parallel_dims.ep_enabled: _warmup_muon_mesh(parallel_dims.get_mesh("dp_shard_mod_ep")) return optimizer diff --git a/src/prime_rl/trainer/pretrain/balance.py b/src/prime_rl/trainer/pretrain/balance.py deleted file mode 100644 index a999294495b..00000000000 --- a/src/prime_rl/trainer/pretrain/balance.py +++ /dev/null @@ -1,81 +0,0 @@ -import torch -import torch.distributed as dist -from torch import Tensor, nn - -from prime_rl.configs.pretrain import MoEBalanceConfig -from prime_rl.trainer.model import is_tt_moe_model, iter_moe_blocks -from prime_rl.trainer.models.layers.moe import MoE -from prime_rl.trainer.models.prime_moe.moe import histogram_quantile_biases - - -def setup_moe_balance(model: nn.Module, balance: MoEBalanceConfig, micro_batch_size: int) -> list[MoE]: - """Prepare the MoE layers' load-balancing state for training. Returns the layers with the sequence-wise balance - loss.""" - if is_tt_moe_model(model) and model.config.model_type != "prime_moe": - raise ValueError(f"Pretraining MoE models use model_type 'prime_moe', got {model.config.model_type!r}") - for moe in iter_moe_blocks(model): - # Allocated uninitialized with the rest of the model. - if moe.routed_k_counts is not None: - moe.routed_k_counts.zero_() - moe.router.quantile_sum.zero_() - moe.router.quantile_count.zero_() - moe.router.quantile_balance = balance.bias_update == "quantile" - if moe.router.quantile_balance and balance.quantile_estimator == "histogram": - device = moe.router.selection_bias.device - moe.router.quantile_hist = torch.zeros(moe.router.num_experts, balance.quantile_bins, device=device) - # The first step's grid; later steps use the previous step's global range. - moe.router.quantile_grid = torch.tensor([-1.0, 1.0], device=device) - moe.router.quantile_range = torch.tensor([float("inf"), float("-inf")], device=device) - moes_with_seq_balance = list(iter_moe_blocks(model)) if balance.seq_aux_coeff > 0 else [] - for moe in moes_with_seq_balance: - # A micro-batch is `micro_batch_size` packed rows concatenated into one sample; the sequences are the rows. - moe.balance_rows = micro_batch_size - return moes_with_seq_balance - - -@torch.no_grad() -def update_selection_biases(model: nn.Module, tokens_per_expert: Tensor, rate: float) -> None: - """Aux-loss-free load balancing (https://arxiv.org/abs/2408.15664): move each layer's selection bias by - ``rate`` towards underloaded experts, given the step's global ``[num_moe_layers, num_experts]`` counts. - The update is mean-centred (as in Moonlight and torchtitan), which leaves the selection unchanged but keeps - the biases from drifting together.""" - for moe, load in zip(iter_moe_blocks(model), tokens_per_expert): - if moe.router.selection_bias is not None: - delta = torch.sign(load.mean() - load) - moe.router.selection_bias += rate * (delta - delta.mean()) - - -@torch.no_grad() -def quantile_update_selection_biases(model: nn.Module, group: dist.ProcessGroup) -> None: - """Quantile balancing: set each layer's selection bias to minus the step's balancing thresholds of its experts - (``PrimeRouter.record_quantiles``), averaged over the micro-batches and ranks and mean-centred.""" - routers = [moe.router for moe in iter_moe_blocks(model)] - sums = torch.stack([torch.cat([router.quantile_sum, router.quantile_count[None]]) for router in routers]) - dist.all_reduce(sums, group=group) - for router, row in zip(routers, sums): - thresholds = row[:-1] / row[-1] - router.selection_bias.copy_(thresholds.mean() - thresholds) - router.quantile_sum.zero_() - router.quantile_count.zero_() - - -@torch.no_grad() -def histogram_quantile_update_selection_biases(model: nn.Module, group: dist.ProcessGroup) -> dict[str, float]: - """Quantile balancing with the histogram estimator: sum the step's histograms over the ranks, set each layer's - selection bias from them (``histogram_quantile_biases``) and use the step's global range of the required biases - as the next step's grid. Returns the grid's mean bin width and the share of values in its edge bins, which also - hold the values outside it.""" - routers = [moe.router for moe in iter_moe_blocks(model)] - hists = torch.stack([router.quantile_hist for router in routers]) - dist.all_reduce(hists, group=group) - ranges = torch.stack([torch.stack([-router.quantile_range[0], router.quantile_range[1]]) for router in routers]) - dist.all_reduce(ranges, op=dist.ReduceOp.MAX, group=group) - bin_width = clamped = 0.0 - for router, hist, (neg_lo, hi) in zip(routers, hists, ranges): - router.selection_bias.copy_(histogram_quantile_biases(hist, router.quantile_grid, router.top_k)) - bin_width += ((router.quantile_grid[1] - router.quantile_grid[0]) / hist.shape[1]).item() - clamped += ((hist[:, 0].sum() + hist[:, -1].sum()) / hist.sum()).item() - router.quantile_grid = torch.stack([-neg_lo, torch.maximum(hi, -neg_lo + 1e-6)]) - router.quantile_hist.zero_() - router.quantile_range.copy_(torch.tensor([float("inf"), float("-inf")])) - return {"router/qb/bin_width": bin_width / len(routers), "router/qb/edge_share": clamped / len(routers)} diff --git a/src/prime_rl/trainer/pretrain/optim.py b/src/prime_rl/trainer/pretrain/optim.py deleted file mode 100644 index 53802a87dff..00000000000 --- a/src/prime_rl/trainer/pretrain/optim.py +++ /dev/null @@ -1,22 +0,0 @@ -from torch import nn - -from prime_rl.configs.trainer import AdamWConfig, MuonConfig -from prime_rl.trainer.optim import _default_muon_param_groups -from prime_rl.trainer.parallel_dims import ParallelDims - - -def param_groups(model: nn.Module, optim: AdamWConfig | MuonConfig, parallel_dims: ParallelDims) -> list[dict] | None: - """The default Muon groups, with the Gated Residual projections (``hyper_connection``) moved to the AdamW group. - None (the default groups) without Muon or Gated Residual. - - These projections produce sigmoid gate logits (the write projection only one scalar per stream), not hidden - activations. Muon's orthogonalized update moves every direction by the same amount, which makes the gates noisy: - on Muon they cost +0.028 held-out loss at d384 (4 seeds).""" - named_params = list(model.named_parameters()) - hyper_connection = [param for name, param in named_params if "hyper_connection" in name and param.requires_grad] - if optim.type != "muon" or not hyper_connection: - return None - rest = [(name, param) for name, param in named_params if "hyper_connection" not in name] - groups = _default_muon_param_groups(optim, rest, parallel_dims, optim.lr) - next(group for group in groups if group["algorithm"] == "adamw")["params"].extend(hyper_connection) - return groups diff --git a/src/prime_rl/trainer/scheduler.py b/src/prime_rl/trainer/scheduler.py index 7fb18fab87f..25bfb93f6c3 100644 --- a/src/prime_rl/trainer/scheduler.py +++ b/src/prime_rl/trainer/scheduler.py @@ -1,12 +1,9 @@ from __future__ import annotations -import math -from collections.abc import Callable - from torch.optim import Optimizer -from torch.optim.lr_scheduler import ConstantLR, CosineAnnealingLR, LambdaLR, LinearLR, LRScheduler, SequentialLR +from torch.optim.lr_scheduler import ConstantLR, CosineAnnealingLR, LinearLR, LRScheduler, SequentialLR -from prime_rl.configs.trainer import SchedulerConfig, resolve_steps +from prime_rl.configs.trainer import SchedulerConfig from prime_rl.trainer.optim import OffloadOptimizer, OptimizerLike @@ -22,20 +19,8 @@ def setup_constant_scheduler(optimizer: Optimizer) -> LRScheduler: return ConstantLR(optimizer, factor=1.0) -DECAY_SHAPES: dict[str, Callable[[float], float]] = { - "cosine": lambda progress: 0.5 * (1 + math.cos(math.pi * progress)), - "1-sqrt": lambda progress: 1 - math.sqrt(progress), -} - - def setup_linear_scheduler( - optimizer: Optimizer, - max_steps: int | None, - warmup_steps: int, - decay_steps: int, - lr: float, - min_lr: float, - decay_shape: str = "linear", + optimizer: Optimizer, max_steps: int | None, warmup_steps: int, decay_steps: int, lr: float, min_lr: float ) -> LRScheduler: """Create a linear (WSD) learning rate scheduler.""" # Create schedulers for each phase @@ -61,16 +46,7 @@ def setup_linear_scheduler( ) constant_steps = decay_start_step - warmup_steps constant_scheduler = LinearLR(optimizer, start_factor=1.0, end_factor=1.0, total_iters=constant_steps) - if decay_shape == "linear": - decay_scheduler = LinearLR( - optimizer, start_factor=1.0, end_factor=min_lr_factor, total_iters=decay_steps - 1 - ) - else: - shape = DECAY_SHAPES[decay_shape] - decay_scheduler = LambdaLR( - optimizer, - lambda step: min_lr_factor + (1 - min_lr_factor) * shape(min(step / max(decay_steps - 1, 1), 1.0)), - ) + decay_scheduler = LinearLR(optimizer, start_factor=1.0, end_factor=min_lr_factor, total_iters=decay_steps - 1) schedulers.append(constant_scheduler) schedulers.append(decay_scheduler) milestones.append(decay_start_step) @@ -128,11 +104,10 @@ def setup_scheduler( return setup_linear_scheduler( base_optimizer, max_steps=max_steps, - warmup_steps=resolve_steps(scheduler_config.warmup_steps, max_steps), - decay_steps=resolve_steps(scheduler_config.decay_steps, max_steps), + warmup_steps=scheduler_config.warmup_steps, + decay_steps=scheduler_config.decay_steps, lr=lr, min_lr=scheduler_config.min_lr, - decay_shape=scheduler_config.decay_shape, ) case "cosine": return setup_cosine_scheduler( diff --git a/tests/integration/test_pretrain_moe.py b/tests/integration/test_pretrain_moe.py deleted file mode 100644 index 4d6f023d27e..00000000000 --- a/tests/integration/test_pretrain_moe.py +++ /dev/null @@ -1,106 +0,0 @@ -from pathlib import Path -from typing import Callable - -import numpy as np -import pytest -from datasets import load_dataset -from transformers import AutoTokenizer - -from tests.conftest import ProcessResult -from tests.utils import check_loss_goes_down, strip_escape_codes - -pytestmark = [pytest.mark.slow, pytest.mark.gpu] - -RUN_NAME = "pretrain-moe" -TIMEOUT = 300 # 5 minutes - - -@pytest.fixture(scope="module") -def run_dir(output_dir: Path) -> Path: - return output_dir / RUN_NAME - - -@pytest.fixture(scope="module") -def wandb_name(branch_name: str) -> str: - return f"test-pretrain-moe:{branch_name}" - - -@pytest.fixture(scope="module") -def token_dir(output_dir: Path) -> Path: - """Reverse-Text-SFT as one datatrove token shard: every conversation one document, ending with EOS.""" - tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") - documents = [ - tokenizer.encode("\n".join(message["content"] for message in row["prompt"] + row["completion"])) - + [tokenizer.eos_token_id] - for row in load_dataset("PrimeIntellect/Reverse-Text-SFT", split="train") - ] - path = output_dir / "pretrain-tokens" - path.mkdir(parents=True, exist_ok=True) - np.array([token for document in documents for token in document], dtype=" list[str]: - return [ - "uv", - "run", - "pretrain", - "@", - f"configs/ci/integration/pretrain-moe/{config}", - "--data.paths", - token_dir.as_posix(), - "--deployment.num-train-gpus", - "2", - "--monitors.wandb.project", - wandb_project, - "--monitors.wandb.name", - wandb_name, - "--output-dir", - output_dir.as_posix(), - "--run.name", - RUN_NAME, - ] - - -@pytest.fixture(scope="module") -def pretrain_process( - run_process: Callable[..., ProcessResult], token_dir: Path, wandb_project: str, wandb_name: str, output_dir: Path -) -> ProcessResult: - cmd = pretrain_command("start.toml", token_dir, wandb_project, wandb_name, output_dir) - return run_process([*cmd, "--clean"], timeout=TIMEOUT) - - -@pytest.fixture(scope="module") -def pretrain_resume_process( - pretrain_process: ProcessResult, - run_process: Callable[..., ProcessResult], - token_dir: Path, - wandb_project: str, - wandb_name: str, - output_dir: Path, -) -> ProcessResult: - cmd = pretrain_command("resume.toml", token_dir, wandb_project, f"{wandb_name}-resume", output_dir) - return run_process(cmd, timeout=TIMEOUT) - - -def trainer_log(run_dir: Path) -> list[str]: - with open(run_dir / "logs" / "latest" / "trainer.log") as f: - return strip_escape_codes(f.read()).splitlines() - - -def test_no_error(pretrain_process: ProcessResult): - assert pretrain_process.returncode == 0, f"Process has non-zero return code ({pretrain_process})" - - -def test_loss_goes_down(pretrain_process: ProcessResult, run_dir: Path): - check_loss_goes_down(trainer_log(run_dir)) - - -def test_no_error_resume(pretrain_resume_process: ProcessResult): - assert pretrain_resume_process.returncode == 0, f"Process has non-zero return code ({pretrain_resume_process})" - - -def test_loss_goes_down_resume(pretrain_resume_process: ProcessResult, run_dir: Path): - check_loss_goes_down(trainer_log(run_dir)) diff --git a/tests/unit/train/models/test_init.py b/tests/unit/train/models/test_init.py new file mode 100644 index 00000000000..6ae4cd0a252 --- /dev/null +++ b/tests/unit/train/models/test_init.py @@ -0,0 +1,21 @@ +import itertools + +import torch + +from prime_rl.trainer.models.init import normal_shard + + +def test_normal_shard_is_independent_of_sharding(): + shape = (37, 5, 3) + full = normal_shard("w", shape, shape, (0, 0, 0), seed=0, device=torch.device("cpu")) + for parts in (2, 3, 8): + bounds = [len(chunk) for chunk in torch.arange(shape[0]).chunk(parts)] + offsets = [0, *itertools.accumulate(bounds)][:-1] + shards = [ + normal_shard("w", shape, (rows, *shape[1:]), (offset, 0, 0), seed=0, device=torch.device("cpu")) + for rows, offset in zip(bounds, offsets) + ] + assert torch.equal(torch.cat(shards), full) + column_shard = normal_shard("w", shape, (37, 2, 3), (0, 3, 0), seed=0, device=torch.device("cpu")) + assert torch.equal(column_shard, full[:, 3:5]) + assert not torch.equal(normal_shard("w", shape, shape, (0, 0, 0), seed=1, device=torch.device("cpu")), full) diff --git a/tests/unit/train/pretrain/test_init.py b/tests/unit/train/pretrain/test_init.py deleted file mode 100644 index 77d90922070..00000000000 --- a/tests/unit/train/pretrain/test_init.py +++ /dev/null @@ -1,68 +0,0 @@ -import itertools -from types import SimpleNamespace - -import torch - -from prime_rl.configs.trainer import MuonConfig -from prime_rl.trainer.models.init import normal_shard -from prime_rl.trainer.models.prime_moe import PrimeMoeConfig, PrimeMoeForCausalLM -from prime_rl.trainer.pretrain.optim import param_groups -from prime_rl.trainer.utils import count_parameters - - -def test_normal_shard_is_independent_of_sharding(): - shape = (37, 5, 3) - full = normal_shard("w", shape, shape, (0, 0, 0), seed=0, device=torch.device("cpu")) - for parts in (2, 3, 8): - bounds = [len(chunk) for chunk in torch.arange(shape[0]).chunk(parts)] - offsets = [0, *itertools.accumulate(bounds)][:-1] - shards = [ - normal_shard("w", shape, (rows, *shape[1:]), (offset, 0, 0), seed=0, device=torch.device("cpu")) - for rows, offset in zip(bounds, offsets) - ] - assert torch.equal(torch.cat(shards), full) - column_shard = normal_shard("w", shape, (37, 2, 3), (0, 3, 0), seed=0, device=torch.device("cpu")) - assert torch.equal(column_shard, full[:, 3:5]) - assert not torch.equal(normal_shard("w", shape, shape, (0, 0, 0), seed=1, device=torch.device("cpu")), full) - - -def small_moe(**kwargs) -> PrimeMoeForCausalLM: - config = PrimeMoeConfig( - vocab_size=128, - hidden_size=512, - intermediate_size=1024, - moe_intermediate_size=128, - num_attention_heads=8, - num_key_value_heads=2, - n_routed_experts=8, - num_experts_per_tok=2, - n_shared_experts=1, - num_hidden_layers=2, - first_k_dense_replace=1, - **kwargs, - ) - config._attn_implementation = "flash_attention_2" - with torch.device("meta"): - return PrimeMoeForCausalLM._from_config(config) - - -def test_count_parameters_moe(): - model = small_moe() - counts = count_parameters(model) - embedding = model.model.embed_tokens.weight.numel() + model.lm_head.weight.numel() - experts = sum(param.numel() for param in model.model.layers[1].mlp.experts.parameters()) - assert counts["total"] == sum(param.numel() for param in model.parameters()) - assert counts["embedding"] == embedding - assert counts["non_embedding"] == counts["total"] - embedding - assert counts["active_non_embedding"] == counts["non_embedding"] - experts * 6 // 8 - - -def test_gated_residual_projections_train_with_adamw(): - model = small_moe(hc_count=4, hc_lowrank=32) - groups = param_groups(model, MuonConfig(lr=1.0), SimpleNamespace(ep_enabled=False)) - algorithm = {param: group["algorithm"] for group in groups for param in group["params"]} - assert sum(len(group["params"]) for group in groups) == len(algorithm) == len(list(model.parameters())) - hyper_connection = [param for name, param in model.named_parameters() if "hyper_connection" in name] - assert hyper_connection and all(algorithm[param] == "adamw" for param in hyper_connection) - assert algorithm[model.model.layers[0].self_attn.q_proj.weight] == "muon" - assert param_groups(small_moe(), MuonConfig(lr=1.0), SimpleNamespace(ep_enabled=False)) is None diff --git a/tests/unit/train/pretrain/test_moe_balance.py b/tests/unit/train/pretrain/test_moe_balance.py deleted file mode 100644 index a2cdf664d84..00000000000 --- a/tests/unit/train/pretrain/test_moe_balance.py +++ /dev/null @@ -1,82 +0,0 @@ -import torch - -from prime_rl.trainer.models.layers.moe import MoE, MoEArgs -from prime_rl.trainer.models.prime_moe.moe import PrimeRouter, histogram_quantile_biases, sequence_balance -from prime_rl.trainer.pretrain.train import update_selection_biases - - -def test_update_selection_biases_moves_towards_underloaded_experts(): - moe = MoE.from_args(MoEArgs(num_experts=4), dim=4, hidden_dim=8, shared_expert=None) - layer = torch.nn.Module() - layer.mlp = moe - model = torch.nn.Module() - model.model = torch.nn.Module() - model.model.layers = torch.nn.ModuleList([layer]) - - update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]]), rate=0.1) - update_selection_biases(model, torch.tensor([[9.0, 1.0, 1.0, 1.0]]), rate=0.1) - - # sign(mean - load) = [-1, 1, 1, 1], centred to [-1.5, 0.5, 0.5, 0.5] per step - torch.testing.assert_close(moe.router.selection_bias, torch.tensor([-0.3, 0.1, 0.1, 0.1])) - - -def test_sequence_balance(): - scores = torch.tensor([[0.9, 0.1], [0.9, 0.1], [0.9, 0.1], [0.1, 0.9]]) - selected = torch.tensor([[0], [0], [0], [1]]) - - # row 0 sends both tokens to expert 0: f = [2, 0], P = [0.9, 0.1]; row 1 is balanced: f = [1, 1], P = [0.5, 0.5] - torch.testing.assert_close(sequence_balance(scores, selected, num_rows=2), torch.tensor([1.8, 1.0])) - # the same tokens as one sequence (a packed sample of both rows) hide row 0's imbalance: f = [1.5, 0.5], P = [0.7, 0.3] - torch.testing.assert_close(sequence_balance(scores, selected, num_rows=1), torch.tensor([1.2])) - - -def test_quantile_balancing_balances_the_recorded_batch(): - torch.manual_seed(0) - router = PrimeRouter( - dim=16, num_experts=6, num_null_experts=2, top_k=2, route_norm=True, route_scale=1.0, bias_on_logits=True - ) - with torch.no_grad(): - router.gate.weight[0] += 0.5 # expert 0 is preferred - router.quantile_balance = True - x = torch.randn(400, 16) - counts = router(x)[2] - assert counts.max() / counts.float().mean() > 1.4 - - # One quantile update brings every expert, null experts included, close to 400 * 2 / 8 = 100 tokens. - thresholds = router.quantile_sum / router.quantile_count - router.selection_bias.copy_(thresholds.mean() - thresholds) - with torch.no_grad(): - counts = router(x)[2] - assert counts.max() / counts.float().mean() < 1.15 - - -def test_histogram_quantile_balancing_balances_the_recorded_batch_however_it_is_split(): - torch.manual_seed(0) - router = PrimeRouter( - dim=16, num_experts=6, num_null_experts=2, top_k=2, route_norm=True, route_scale=1.0, bias_on_logits=True - ) - with torch.no_grad(): - router.gate.weight[0] += 0.5 # expert 0 is preferred - router.quantile_balance = True - router.quantile_grid = torch.tensor([-4.0, 4.0]) - x = torch.randn(400, 16) - hists = [] - for chunks in (1, 4): - router.quantile_hist = torch.zeros(8, 1024) - router.quantile_range = torch.tensor([float("inf"), float("-inf")]) - for part in x.chunk(chunks): - router(part) - hists.append(router.quantile_hist) - torch.testing.assert_close(hists[0], hists[1]) - # Each token counts once per expert, in the bin of its required bias. - biased = router.gate(x) + router.selection_bias - required = biased.topk(router.top_k + 1, dim=1).values[:, -1:] - router.gate(x) - bins = ((required + 4.0) * (1024 / 8.0)).floor().clamp(0, 1023).long() - expected = torch.stack([torch.bincount(bins[:, j], minlength=1024) for j in range(8)]).float() - torch.testing.assert_close(hists[0], expected) - - # One update brings every expert, null experts included, close to 400 * 2 / 8 = 100 tokens. - router.selection_bias.copy_(histogram_quantile_biases(hists[0], router.quantile_grid, router.top_k)) - with torch.no_grad(): - counts = router(x)[2] - assert counts.max() / counts.float().mean() < 1.15 diff --git a/tests/unit/train/test_scheduler.py b/tests/unit/train/test_scheduler.py deleted file mode 100644 index 57df2dd3c46..00000000000 --- a/tests/unit/train/test_scheduler.py +++ /dev/null @@ -1,28 +0,0 @@ -import pytest -import torch - -from prime_rl.configs.trainer import LinearSchedulerConfig -from prime_rl.trainer.scheduler import setup_scheduler - - -def lrs(config: LinearSchedulerConfig, max_steps: int) -> list[list[float]]: - params = [torch.nn.Parameter(torch.zeros(1)), torch.nn.Parameter(torch.zeros(1))] - optimizer = torch.optim.AdamW([{"params": params[:1], "lr": 1.0}, {"params": params[1:], "lr": 0.5}]) - scheduler = setup_scheduler(optimizer, config, max_steps, lr=1.0) - values = [] - for _ in range(max_steps): - values.append([group["lr"] for group in optimizer.param_groups]) - optimizer.step() - scheduler.step() - return values - - -def test_linear_scheduler_fractions_and_decay_shapes(): - linear = lrs(LinearSchedulerConfig(warmup_steps=0.1, decay_steps=0.9), max_steps=20) - assert linear[2] == [1.0, 0.5] and linear[-1][0] < 1e-7 - shaped = lrs(LinearSchedulerConfig(warmup_steps=2, decay_steps=0.5, decay_shape="1-sqrt"), max_steps=20) - # The decay runs over the last 10 updates: 1 at its start, 1 - sqrt(5/9) five updates in, ~0 at the end. - assert shaped[10] == [1.0, 0.5] - assert shaped[15][0] == pytest.approx(1 - (5 / 9) ** 0.5, abs=1e-6) - assert shaped[15][1] == pytest.approx(shaped[15][0] / 2) - assert shaped[-1][0] < 1e-7 From 060d730f928487709a02a14a27e3b678f204b1a0 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 22 Oct 2026 23:11:09 +0200 Subject: [PATCH 125/132] refactor(pretrain): pretraining is the SFT trainer with token data The separate pretrain trainer, config, launcher and entrypoint are gone. Pretraining runs with `uv run sft`: - data.type = "tokens": datatrove token shards as SFT data (TokenDataset); CatDataset keeps every document of a window as its own segment - model.init: initialize the weights from scratch instead of loading them - [ll_eval]: the log-likelihood evals, scored inside the trainer - the micro-batch loop moves to trainer/micro_batches.py: no host sync per micro-batch, the step statistics are reduced once, with the same gradients and metric names as before; model.defer_grad_sync reduce-scatters the gradients once per optimizer step - extra metrics for every run: params/*, loss/spike, optim/grad_norm/, optim/grad_norm_spike, time/data/*, expert_load/{min_mean,entropy,dead,layer_N/*} - parameter counts are methods of PreTrainedModelPrimeRL, built on one embedding rule --- .github/workflows/tests.yml | 11 +- configs/debug/fake/pretrain-model/config.json | 28 -- configs/debug/fake/pretrain.toml | 17 - .../src/prime_rl/configs/pretrain.py | 221 --------- .../src/prime_rl/configs/sft.py | 273 ++++++++++- .../src/prime_rl/configs/shared.py | 42 +- .../src/prime_rl/configs/trainer.py | 103 +---- pyproject.toml | 1 - src/prime_rl/dashboard/server.py | 4 +- src/prime_rl/dashboard/static/app.js | 12 - src/prime_rl/entrypoints/pretrain.py | 94 ---- src/prime_rl/entrypoints/sft.py | 170 +++---- src/prime_rl/monitors/__init__.py | 2 +- src/prime_rl/monitors/prime.py | 5 +- src/prime_rl/monitors/wandb/monitor.py | 2 +- src/prime_rl/monitors/wandb/overview.py | 38 +- .../templates/multi_node_sft.sbatch.j2 | 4 +- .../templates/single_node_sft.sbatch.j2 | 2 +- .../trainer/{pretrain => }/metrics.py | 34 +- src/prime_rl/trainer/micro_batches.py | 140 ++++++ src/prime_rl/trainer/model.py | 21 +- src/prime_rl/trainer/models/base.py | 18 + src/prime_rl/trainer/pretrain/bytes.py | 33 -- src/prime_rl/trainer/pretrain/data.py | 254 ----------- src/prime_rl/trainer/pretrain/step.py | 123 ----- src/prime_rl/trainer/pretrain/train.py | 430 ------------------ src/prime_rl/trainer/rl/train.py | 11 +- src/prime_rl/trainer/sft/data.py | 157 ++++++- .../{pretrain/suite.py => sft/ll_eval.py} | 86 +++- src/prime_rl/trainer/sft/train.py | 201 ++++---- src/prime_rl/trainer/utils.py | 28 -- src/prime_rl/utils/pathing.py | 14 +- tests/unit/test_configs.py | 2 - tests/unit/train/pretrain/test_bytes.py | 12 - tests/unit/train/pretrain/test_data.py | 79 ---- tests/unit/train/pretrain/test_launcher.py | 33 -- .../train/pretrain/test_pretrain_metrics.py | 34 -- tests/unit/train/pretrain/test_suite.py | 161 ------- tests/unit/train/sft/test_ll_eval.py | 63 +++ tests/unit/train/sft/test_token_dataset.py | 38 ++ 40 files changed, 978 insertions(+), 2023 deletions(-) delete mode 100644 configs/debug/fake/pretrain-model/config.json delete mode 100644 configs/debug/fake/pretrain.toml delete mode 100644 packages/prime-rl-configs/src/prime_rl/configs/pretrain.py delete mode 100644 src/prime_rl/entrypoints/pretrain.py rename src/prime_rl/trainer/{pretrain => }/metrics.py (79%) create mode 100644 src/prime_rl/trainer/micro_batches.py delete mode 100644 src/prime_rl/trainer/pretrain/bytes.py delete mode 100644 src/prime_rl/trainer/pretrain/data.py delete mode 100644 src/prime_rl/trainer/pretrain/step.py delete mode 100644 src/prime_rl/trainer/pretrain/train.py rename src/prime_rl/trainer/{pretrain/suite.py => sft/ll_eval.py} (80%) delete mode 100644 tests/unit/train/pretrain/test_bytes.py delete mode 100644 tests/unit/train/pretrain/test_data.py delete mode 100644 tests/unit/train/pretrain/test_launcher.py delete mode 100644 tests/unit/train/pretrain/test_pretrain_metrics.py delete mode 100644 tests/unit/train/pretrain/test_suite.py create mode 100644 tests/unit/train/sft/test_ll_eval.py create mode 100644 tests/unit/train/sft/test_token_dataset.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 3adf9a69988..6e90c6aa2db 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -170,13 +170,6 @@ jobs: - 'tests/integration/test_alphabet_sort.py' - 'configs/ci/integration/reverse-text{,-lora,-moe,-rl-opd,-rl-sft}/**' - 'configs/ci/integration/alphabet_sort.toml' - pretrain: - - *runtime - - 'packages/prime-rl-configs/src/prime_rl/configs/{pretrain,monitors}.py' - - 'src/prime_rl/entrypoints/pretrain.py' - - 'src/prime_rl/monitors/**' - - 'tests/integration/test_pretrain_moe.py' - - 'configs/ci/integration/pretrain-moe/**' eval: - *runtime - 'packages/prime-rl-configs/src/prime_rl/configs/{eval,orchestrator,inference,env_server,algorithm,monitors}.py' @@ -191,14 +184,12 @@ jobs: RUN_ALL: ${{ github.event_name != 'pull_request' }} SFT: ${{ steps.filter.outputs.sft || 'false' }} RL: ${{ steps.filter.outputs.rl || 'false' }} - PRETRAIN: ${{ steps.filter.outputs.pretrain || 'false' }} run: | - tests=$(jq -cn --argjson all "$RUN_ALL" --argjson sft "$SFT" --argjson rl "$RL" --argjson pretrain "$PRETRAIN" '[ + tests=$(jq -cn --argjson all "$RUN_ALL" --argjson sft "$SFT" --argjson rl "$RL" '[ if $all or $sft then {test: "reverse_text_sft", name: "Reverse Text SFT"}, {test: "reverse_text_sft_lora", name: "Reverse Text SFT LoRA"} else empty end, - if $all or $pretrain then {test: "pretrain_moe", name: "Pretrain MoE"} else empty end, if $all or $rl then {test: "reverse_text", name: "Reverse Text"}, {test: "reverse_text_rl_opd", name: "Reverse Text RL OPD"}, diff --git a/configs/debug/fake/pretrain-model/config.json b/configs/debug/fake/pretrain-model/config.json deleted file mode 100644 index e991ed03b69..00000000000 --- a/configs/debug/fake/pretrain-model/config.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "attention_bias": false, - "attention_dropout": 0.0, - "bos_token_id": null, - "eos_token_id": 151645, - "head_dim": 64, - "hidden_act": "silu", - "hidden_size": 256, - "initializer_range": 0.02, - "intermediate_size": 1024, - "max_position_embeddings": 2048, - "mlp_bias": false, - "model_type": "llama", - "num_attention_heads": 4, - "num_hidden_layers": 2, - "num_key_value_heads": 2, - "pad_token_id": 151643, - "pretraining_tp": 1, - "rms_norm_eps": 1e-06, - "rope_parameters": { - "rope_theta": 10000.0, - "rope_type": "default" - }, - "tie_word_embeddings": false, - "transformers_version": "5.6.2", - "use_cache": true, - "vocab_size": 151680 -} diff --git a/configs/debug/fake/pretrain.toml b/configs/debug/fake/pretrain.toml deleted file mode 100644 index 1c778cd9a6e..00000000000 --- a/configs/debug/fake/pretrain.toml +++ /dev/null @@ -1,17 +0,0 @@ -max_steps = 5 - -[model] -name = "configs/debug/fake/pretrain-model" -compile = "None" - -[tokenizer] -name = "PrimeIntellect/Qwen3-0.6B" - -[data] -type = "fake" -seq_len = 512 -batch_size = 4 -micro_batch_size = 1 - -[optim] -lr = 1e-3 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py b/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py deleted file mode 100644 index 2106d89041b..00000000000 --- a/packages/prime-rl-configs/src/prime_rl/configs/pretrain.py +++ /dev/null @@ -1,221 +0,0 @@ -import uuid -from pathlib import Path -from typing import Annotated, Literal, TypeAlias - -from pydantic import Field, model_validator - -from prime_rl.configs.sft import BaseDataConfig -from prime_rl.configs.trainer import ( - AdamWConfig, - CheckpointConfig, - InitConfig, - LinearSchedulerConfig, - ModelConfig, - MoERuntimeConfig, - MuonConfig, - SchedulerConfig, - TrainerLaunchConfig, -) -from prime_rl.utils.config import BaseConfig - - -class MoEBalanceConfig(BaseConfig): - """MoE load balancing as in DeepSeek-V3: an aux-loss-free selection bias plus a small sequence-wise balance - loss.""" - - bias_update: Literal["sign", "quantile"] = "quantile" - """``quantile`` (default): quantile balancing (Su 2026; Kimi K3, Marin): every step, each layer's bias is set to - the value with which every expert would have been selected by ``top_k / num_experts`` of the step's tokens (null - experts included), estimated as ``quantile_estimator`` sets; no rate. Marin found it - works best on the router logits (``router_bias_on_logits`` in the model config, on by default). ``sign``: every - step, each layer's bias moves by ``bias_update_rate`` towards underloaded experts.""" - - quantile_estimator: Literal["histogram", "micro_batch_mean"] = "histogram" - """How ``quantile`` estimates the thresholds. ``histogram`` (default; Kimi K3 report, §2.3.3): per expert, a - histogram of the step's margins summed over the micro-batches and ranks (one all-reduce per step), read with in-bin - interpolation; the quantile of the pooled step, independent of how the tokens are split. Its grid is the previous - step's global margin range (``[-1, 1]`` on the first step after a start or resume). ``micro_batch_mean``: each - micro-batch's exact quantile, averaged over the micro-batches and ranks.""" - - quantile_bins: int = Field(1024, ge=2) - """Bins of the ``histogram`` estimator's grid.""" - - bias_update_rate: float = Field(3e-3, ge=0) - """Step size of the sign update (gamma): every step, each layer's bias moves by this towards underloaded experts, - mean-centred. 0 turns the update off.""" - - bias_freeze: float = Field(0.1, ge=0, lt=1) - """Fraction of ``max_steps`` at the end of training during which the biases stop updating, as DeepSeek-V3, - GLM-4.5 and Moonlight do in their final stage.""" - - seq_aux_coeff: float = Field(1e-3, ge=0) - """Weight of the sequence-wise balance loss (alpha), computed per packed row. 0 turns it off.""" - - -class PretrainMoEConfig(MoERuntimeConfig): - balance: MoEBalanceConfig = MoEBalanceConfig() - """Load balancing of the routed experts.""" - - -class PretrainModelConfig(ModelConfig): - init: InitConfig = InitConfig() - """Random weight initialization. Used unless the run resumes from a checkpoint.""" - - moe: PretrainMoEConfig = PretrainMoEConfig() - - offload: Literal["none", "optimizer"] = "none" - """``optimizer`` keeps the optimizer states in CPU RAM between steps, as in SFT and RL. ``full`` (CPU optimizer - steps during backward) is not supported by the pretraining step.""" - - -class BaseBatchConfig(BaseDataConfig): - seq_len: int = Field(8192, ge=1) - """Tokens per row. Documents continue across rows; attention restarts at every document and every row.""" - - seed: int = 0 - """Seed for the per-epoch window order of the pool (and the fake tokens).""" - - -class FakeDataConfig(BaseBatchConfig): - """Uniformly random token ids, for smoke tests.""" - - type: Literal["fake"] = "fake" - - -class TokensDataConfig(BaseBatchConfig): - """Rows of pre-tokenized datatrove shards: random windows of one token pool. Mixtures are made offline, as shard - folders.""" - - type: Literal["tokens"] = "tokens" - - paths: list[Path] = Field(min_length=1) - """Folders of datatrove tokenized shards: ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``. - The files of all folders (in folder order, then file name order) are read as one token pool.""" - - max_epochs: float = Field(1.0, gt=0) - """Most passes over the pool that the run (``max_steps``) may make. Checked at startup.""" - - document_masking: bool = True - """Attention and positions restart at every document. If False, a row is one sequence: attention spans the - whole row and positions run over it.""" - - -PretrainDataConfig: TypeAlias = Annotated[FakeDataConfig | TokensDataConfig, Field(discriminator="type")] - - -class TokenEvalSetConfig(BaseConfig): - """A held-out folder of datatrove tokenized shards, scored like a ``ppl`` set: every document on its own (in - windows of each of ``seq_lens``), all documents of the folder.""" - - name: str - """Names the metrics, ``eval//@``.""" - - path: Path - """Folder of ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``, as for training.""" - - -class EvalsConfig(BaseConfig): - """The evals, scored inside the trainer: forward passes without gradients on the training model, every - ``interval`` steps. The sets are log-likelihood evals in the format of the pretrain-evals offline suite, and - held-out token shards.""" - - interval: int = Field(ge=1) - """Score the sets every N steps (after that step's optimizer update), and at the last step.""" - - eval_on_start: bool = True - """Also score the weights the run starts from (the initial weights, logged as step 0, or the checkpoint it - resumes from, logged at that step).""" - - data_dir: Path | None = None - """Directory of the pretrain-evals sets, one directory per set with ``test.parquet`` and ``meta.json``.""" - - sets: list[str | TokenEvalSetConfig] = Field(min_length=1) - """The sets to score: the name of a pretrain-evals set in ``data_dir``, or a held-out token folder.""" - - seq_lens: list[int] | None = None - """Context lengths to score every set at: the window length for ``ppl`` and token sets and the context limit for - ``ctxcont`` sets. Defaults to ``data.seq_len``. Metrics are named ``eval//@`` (e.g. - ``bpb@4k``). Requests that are the same at several lengths are scored once.""" - - micro_batch_tokens: int | None = Field(None, ge=1) - """Tokens per packed eval row (one forward pass on one GPU). Defaults to the training micro-batch - (``data.micro_batch_size * data.seq_len``) or the longest of ``seq_lens``, whichever is larger; must be at least - the longest of ``seq_lens``.""" - - save_arrays: bool = False - """Write per-document arrays (npz) and summaries to ``/evals/step_N/``.""" - - @model_validator(mode="after") - def validate_data_dir(self): - if self.data_dir is None and any(isinstance(eval_set, str) for eval_set in self.sets): - raise ValueError("evals.data_dir is required for pretrain-evals sets") - return self - - -PretrainOptimizerConfig: TypeAlias = Annotated[AdamWConfig | MuonConfig, Field(discriminator="type")] - - -class PretrainConfig(TrainerLaunchConfig): - # A trainable router, and weight decay 0.1. - auto_freeze_moe_router = False - auto_weight_decay = 0.1 - - model: PretrainModelConfig = PretrainModelConfig() - """The model. ``model.name`` is an HF model id or a local model directory with a ``config.json`` (as in SFT and - RL); its weights are not loaded: the model is initialized from scratch (``model.init``) or resumed.""" - - data: PretrainDataConfig - - evals: EvalsConfig | None = None - """The evals, scored inside the trainer. If None, they never run.""" - - optim: PretrainOptimizerConfig = MuonConfig() - """Optimizer. ``optim.lr`` is required.""" - - scheduler: SchedulerConfig = LinearSchedulerConfig(warmup_steps=0.1, decay_steps=0.9) - """Learning-rate schedule. Defaults to a warmup over 10% of the steps and a linear decay to 0 over the rest.""" - - ckpt: CheckpointConfig | None = CheckpointConfig() - """Checkpointing. Defaults to a checkpoint at the end of training.""" - - max_steps: int = Field(ge=1) - """Training steps.""" - - def default_run_name(self) -> str: - return f"{Path(self.model.name).name}--{uuid.uuid4().hex[:8]}" - - @model_validator(mode="after") - def validate_unsupported(self): - unsupported = [name for name in ("trace_path", "memory_profiler_path", "control") if getattr(self, name)] - if unsupported: - raise ValueError(f"Pretraining does not support {', '.join(unsupported)} yet") - return self - - @model_validator(mode="after") - def validate_lr(self): - if "lr" not in self.optim.model_fields_set: - raise ValueError("optim.lr is required") - return self - - @model_validator(mode="after") - def validate_cp(self): - if self.model.cp > 1: - if self.data.seq_len % self.model.cp != 0: - raise ValueError("data.seq_len must be divisible by model.cp") - if self.data.micro_batch_size != 1: - raise ValueError("data.micro_batch_size must be 1 with context parallelism") - if self.evals is not None: - raise ValueError("[evals] does not support context parallelism yet") - return self - - @model_validator(mode="after") - def auto_setup_evals(self): - if self.evals is not None: - if self.evals.seq_lens is None: - self.evals.seq_lens = [self.data.seq_len] - longest = max(self.evals.seq_lens) - if self.evals.micro_batch_tokens is None: - self.evals.micro_batch_tokens = max(self.data.micro_batch_size * self.data.seq_len, longest) - elif self.evals.micro_batch_tokens < longest: - raise ValueError(f"evals.micro_batch_tokens must be at least the longest of evals.seq_lens ({longest})") - return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 1e7fe2914f2..c58468a1d0b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -1,22 +1,27 @@ import uuid import warnings +from pathlib import Path from typing import Annotated, Literal, TypeAlias from urllib.parse import urlparse -from pydantic import Field, model_validator +from pydantic import AliasChoices, Field, model_validator from renderers import AutoRendererConfig, DefaultRendererConfig, RendererConfig from renderers.base import MODEL_RENDERER_MAP from prime_rl.configs.eval import SFTOnlineEvalConfig from prime_rl.configs.inference import InferenceConfig from prime_rl.configs.inference import WeightBroadcastConfig as InferenceWeightBroadcastConfig +from prime_rl.configs.monitors import TrainMonitorsConfig from prime_rl.configs.shared import ( FileSystemWeightBroadcastConfig, NCCLWeightBroadcastConfig, + ResumeConfig, + RunConfig, + SlurmConfig, WeightBroadcastConfig, ) -from prime_rl.configs.trainer import TrainerLaunchConfig, mega_moe_max_tokens_per_rank -from prime_rl.utils.config import BaseConfig +from prime_rl.configs.trainer import BaseTrainerConfig, InitConfig, ModelConfig, mega_moe_max_tokens_per_rank +from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource class BaseDataConfig(BaseConfig): @@ -32,6 +37,11 @@ class BaseDataConfig(BaseConfig): num_workers: int = Field(1, ge=1) """Number of dataloader worker processes. Batches are prepared and pinned in the background so tokenization/packing overlaps with training.""" + pad_to_seq_len: bool = False + """Pad every packed row to ``seq_len * micro_batch_size`` instead of to a multiple of ``model.cp``. Rows then have + one static shape, which avoids recompilation on dynamic sequence lengths (torch.compile, DeepGEMM JIT kernels, ...) + at the cost of extra padding.""" + @model_validator(mode="after") def validate_batch_size(self): if self.batch_size % self.micro_batch_size != 0: @@ -41,16 +51,7 @@ def validate_batch_size(self): return self -class PackedDataConfig(BaseDataConfig): - """Samples packed into rows of ``seq_len * micro_batch_size`` tokens.""" - - pad_to_seq_len: bool = False - """Pad every packed row to ``seq_len * micro_batch_size`` instead of to a multiple of ``model.cp``. Rows then have - one static shape, which avoids recompilation on dynamic sequence lengths (torch.compile, DeepGEMM JIT kernels, ...) - at the cost of extra padding.""" - - -class FakeDataConfig(PackedDataConfig): +class FakeDataConfig(BaseDataConfig): type: Literal["fake"] = "fake" length: Literal["fixed", "variable"] = "fixed" @@ -99,7 +100,7 @@ class SFTColumnsConfig(BaseConfig): """Per-sample renderer arguments as ``renderer field = dataset column``, e.g. ``reasoning_effort = "effort"``. A row's non-null value overrides the ``[renderer]`` setting; rows and datasets without the column use it unchanged.""" -class SFTDataConfig(PackedDataConfig): +class SFTDataConfig(BaseDataConfig): type: Literal["sft"] = "sft" name: str = "PrimeIntellect/Reverse-Text-SFT" @@ -154,6 +155,25 @@ def validate_subsets_and_splits(self): return self +class TokensDataConfig(BaseDataConfig): + """Pretraining data: ``seq_len`` windows of pre-tokenized datatrove shards, read as one token pool. Every token + is trained on. Mixtures are made offline, as shard folders.""" + + type: Literal["tokens"] = "tokens" + + paths: list[Path] = Field(min_length=1) + """Folders of datatrove tokenized shards: ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``. + The files of all folders (in folder order, then file name order) are read as one token pool.""" + + seed: int = 0 + """Seed for the window order: each epoch reads every window of the pool once, in a new random order and from a + new random offset.""" + + document_masking: bool = True + """Attention and positions restart at every document. If False, a window is one sequence: attention spans the + whole window and positions run over it.""" + + class SFTValConfig(BaseConfig): interval: int = Field(50, ge=1) """Run validation every N training steps.""" @@ -164,14 +184,114 @@ class SFTValConfig(BaseConfig): data: SFTDataConfig -DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig, Field(discriminator="type")] +DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig | TokensDataConfig, Field(discriminator="type")] + + +class TokenEvalSetConfig(BaseConfig): + """A held-out folder of datatrove tokenized shards, scored like a ``ppl`` set: every document on its own, in + windows of each of ``seq_lens``.""" + + name: str + """Names the metrics, ``eval//@``.""" + + path: Path + """Folder of ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``, as for ``tokens`` data.""" + + +class LogLikelihoodEvalConfig(BaseConfig): + """Log-likelihood evals scored inside the trainer: forward passes without gradients on the training model, every + ``interval`` steps. A set is a pretrain-evals set (``ppl``, ``ctxcont`` or ``pairs``) or a folder of held-out + token shards.""" + + interval: int = Field(ge=1) + """Score the sets every N steps (after that step's optimizer update), and at the last step.""" + + eval_on_start: bool = True + """Also score the weights the run starts from (logged at step 0, or at the step a resumed run starts from).""" + + data_dir: Path | None = None + """Directory of the pretrain-evals sets, one directory per set with ``test.parquet`` and ``meta.json``.""" + + sets: list[str | TokenEvalSetConfig] = Field(min_length=1) + """The sets to score: the name of a pretrain-evals set in ``data_dir``, or a held-out token folder.""" + + seq_lens: list[int] | None = None + """Context lengths to score every set at: the window length for ``ppl`` and token sets and the context limit for + ``ctxcont`` sets. Defaults to ``data.seq_len``. Metrics are named ``eval//@`` (e.g. + ``bpb@4k``).""" + + micro_batch_tokens: int | None = Field(None, ge=1) + """Tokens per packed eval row (one forward pass on one GPU). Defaults to the training micro-batch + (``data.micro_batch_size * data.seq_len``) or the longest of ``seq_lens``, whichever is larger; must be at least + the longest of ``seq_lens``.""" + + save_arrays: bool = False + """Write per-document arrays (npz) and summaries to ``/ll_eval/step_N/``.""" + + @model_validator(mode="after") + def validate_data_dir(self): + if self.data_dir is None and any(isinstance(eval_set, str) for eval_set in self.sets): + raise ValueError("ll_eval.data_dir is required for pretrain-evals sets") + return self + + +class SFTModelConfig(ModelConfig): + init: InitConfig | None = None + """Initialize the weights from scratch (pretraining) instead of loading them from ``name``, which then only + supplies the model config. A resumed run loads the weights from the checkpoint.""" + + defer_grad_sync: bool = False + """Reduce-scatter the gradients once per optimizer step instead of after every micro-batch. Saves communication + with gradient accumulation, but keeps the unsharded parameters and gradients in memory across the micro-batches.""" + + +class BaseDeploymentConfig(BaseConfig): + gpus_per_node: int = 8 + """GPUs per node.""" + + +class SingleNodeDeploymentConfig(BaseDeploymentConfig): + type: Literal["single_node"] = "single_node" + + num_train_gpus: int = 1 + """GPUs allocated to the trainer.""" + + num_infer_gpus: int = Field(1, validation_alias=AliasChoices("num_infer_gpus", "num_eval_gpus")) + """GPUs allocated to inference for online evals (alias: ``num_eval_gpus``). Only used when an ``[inference]`` block is configured.""" + + @model_validator(mode="after") + def validate_gpu_count(self): + if self.num_train_gpus > self.gpus_per_node: + raise ValueError(f"num_train_gpus ({self.num_train_gpus}) exceeds gpus_per_node ({self.gpus_per_node}).") + return self + + +class MultiNodeDeploymentConfig(BaseDeploymentConfig): + type: Literal["multi_node"] = "multi_node" + + num_train_nodes: int = Field(2, ge=1) + """Training nodes.""" + + num_infer_nodes: int = Field(0, ge=0, validation_alias=AliasChoices("num_infer_nodes", "num_eval_nodes")) + """Inference nodes for online evals (alias: ``num_eval_nodes``). These nodes share + one SLURM allocation with the trainer nodes.""" + + nodes_per_fsdp_group: int | None = None + """Nodes per FSDP island. Auto-sets ``model.dp_replicate = num_train_nodes / nodes_per_fsdp_group``.""" + +SFTDeploymentConfig: TypeAlias = Annotated[ + SingleNodeDeploymentConfig | MultiNodeDeploymentConfig, Field(discriminator="type") +] -class SFTConfig(TrainerLaunchConfig): + +class SFTConfig(BaseTrainerConfig): auto_freeze_moe_router = False # Standard L2 regularization for supervised training. auto_weight_decay = 0.01 + model: SFTModelConfig = SFTModelConfig() + renderer: RendererConfig = AutoRendererConfig() """Renderer config. Defaults to auto-selecting from the tokenizer model name.""" @@ -180,6 +300,9 @@ class SFTConfig(TrainerLaunchConfig): val: SFTValConfig | None = None """Validation configuration. If None, no validation runs.""" + ll_eval: LogLikelihoodEvalConfig | None = None + """Log-likelihood evals scored inside the trainer. If None, they never run.""" + eval: SFTOnlineEvalConfig | None = None """Online evaluation configuration: rollout-based evals against a live inference server that receives the trainer's weight broadcasts. If None, no online evals run.""" @@ -192,19 +315,69 @@ class SFTConfig(TrainerLaunchConfig): """Trainer-to-inference weight transport for online evals. Defaults to NCCL. LoRA and external inference use filesystem broadcast.""" - def default_run_name(self) -> str: - """``----``.""" - dataset = str(getattr(self.data, "name", "")).split("/")[-1] - model = self.model.name.split("/")[-1] - parts = [part for part in (dataset, model) if part] - return "--".join([*parts, uuid.uuid4().hex[:8]]).lower() + resume: ResumeConfig | None = None + """Resume the run from a checkpoint (point at it with the previous run's ``run.name``). Without ``[ckpt]`` the run loads the checkpoint but saves no new ones. If None, does not resume.""" + + monitors: TrainMonitorsConfig = TrainMonitorsConfig() + """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" + + run: RunConfig = Field(default_factory=RunConfig) + """Run metadata. ``run.name`` names the run directory under ``output_dir``.""" + + output_dir: Path = Field(default_factory=default_output_dir) + """Directory that groups related runs. Each run writes its artifacts (checkpoints, logs, ...) to ``output_dir / run.name``. Should be a persistent directory with enough disk space. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``.""" + + clean: bool = False + """Delete the run directory (``output_dir / run.name``) before starting training. Required to overwrite a run directory that contains artifacts from a previous run when not resuming.""" + + @property + def run_dir(self) -> Path: + assert self.run.dir is not None # resolved at construction + return self.output_dir / self.run.dir + + @model_validator(mode="after") + def auto_setup_run_identity(self): + """Auto-generate the run name (``----``) when unset and + default the run directory, W&B run name and platform run name to it when not + set explicitly.""" + if self.run.name is None: + dataset = str(getattr(self.data, "name", "")).split("/")[-1] + model = self.model.name.split("/")[-1] + parts = [part for part in (dataset, model) if part] + self.run.name = "--".join([*parts, uuid.uuid4().hex[:8]]).lower() + if self.run.dir is None: + self.run.dir = self.run.name + if self.monitors.wandb is not None and self.monitors.wandb.name is None: + self.monitors.wandb.name = self.run.name + if self.monitors.prime is not None and self.monitors.prime.name is None: + self.monitors.prime.name = self.run.name + return self + + deployment: SFTDeploymentConfig = SingleNodeDeploymentConfig() + + slurm: SlurmConfig | None = None + """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" dashboard: bool = True """Make sure a local dashboard daemon serves this run's output dir (started on demand in interactive sessions; an already-running daemon's URL is logged).""" + dry_run: bool = False + """Only validate and dump resolved configs, then exit early.""" + ### Pre-validation normalization + @model_validator(mode="before") + @classmethod + def normalize_deployment(cls, data): + if not isinstance(data, dict): + return data + deployment = data.get("deployment") + if isinstance(deployment, dict) and deployment.get("type") == "multi_node": + for key in ("num_train_gpus", "num_infer_gpus"): + deployment.pop(key, None) + return data + @model_validator(mode="before") @classmethod def propagate_model_to_inference(cls, data): @@ -223,6 +396,12 @@ def propagate_model_to_inference(cls, data): ### Validate configs (e.g. raise for unsupported (combinations of) configs) + @model_validator(mode="after") + def validate_deployment(self): + if self.deployment.type == "multi_node" and self.slurm is None: + raise ValueError("Must use SLURM for multi-node deployment.") + return self + @model_validator(mode="after") def auto_setup_online_eval(self): """Wire online evals and the trainer-to-inference weight transport.""" @@ -357,7 +536,7 @@ def auto_setup_online_eval(self): @model_validator(mode="after") def validate_typed_renderer(self): """Require a typed renderer whenever SFT renders real samples.""" - if self.data.type == "fake" and self.val is None: + if self.data.type in ("fake", "tokens") and self.val is None: return self model_id = self.tokenizer.name or self.model.name @@ -404,6 +583,27 @@ def validate_cp_micro_batch_size(self): raise ValueError("Validation micro batch size must be 1 when CP is enabled") return self + @model_validator(mode="after") + def validate_defer_grad_sync(self): + if self.model.defer_grad_sync and self.model.offload == "full": + raise ValueError("model.defer_grad_sync is not supported with model.offload = 'full'") + return self + + @model_validator(mode="after") + def auto_setup_ll_eval(self): + if self.ll_eval is None: + return self + if self.model.cp > 1: + raise ValueError("[ll_eval] does not support context parallelism yet") + if self.ll_eval.seq_lens is None: + self.ll_eval.seq_lens = [self.data.seq_len] + longest = max(self.ll_eval.seq_lens) + if self.ll_eval.micro_batch_tokens is None: + self.ll_eval.micro_batch_tokens = max(self.data.micro_batch_size * self.data.seq_len, longest) + elif self.ll_eval.micro_batch_tokens < longest: + raise ValueError(f"ll_eval.micro_batch_tokens must be at least the longest of ll_eval.seq_lens ({longest})") + return self + @model_validator(mode="after") def validate_vlm_constraints(self): if self.model.vlm is None: @@ -417,3 +617,30 @@ def validate_vlm_constraints(self): if self.val is not None and self.val.data.micro_batch_size != 1: raise ValueError("VLM SFT requires val.data.micro_batch_size = 1.") return self + + ### Auto-setup and validate shared configs + + @model_validator(mode="after") + def auto_setup_deployment(self): + if self.deployment.type == "multi_node": + if self.deployment.nodes_per_fsdp_group is not None: + if self.deployment.num_train_nodes % self.deployment.nodes_per_fsdp_group != 0: + raise ValueError( + f"deployment.num_train_nodes ({self.deployment.num_train_nodes}) must be divisible by " + f"deployment.nodes_per_fsdp_group ({self.deployment.nodes_per_fsdp_group})" + ) + self.model.dp_replicate = self.deployment.num_train_nodes // self.deployment.nodes_per_fsdp_group + return self + + @model_validator(mode="after") + def auto_setup_slurm_template(self): + if self.slurm is not None and self.slurm.job_name is None: + self.slurm.job_name = self.run.name + if self.slurm is not None and self.slurm.template_path is None: + templates_dir = find_package_resource("templates") + if templates_dir is not None: + if self.deployment.type == "single_node": + self.slurm.template_path = templates_dir / "single_node_sft.sbatch.j2" + else: + self.slurm.template_path = templates_dir / "multi_node_sft.sbatch.j2" + return self diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index d9a8740e943..1ce168817a9 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -3,7 +3,7 @@ from pathlib import Path from typing import Annotated, Literal, TypeAlias -from pydantic import AfterValidator, AliasChoices, Field, model_validator +from pydantic import AfterValidator, Field, model_validator from verifiers.v1.configs.client import BaseClientConfig from prime_rl.utils.config import BaseConfig @@ -298,43 +298,3 @@ class ZMQTransportConfig(BaseTransportConfig): TransportConfig: TypeAlias = Annotated[FileSystemTransportConfig | ZMQTransportConfig, Field(discriminator="type")] - - -class BaseDeploymentConfig(BaseConfig): - gpus_per_node: int = 8 - """GPUs per node.""" - - -class SingleNodeDeploymentConfig(BaseDeploymentConfig): - type: Literal["single_node"] = "single_node" - - num_train_gpus: int = 1 - """GPUs allocated to the trainer.""" - - num_infer_gpus: int = Field(1, validation_alias=AliasChoices("num_infer_gpus", "num_eval_gpus")) - """GPUs allocated to inference for online evals (alias: ``num_eval_gpus``). Only used when an ``[inference]`` block is configured.""" - - @model_validator(mode="after") - def validate_gpu_count(self): - if self.num_train_gpus > self.gpus_per_node: - raise ValueError(f"num_train_gpus ({self.num_train_gpus}) exceeds gpus_per_node ({self.gpus_per_node}).") - return self - - -class MultiNodeDeploymentConfig(BaseDeploymentConfig): - type: Literal["multi_node"] = "multi_node" - - num_train_nodes: int = Field(2, ge=1) - """Training nodes.""" - - num_infer_nodes: int = Field(0, ge=0, validation_alias=AliasChoices("num_infer_nodes", "num_eval_nodes")) - """Inference nodes for online evals (alias: ``num_eval_nodes``). These nodes share - one SLURM allocation with the trainer nodes.""" - - nodes_per_fsdp_group: int | None = None - """Nodes per FSDP island. Auto-sets ``model.dp_replicate = num_train_nodes / nodes_per_fsdp_group``.""" - - -SFTDeploymentConfig: TypeAlias = Annotated[ - SingleNodeDeploymentConfig | MultiNodeDeploymentConfig, Field(discriminator="type") -] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 96e0969d74c..01b01e6749a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -1,12 +1,11 @@ import re -import uuid import warnings from pathlib import Path from typing import Annotated, Any, ClassVar, Literal, TypeAlias from pydantic import Field, field_validator, model_validator -from prime_rl.configs.monitors import MonitorsConfig, TrainMonitorsConfig +from prime_rl.configs.monitors import MonitorsConfig from prime_rl.configs.shared import ( BaseModelConfig, ControlConfig, @@ -14,16 +13,12 @@ FileSystemWeightBroadcastConfig, HeartbeatConfig, ResumeConfig, - RunConfig, - SFTDeploymentConfig, - SingleNodeDeploymentConfig, - SlurmConfig, TrainerLogConfig, TransportConfig, WeightBroadcastConfig, ZMQTransportConfig, ) -from prime_rl.utils.config import BaseConfig, default_output_dir, find_package_resource +from prime_rl.utils.config import BaseConfig, default_output_dir # -- Shared trainer configs (used by both SFT and RL trainers) -- @@ -268,7 +263,7 @@ class MoERuntimeConfig(BaseConfig): class InitConfig(BaseConfig): """Random weight initialization: every weight matrix and embedding from N(0, std). Norm weights start at 1, - biases at 0.""" + biases and other 1-D parameters at 0.""" std: float = Field(0.02, gt=0) """Standard deviation.""" @@ -736,98 +731,6 @@ def auto_setup_tokenizer(self): return self -class TrainerLaunchConfig(BaseTrainerConfig): - """Fields and validators of the trainers that launch on their own (SFT and pretraining): run identity, output - directory, monitors, resume, deployment and SLURM.""" - - resume: ResumeConfig | None = None - """Resume the run from a checkpoint (point at it with the previous run's ``run.name``). Without ``[ckpt]`` the run loads the checkpoint but saves no new ones. If None, does not resume.""" - - monitors: TrainMonitorsConfig = TrainMonitorsConfig() - """Metric monitors (``monitors.wandb``, ``monitors.file``, ``monitors.prime``).""" - - run: RunConfig = Field(default_factory=RunConfig) - """Run metadata. ``run.name`` names the run directory under ``output_dir``.""" - - output_dir: Path = Field(default_factory=default_output_dir) - """Directory that groups related runs. Each run writes its artifacts (checkpoints, logs, ...) to ``output_dir / run.name``. Should be a persistent directory with enough disk space. Defaults to ``$PRL_OUTPUT_DIR`` if set, else ``outputs``.""" - - clean: bool = False - """Delete the run directory (``output_dir / run.name``) before starting training. Required to overwrite a run directory that contains artifacts from a previous run when not resuming.""" - - deployment: SFTDeploymentConfig = SingleNodeDeploymentConfig() - - slurm: SlurmConfig | None = None - """SLURM configuration. When set, the run is submitted as a SLURM job instead of running locally.""" - - dry_run: bool = False - """Only validate and dump resolved configs, then exit early.""" - - @property - def run_dir(self) -> Path: - assert self.run.dir is not None # resolved at construction - return self.output_dir / self.run.dir - - def default_run_name(self) -> str: - return f"{self.model.name.split('/')[-1]}--{uuid.uuid4().hex[:8]}".lower() - - @model_validator(mode="before") - @classmethod - def normalize_deployment(cls, data): - if not isinstance(data, dict): - return data - deployment = data.get("deployment") - if isinstance(deployment, dict) and deployment.get("type") == "multi_node": - for key in ("num_train_gpus", "num_infer_gpus"): - deployment.pop(key, None) - return data - - @model_validator(mode="after") - def auto_setup_run_identity(self): - """Default the run name (``default_run_name``) and the run directory, W&B run name and platform run name to it - when not set explicitly.""" - if self.run.name is None: - self.run.name = self.default_run_name() - if self.run.dir is None: - self.run.dir = self.run.name - if self.monitors.wandb is not None and self.monitors.wandb.name is None: - self.monitors.wandb.name = self.run.name - if self.monitors.prime is not None and self.monitors.prime.name is None: - self.monitors.prime.name = self.run.name - return self - - @model_validator(mode="after") - def validate_deployment(self): - if self.deployment.type == "multi_node" and self.slurm is None: - raise ValueError("Must use SLURM for multi-node deployment.") - return self - - @model_validator(mode="after") - def auto_setup_deployment(self): - if self.deployment.type == "multi_node": - if self.deployment.nodes_per_fsdp_group is not None: - if self.deployment.num_train_nodes % self.deployment.nodes_per_fsdp_group != 0: - raise ValueError( - f"deployment.num_train_nodes ({self.deployment.num_train_nodes}) must be divisible by " - f"deployment.nodes_per_fsdp_group ({self.deployment.nodes_per_fsdp_group})" - ) - self.model.dp_replicate = self.deployment.num_train_nodes // self.deployment.nodes_per_fsdp_group - return self - - @model_validator(mode="after") - def auto_setup_slurm_template(self): - if self.slurm is not None and self.slurm.job_name is None: - self.slurm.job_name = self.run.name - if self.slurm is not None and self.slurm.template_path is None: - templates_dir = find_package_resource("templates") - if templates_dir is not None: - if self.deployment.type == "single_node": - self.slurm.template_path = templates_dir / "single_node_sft.sbatch.j2" - else: - self.slurm.template_path = templates_dir / "multi_node_sft.sbatch.j2" - return self - - class TrainerConfig(BaseTrainerConfig): # RL keeps the pretrained router. auto_freeze_moe_router = True diff --git a/pyproject.toml b/pyproject.toml index 5174a2910df..46380f46a46 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,6 @@ dependencies = [ [project.scripts] rl = "prime_rl.entrypoints.rl:main" sft = "prime_rl.entrypoints.sft:main" -pretrain = "prime_rl.entrypoints.pretrain:main" inference = "prime_rl.entrypoints.inference:main" trainer = "prime_rl.entrypoints.trainer:main" orchestrator = "prime_rl.entrypoints.orchestrator:main" diff --git a/src/prime_rl/dashboard/server.py b/src/prime_rl/dashboard/server.py index 17e27c3ac63..2681f70e73b 100755 --- a/src/prime_rl/dashboard/server.py +++ b/src/prime_rl/dashboard/server.py @@ -246,8 +246,6 @@ def main_config(run_dir: Path) -> tuple[str, dict]: return "rl", read_json(configs / "orchestrator.json") or read_json(configs / "trainer.json") if (configs / "eval.json").exists(): return "eval", read_json(configs / "eval.json") - if (configs / "pretrain.json").exists(): - return "pretrain", read_json(configs / "pretrain.json") return "other", {} @@ -479,7 +477,7 @@ def read_log(run: str, file: str, start: int | None = None, end: int | None = No # ------------------------------------------------------------------------- configs -CONFIG_ORDER = ["rl", "sft", "eval", "pretrain", "orchestrator", "trainer", "inference"] +CONFIG_ORDER = ["rl", "sft", "eval", "orchestrator", "trainer", "inference"] def config_rank(name: str) -> tuple[int, str]: diff --git a/src/prime_rl/dashboard/static/app.js b/src/prime_rl/dashboard/static/app.js index 41d27b41615..c53995009ef 100644 --- a/src/prime_rl/dashboard/static/app.js +++ b/src/prime_rl/dashboard/static/app.js @@ -446,11 +446,6 @@ const SFT_TRAIN_METRICS = ["loss/mean", "loss/perplexity", "val/loss", "val/perp // max_vio only exists on MoE models: a regex panel stays away on dense ones const SFT_STABILITY_PANELS = [{ metric: "optim/grad_norm" }, { metric: "optim/lr" }, { metric: "loss/nan_count" }, { regex: "max_vio/(mean|max)" }]; const SFT_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memory", "time/forward_backward"]; -// Pretrain flavor (overview.py PRETRAIN_*): bpb of the evals (eval/) -const PRETRAIN_TRAIN_PANELS = [{ metric: "loss/mean" }, { metric: "loss/spike" }, { regex: "eval/.*/bpb" }]; -const PRETRAIN_STABILITY_PANELS = [{ metric: "optim/grad_norm" }, { metric: "optim/grad_norm_spike" }, { metric: "optim/lr" }, { metric: "loss/nan_count" }, { regex: "optim/grad_norm/.*" }]; -const PRETRAIN_MOE_METRICS = ["expert_load/max_mean/mean", "expert_load/max_mean/max", "expert_load/micro_max_mean/mean", "expert_load/micro_max_mean/max", "expert_load/entropy/min", "expert_load/dead", "router/confidence", "router/bias/absmax", "loss/seq_balance"]; -const PRETRAIN_PERFORMANCE_METRICS = ["perf/mfu", "perf/throughput", "perf/peak_memory/max", "time/step", "time/data/max", "time/eval_suite"]; // Multi-series inference panels (overview.py INFERENCE_PANELS): fleet aggregate // paired with the cross-engine tail that flags a single sick engine. @@ -1462,13 +1457,6 @@ function buildSections(meta) { }); const sections = []; const evalEnvs = meta.eval_envs || []; - if (meta.type === "pretrain") { - sections.push({ name: "train", panels: PRETRAIN_TRAIN_PANELS }); - sections.push({ name: "stability", panels: PRETRAIN_STABILITY_PANELS }); - sections.push({ name: "moe", panels: PRETRAIN_MOE_METRICS.map((m) => ({ metric: m })) }); - sections.push({ name: "performance", panels: PRETRAIN_PERFORMANCE_METRICS.map((m) => ({ metric: m })) }); - return sections; - } if (meta.type === "sft") { // the val panes show once a validation value is logged, not on the config alone sections.push({ name: "train", panels: SFT_TRAIN_METRICS.map((m) => (m.startsWith("val/") ? { regex: escRe(m) } : { metric: m })) }); diff --git a/src/prime_rl/entrypoints/pretrain.py b/src/prime_rl/entrypoints/pretrain.py deleted file mode 100644 index bbd24036ab2..00000000000 --- a/src/prime_rl/entrypoints/pretrain.py +++ /dev/null @@ -1,94 +0,0 @@ -import os -from pathlib import Path - -from prime_rl.configs.pretrain import PretrainConfig -from prime_rl.entrypoints.sft import submit_slurm_script, write_config, write_trainer_slurm_script -from prime_rl.utils.config import cli -from prime_rl.utils.logger import get_logger, setup_logger -from prime_rl.utils.pathing import ( - format_config_message, - format_log_message, - get_launcher_dir, - prepare_attempt_dirs, - prepare_run, - write_launch_artifacts, -) -from prime_rl.utils.process import ( - DEFAULT_COMMON_ENV_VARS, - DEFAULT_TRAINER_ENV_VARS, - ProcessGroup, - set_proc_title, - torchrun_cmd, -) - -PRETRAIN_CONFIG = "pretrain.json" -PRETRAIN_SBATCH = "pretrain.sbatch" - - -def pretrain_slurm(config: PretrainConfig, config_path: Path, log_dir: Path) -> None: - logger = get_logger() - script_path = get_launcher_dir(config.run_dir) / PRETRAIN_SBATCH - write_trainer_slurm_script(config, "pretrain", config_path, log_dir, script_path) - logger.info(f"Wrote SLURM script to {script_path}") - num_nodes = config.deployment.num_train_nodes if config.deployment.type == "multi_node" else 1 - log_message = format_log_message(log_dir=log_dir, trainer=True, num_train_nodes=num_nodes) - if config.dry_run: - logger.success(f"Dry run complete. To submit manually:\n\n sbatch {script_path}\n\n{log_message}") - return - submit_slurm_script(script_path, log_message, logger) - - -def pretrain(config: PretrainConfig): - logger = setup_logger(config.log.level or "info", json_logging=config.log.json_logging) - prepare_run(config, config.ckpt.output_dir if config.ckpt else None) - - logger.info("Starting pretraining run") - config_dir, log_dir = prepare_attempt_dirs(config.run_dir) - write_launch_artifacts(config_dir, "pretrain") - - if not config.dry_run: - from prime_rl.trainer.model import pre_download_model - - pre_download_model(config.model.name, skip_weights=True) - - config_path = config_dir / PRETRAIN_CONFIG - exclude = {"dry_run", "clean"} - if config.slurm is not None: - # A single-node job runs this launcher again on the node; the trainers of a - # multi-node job read the config directly. - exclude |= {"slurm", "deployment"} if config.deployment.type == "multi_node" else {"slurm"} - write_config(config, config_path, exclude=exclude) - logger.info(f"Configs:\n{format_config_message(config_dir, 'pretrain', [('Pretrain', config_path)])}") - if config.slurm is not None: - pretrain_slurm(config, config_path, log_dir) - return - if config.dry_run: - logger.success("Dry run complete. To start the run, remove --dry-run from your command.") - return - - logger.info(format_log_message(log_dir=log_dir, trainer=True)) - with ProcessGroup() as processes: - processes.start( - "trainer", - torchrun_cmd( - "prime_rl.trainer.pretrain.train", - config_path, - config.deployment.num_train_gpus, - log_dir, - config.log.ranks_filter, - ), - env={**os.environ, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, **config.env_vars}, - log_path=log_dir / "trainer.log", - ) - logger.success(f"Started trainer with {config.deployment.num_train_gpus} GPU(s)") - processes.wait("trainer") - logger.success("Pretraining finished!") - - -def main(): - set_proc_title("Pretrain") - pretrain(cli(PretrainConfig)) - - -if __name__ == "__main__": - main() diff --git a/src/prime_rl/entrypoints/sft.py b/src/prime_rl/entrypoints/sft.py index 6459768d22e..57912b853ec 100644 --- a/src/prime_rl/entrypoints/sft.py +++ b/src/prime_rl/entrypoints/sft.py @@ -7,7 +7,6 @@ from prime_rl.configs.eval import SFTOnlineEvalConfig from prime_rl.configs.monitors import EvalMonitorsConfig, PrimeEvalMonitorConfig, TrainMonitorsConfig from prime_rl.configs.orchestrator import OnlineEvalSourceConfig -from prime_rl.configs.pretrain import PretrainConfig from prime_rl.configs.sft import SFTConfig from prime_rl.configs.shared import LogConfig from prime_rl.entrypoints.dashboard import ensure_dashboard, log_dashboard_url @@ -23,7 +22,7 @@ get_launcher_log_dir, prepare_attempt_dirs, prepare_run, - resolve_resume_step, + resolve_latest_ckpt_step, write_env_server_config, write_launch_artifacts, ) @@ -59,6 +58,16 @@ def get_ckpt_base(config: SFTConfig) -> Path: return (config.ckpt.output_dir if config.ckpt else None) or config.run_dir +def resolve_resume_step(config: SFTConfig) -> int | None: + if config.resume is None: + return None + if config.resume.dir is not None: + return config.resume.dir_step + if config.resume.step is not None: + return config.resume.step + return resolve_latest_ckpt_step(get_ckpt_dir(get_ckpt_base(config))) + + def build_online_eval_monitors(monitors: TrainMonitorsConfig) -> EvalMonitorsConfig: """The online-eval process's monitors: same W&B and file config as the trainer's, with the platform monitor switched to the evaluation flavor under the same name.""" @@ -81,7 +90,7 @@ def build_online_eval_config(config: SFTConfig) -> SFTOnlineEvalConfig: weight_broadcast=config.weight_broadcast, broadcasts_dir=get_broadcast_dir(config.run_dir), max_steps=config.max_steps, - resume_step=resolve_resume_step(config.resume, get_ckpt_dir(get_ckpt_base(config))), + resume_step=resolve_resume_step(config), # Same run dir as the trainer: the dashboard reads the eval's # file-monitor artifacts from the run directory, so a separate # output dir would hide them. The train-vs-eval record collision @@ -107,7 +116,7 @@ def sft_config_components(config: SFTConfig, config_dir: Path) -> list[tuple[str return components -def write_config(config: SFTConfig | PretrainConfig, config_path: Path, exclude: set[str] | None = None) -> None: +def write_config(config: SFTConfig, config_path: Path, exclude: set[str] | None = None) -> None: """Write resolved config to disk, excluding launcher-only fields.""" config_path.parent.mkdir(parents=True, exist_ok=True) with open(config_path, "w") as f: @@ -138,17 +147,10 @@ def write_eval_subconfigs(config: SFTConfig, config_dir: Path, strip_router: boo write_env_server_config(config_dir, "eval", source, config.log) -def write_trainer_slurm_script( - config: SFTConfig | PretrainConfig, - entrypoint: str, - config_path: Path, - log_dir: Path, - script_path: Path, - **template_vars, +def write_slurm_script( + config: SFTConfig, config_path: Path, log_dir: Path, script_path: Path, prl_run_id: str, wandb_shared: bool ) -> None: - """Write the SLURM script of a trainer entrypoint (``sft`` or ``pretrain``). Without - ``template_vars`` a multi-node job runs only trainer nodes; SFT passes its online-eval - variables on top.""" + """Write the SLURM script to disk.""" from jinja2 import Environment, FileSystemLoader assert config.slurm is not None @@ -161,82 +163,80 @@ def write_trainer_slurm_script( env = Environment(loader=FileSystemLoader(template_dirs), keep_trailing_newline=True) template = env.get_template(config.slurm.template_path.name) - trainer_vars = {} - if config.deployment.type == "multi_node": - trainer_vars = dict( - trainer_env_vars={**DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, **config.env_vars}, + trainer_env_vars = { + **DEFAULT_COMMON_ENV_VARS, + **DEFAULT_TRAINER_ENV_VARS, + **config.env_vars, + } + + if config.deployment.type == "single_node": + script = template.render( + **config.slurm.template_vars, + config_path=config_path, + config_dir=config_path.parent, + log_dir=log_dir, + output_dir=config.run_dir, + launcher_dir=get_launcher_dir(config.run_dir), + launcher_log_dir=get_launcher_log_dir(config.run_dir), + gpus_per_node=config.deployment.gpus_per_node, + ) + else: + online_eval = config.eval is not None + eval_vars = {} + if online_eval: + assert config.inference is not None + inference_env_vars = { + **DEFAULT_COMMON_ENV_VARS, + **DEFAULT_INFERENCE_ENV_VARS, + **config.env_vars, + **config.inference.env_vars, + } + eval_vars = { + "router": config.inference.router, + "router_port": config.inference.server.port, + "backend_port": config.inference.backend_port, + "data_parallel_rpc_port": config.inference.vllm.data_parallel_rpc_port, + "dp_per_node": config.deployment.gpus_per_node // config.inference.vllm.tensor_parallel_size, + "enable_expert_parallel": config.inference.vllm.enable_expert_parallel, + "inference_env_vars": inference_env_vars, + "online_eval_env_vars": { + **DEFAULT_COMMON_ENV_VARS, + "LOGURU_FORCE_COLORS": "1", + **config.env_vars, + }, + "eval_env_names": [source.resolved_name for source in eval_env_servers(config)], + } + script = template.render( + **config.slurm.template_vars, + config_path=config_path, + config_dir=config_path.parent, + log_dir=log_dir, + output_dir=config.run_dir, + launcher_dir=get_launcher_dir(config.run_dir), + launcher_log_dir=get_launcher_log_dir(config.run_dir), + trainer_env_vars=trainer_env_vars, num_nodes=config.deployment.num_train_nodes, num_train_nodes=config.deployment.num_train_nodes, - num_infer_nodes=0, + num_infer_nodes=config.deployment.num_infer_nodes if online_eval else 0, + gpus_per_node=config.deployment.gpus_per_node, ranks_filter=",".join(map(str, config.log.ranks_filter)), - prl_run_id=os.environ["PRL_RUN_ID"], - wandb_shared=False, + prl_run_id=prl_run_id, + wandb_shared=wandb_shared, run_name=config.run.name, - online_eval=False, - use_nccl_broadcast=False, + online_eval=online_eval, + use_nccl_broadcast=( + config.eval is not None + and config.weight_broadcast is not None + and config.weight_broadcast.type == "nccl" + ), + **eval_vars, ) - script = template.render( - **config.slurm.template_vars, - entrypoint=entrypoint, - config_path=config_path, - config_dir=config_path.parent, - log_dir=log_dir, - output_dir=config.run_dir, - launcher_dir=get_launcher_dir(config.run_dir), - launcher_log_dir=get_launcher_log_dir(config.run_dir), - gpus_per_node=config.deployment.gpus_per_node, - **{**trainer_vars, **template_vars}, - ) script_path.parent.mkdir(parents=True, exist_ok=True) get_launcher_log_dir(config.run_dir).mkdir(parents=True, exist_ok=True) script_path.write_text(script) -def submit_slurm_script(script_path: Path, log_message: str, logger) -> None: - logger.info(f"Submitting: sbatch {script_path}") - result = subprocess.run(["sbatch", str(script_path)], capture_output=True, text=True) - if result.returncode != 0: - logger.error(f"sbatch failed: {result.stderr.strip()}") - sys.exit(1) - logger.success(f"{result.stdout.strip()}\n\n{log_message}") - - -def write_slurm_script( - config: SFTConfig, config_path: Path, log_dir: Path, script_path: Path, wandb_shared: bool -) -> None: - """Write the SLURM script to disk, with the online-eval nodes of a multi-node run.""" - eval_vars = {} - if config.deployment.type == "multi_node" and config.eval is not None: - assert config.inference is not None - inference_env_vars = { - **DEFAULT_COMMON_ENV_VARS, - **DEFAULT_INFERENCE_ENV_VARS, - **config.env_vars, - **config.inference.env_vars, - } - eval_vars = { - "num_infer_nodes": config.deployment.num_infer_nodes, - "wandb_shared": wandb_shared, - "online_eval": True, - "use_nccl_broadcast": config.weight_broadcast is not None and config.weight_broadcast.type == "nccl", - "router": config.inference.router, - "router_port": config.inference.server.port, - "backend_port": config.inference.backend_port, - "data_parallel_rpc_port": config.inference.vllm.data_parallel_rpc_port, - "dp_per_node": config.deployment.gpus_per_node // config.inference.vllm.tensor_parallel_size, - "enable_expert_parallel": config.inference.vllm.enable_expert_parallel, - "inference_env_vars": inference_env_vars, - "online_eval_env_vars": { - **DEFAULT_COMMON_ENV_VARS, - "LOGURU_FORCE_COLORS": "1", - **config.env_vars, - }, - "eval_env_names": [source.resolved_name for source in eval_env_servers(config)], - } - write_trainer_slurm_script(config, "sft", config_path, log_dir, script_path, **eval_vars) - - def sft_slurm(config: SFTConfig): """Run SFT training and its online-eval deployment in one SLURM allocation.""" assert config.slurm is not None @@ -268,7 +268,7 @@ def sft_slurm(config: SFTConfig): write_eval_subconfigs(config, config_dir, strip_router=True) logger.info(f"Configs:\n{format_config_message(config_dir, 'sft', sft_config_components(config, config_dir))}") script_path = launcher_dir / SFT_SBATCH - write_slurm_script(config, config_path, log_dir, script_path, wandb_shared) + write_slurm_script(config, config_path, log_dir, script_path, os.environ["PRL_RUN_ID"], wandb_shared) logger.info(f"Wrote SLURM script to {script_path}") num_nodes = config.deployment.num_train_nodes if config.deployment.type == "multi_node" else 1 @@ -288,7 +288,13 @@ def sft_slurm(config: SFTConfig): dashboard_url = ensure_dashboard(config.output_dir, logger) if config.dashboard else None - submit_slurm_script(script_path, log_message, logger) + logger.info(f"Submitting: sbatch {script_path}") + result = subprocess.run(["sbatch", str(script_path)], capture_output=True, text=True) + if result.returncode != 0: + logger.error(f"sbatch failed: {result.stderr.strip()}") + sys.exit(1) + + logger.success(f"{result.stdout.strip()}\n\n{log_message}") log_dashboard_url(logger, dashboard_url) @@ -441,7 +447,7 @@ def clean_stale_eval_artifacts(config: SFTConfig) -> None: if os.environ.get("NEVER_CLEAN"): logger.warning("NEVER_CLEAN is set - keeping stale weight broadcasts; the online-eval process may replay them") return - resume_step = resolve_resume_step(config.resume, get_ckpt_dir(get_ckpt_base(config))) + resume_step = resolve_resume_step(config) clean_future_steps(config.run_dir, resume_step if resume_step is not None else -1) diff --git a/src/prime_rl/monitors/__init__.py b/src/prime_rl/monitors/__init__.py index 7d7960dc9c1..53ce9413ae0 100644 --- a/src/prime_rl/monitors/__init__.py +++ b/src/prime_rl/monitors/__init__.py @@ -55,7 +55,7 @@ async def setup( run_config: BaseConfig | None = None, train_env_names: list[str] | None = None, eval_env_names: list[str] | None = None, - overview_flavor: Literal["rl", "sft", "eval", "pretrain"] = "rl", + overview_flavor: Literal["rl", "sft", "eval"] = "rl", ) -> None: """Construct, initialize, and register one monitor per non-None config. diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 36384532641..f64be4b8236 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -14,7 +14,6 @@ from prime_sandboxes import Config as PrimeConfig from prime_rl.configs.monitors import PrimeEvalMonitorConfig, PrimeTrainMonitorConfig -from prime_rl.configs.pretrain import PretrainConfig from prime_rl.configs.sft import SFTConfig from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.utils.config import BaseConfig @@ -157,8 +156,8 @@ async def init(self, config: BaseConfig | None = None, output_dir: Path | None = # attached run's failure/completion marking; a clean finish() is # delivered best-effort and does not by itself mark it completed. init_kwargs = {"id": run_id} - elif isinstance(config, SFTConfig | PretrainConfig): - # SFT and pretraining train on data, not on environments, and have no rollouts - the + elif isinstance(config, SFTConfig): + # SFT trains on a dataset, not on environments, and has no rollouts - the # platform run is registered without both. init_kwargs = dict( name=self.config.name, diff --git a/src/prime_rl/monitors/wandb/monitor.py b/src/prime_rl/monitors/wandb/monitor.py index 33ae1c76849..e550ebf19e3 100644 --- a/src/prime_rl/monitors/wandb/monitor.py +++ b/src/prime_rl/monitors/wandb/monitor.py @@ -33,7 +33,7 @@ async def init( config: BaseConfig | None = None, train_env_names: list[str] | None = None, eval_env_names: list[str] | None = None, - overview_flavor: Literal["rl", "sft", "eval", "pretrain"] = "rl", + overview_flavor: Literal["rl", "sft", "eval"] = "rl", ) -> None: # W&B reads the start command off sys.argv; the launcher passes the original # command to subprocesses via $WANDB_ARGS. diff --git a/src/prime_rl/monitors/wandb/overview.py b/src/prime_rl/monitors/wandb/overview.py index 8c6c52a9210..32c949f6345 100644 --- a/src/prime_rl/monitors/wandb/overview.py +++ b/src/prime_rl/monitors/wandb/overview.py @@ -91,31 +91,6 @@ "time/save_ckpt", ] -# Pretrain flavor: loss and held-out evals, training stability, MoE balance, performance. -PRETRAIN_TRAIN_METRICS = ["loss/mean", "loss/spike"] -PRETRAIN_TRAIN_REGEXES = ["eval/.*/bpb"] -PRETRAIN_STABILITY_METRICS = ["optim/grad_norm", "optim/grad_norm_spike", "optim/lr", "loss/nan_count"] -PRETRAIN_STABILITY_REGEXES = ["optim/grad_norm/.*"] -PRETRAIN_MOE_METRICS = [ - "expert_load/max_mean/mean", - "expert_load/max_mean/max", - "expert_load/micro_max_mean/mean", - "expert_load/micro_max_mean/max", - "expert_load/entropy/min", - "expert_load/dead", - "router/confidence", - "router/bias/absmax", - "loss/seq_balance", -] -PRETRAIN_PERFORMANCE_METRICS = [ - "perf/mfu", - "perf/throughput", - "perf/peak_memory/max", - "time/step", - "time/data/max", - "time/eval_suite", -] - # Dense grid: more, smaller panels per row and enough rows that sections don't paginate. COLUMNS = 4 ROWS = 6 @@ -175,17 +150,8 @@ def eval_section(name: str, env_pattern: str) -> ws.Section: def build_sections( - train_envs: Sequence[str] = (), - eval_envs: Sequence[str] = (), - flavor: Literal["rl", "sft", "eval", "pretrain"] = "rl", + train_envs: Sequence[str] = (), eval_envs: Sequence[str] = (), flavor: Literal["rl", "sft", "eval"] = "rl" ) -> list[ws.Section]: - if flavor == "pretrain": - return [ - section("train", metrics=PRETRAIN_TRAIN_METRICS, regexes=PRETRAIN_TRAIN_REGEXES), - section("stability", metrics=PRETRAIN_STABILITY_METRICS, regexes=PRETRAIN_STABILITY_REGEXES), - section("moe", metrics=PRETRAIN_MOE_METRICS), - section("performance", metrics=PRETRAIN_PERFORMANCE_METRICS), - ] # A standalone eval run has no training signal: one section per eval env. if flavor == "eval": if eval_envs: @@ -265,7 +231,7 @@ def ensure_overview_view( name: str = OVERVIEW_NAME, train_envs: Sequence[str] = (), eval_envs: Sequence[str] = (), - flavor: Literal["rl", "sft", "eval", "pretrain"] = "rl", + flavor: Literal["rl", "sft", "eval"] = "rl", ) -> str | None: """Ensure an overview saved view exists for this run's env set. Reuses an existing overview built for the same envs; when the env set is new, creates a fresh versioned view (``overview`` → diff --git a/src/prime_rl/templates/multi_node_sft.sbatch.j2 b/src/prime_rl/templates/multi_node_sft.sbatch.j2 index 6dcb0ca9e44..32027951d71 100644 --- a/src/prime_rl/templates/multi_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_sft.sbatch.j2 @@ -71,7 +71,7 @@ export WANDB_RUN_ID=$PRL_RUN_ID export WANDB_SHARED_MODE=1 export WANDB_SHARED_PRIMARY=trainer export WANDB_SHARED_FINISHER=online-eval -export WANDB_PROGRAM="uv run {{ entrypoint }}" +export WANDB_PROGRAM="uv run sft" {% endif %} # Networking. Inference nodes occupy the first ranks, followed by trainer nodes. @@ -241,7 +241,7 @@ else --tee=3 \ --redirects=3 \ --local-ranks-filter={{ ranks_filter }} \ - -m prime_rl.trainer.{{ entrypoint }}.train \ + -m prime_rl.trainer.sft.train \ @ $TRAIN_CONFIG_PATH \ {%- if use_nccl_broadcast %} --weight-broadcast.host 0.0.0.0 \ diff --git a/src/prime_rl/templates/single_node_sft.sbatch.j2 b/src/prime_rl/templates/single_node_sft.sbatch.j2 index d1183709f45..9723eb491e9 100644 --- a/src/prime_rl/templates/single_node_sft.sbatch.j2 +++ b/src/prime_rl/templates/single_node_sft.sbatch.j2 @@ -41,4 +41,4 @@ uv sync --all-extras --all-packages # Pre-run command {{ pre_run_command }} {% endif %} -uv run {{ entrypoint }} @ {{ config_path }} +uv run sft @ {{ config_path }} diff --git a/src/prime_rl/trainer/pretrain/metrics.py b/src/prime_rl/trainer/metrics.py similarity index 79% rename from src/prime_rl/trainer/pretrain/metrics.py rename to src/prime_rl/trainer/metrics.py index 47f698d7a9f..4c5a3ce9451 100644 --- a/src/prime_rl/trainer/pretrain/metrics.py +++ b/src/prime_rl/trainer/metrics.py @@ -7,7 +7,7 @@ from torch import Tensor, nn from torch.distributed.tensor import DTensor, Partial, Replicate -from prime_rl.trainer.model import compute_expert_load_stats, iter_moe_blocks +from prime_rl.trainer.model import compute_expert_load_stats from prime_rl.trainer.models.layers.moe import MoE from prime_rl.utils.vlm import get_language_model @@ -29,10 +29,9 @@ def reduce_stats(sums: dict[str, Tensor], maxes: dict[str, Tensor], group: dist. return reduced -def take_routing_counts(model: nn.Module) -> tuple[Tensor, Tensor]: +def take_routing_counts(moes: list[MoE]) -> tuple[Tensor, Tensor]: """This rank's `[num_moe_layers, num_experts]` expert counts and `[num_moe_layers]` routing-confidence sums since the last call, and reset them.""" - moes = list(iter_moe_blocks(model)) counts = torch.stack([moe.tokens_per_expert for moe in moes]) confidence = torch.stack([moe.routing_confidence_sum for moe in moes]) for moe in moes: @@ -41,35 +40,6 @@ def take_routing_counts(model: nn.Module) -> tuple[Tensor, Tensor]: return counts, confidence -def routed_k_metrics(model: nn.Module, group: dist.ProcessGroup) -> dict[str, float]: - """Routed experts per training token since the last call, for variable-k routing (null experts), over all ranks: - mean, std, share of tokens with none and with the maximum (the mean over MoE layers), the worst layer's std, and - the busiest rank's routed pairs per token over the mean.""" - moes = [moe for moe in iter_moe_blocks(model) if moe.routed_k_counts is not None] - if not moes: - return {} - counts = torch.stack([moe.routed_k_counts for moe in moes]) - for moe in moes: - moe.routed_k_counts.zero_() - k = torch.arange(counts.shape[1], device=counts.device) - # This rank's routed pairs per token, over the global mean: the load the data puts on one rank. - rank_mean = (counts * k).sum() / counts.sum() - rank_max = rank_mean.clone() - dist.all_reduce(rank_max, op=dist.ReduceOp.MAX, group=group) - dist.all_reduce(counts, group=group) - tokens = counts.sum(dim=1) - mean = (counts * k).sum(dim=1) / tokens - std = ((counts * (k - mean[:, None]) ** 2).sum(dim=1) / tokens).sqrt() - return { - "router/k/mean": mean.mean().item(), - "router/k/std": std.mean().item(), - "router/k/std/max": std.max().item(), - "router/k/zero_frac": (counts[:, 0] / tokens).mean().item(), - "router/k/max_frac": (counts[:, -1] / tokens).mean().item(), - "router/k/rank_max_mean": (rank_max / mean.mean()).item(), - } - - def moe_layer_indices(model: nn.Module) -> list[int]: """The decoder-layer index of every MoE layer, in the order of `iter_moe_blocks`.""" return [ diff --git a/src/prime_rl/trainer/micro_batches.py b/src/prime_rl/trainer/micro_batches.py new file mode 100644 index 00000000000..ac5355be2d9 --- /dev/null +++ b/src/prime_rl/trainer/micro_batches.py @@ -0,0 +1,140 @@ +"""Forward and backward passes over one optimizer step's micro-batches. The step's statistics stay on the GPU until +the last micro-batch and are then reduced with one SUM and one MAX all-reduce: no host sync per micro-batch.""" + +import time +from collections.abc import Callable, Iterator +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch import Tensor, nn +from torch.profiler import record_function + +from prime_rl.trainer.metrics import expert_load_metrics, moe_layer_indices, reduce_stats, take_routing_counts +from prime_rl.trainer.model import compute_max_vio_and_routing_confidence, iter_moe_blocks +from prime_rl.trainer.optim.offload import GradientOffloadManager +from prime_rl.trainer.parallel_dims import ParallelDims +from prime_rl.trainer.perf import PerfCounter +from prime_rl.trainer.utils import begin_backward, finish_backward, scale_gradients_ + + +@dataclass +class StepResult: + loss: float + """Mean loss per unmasked token over the global batch.""" + + nan_count: int + """Micro-batches with a NaN loss over all ranks. They are left out of the loss and the backward pass.""" + + num_tokens: int + """Input tokens of the global batch.""" + + num_flops: float + """Training FLOPs of the global batch.""" + + tokens_per_expert: Tensor | None + """Global ``[num_moe_layers, num_experts]`` token counts of the step. None for dense models.""" + + metrics: dict[str, float] + """The data wait (max and mean over ranks) and, for MoE models, the routing statistics, to log.""" + + +def run_micro_batches( + model: nn.Module, + micro_batches: Iterator[dict], + compute_loss: Callable[[dict], tuple[Tensor, Tensor]], + grad_accum_steps: int, + parallel_dims: ParallelDims, + perf_counter: PerfCounter, + *, + gradient_manager: GradientOffloadManager | None = None, + defer_grad_sync: bool = False, +) -> StepResult: + """Forward and backward passes over the ``grad_accum_steps`` micro-batches of one step, leaving the gradients of + the summed token loss divided by the global token count. ``compute_loss`` returns the summed loss and the number + of unmasked tokens of a micro-batch. With a ``gradient_manager``, the caller has set the gradient scale up front. + With ``defer_grad_sync``, FSDP reduce-scatters the gradients only at the last micro-batch.""" + dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() + ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None + moes = list(iter_moe_blocks(model)) + + zero = torch.zeros((), device="cuda", dtype=torch.float64) + loss_sum, loss_tokens, nan_count = zero.clone(), zero.clone(), zero.clone() + num_tokens, num_flops, data_time = 0, 0, 0.0 + expert_counts, routing_confidence_sums = [], [] + for micro_step in range(grad_accum_steps): + data_start_time = time.perf_counter() + micro_batch = next(micro_batches) + data_time += time.perf_counter() - data_start_time + num_tokens += micro_batch["input_ids"].shape[1] + num_flops += perf_counter.flop_counter(micro_batch["seq_lens"].tolist()) + + with record_function("forward"): + micro_loss_sum, micro_token_count = compute_loss(micro_batch) + loss_tokens += micro_token_count + is_nan = torch.isnan(micro_loss_sum.detach()) + nan_count += is_nan + loss_sum += torch.where(is_nan, 0.0, micro_loss_sum.detach()) + scaled_loss = torch.where(is_nan, torch.nan_to_num(micro_loss_sum, nan=0.0), micro_loss_sum) / grad_accum_steps + + is_last_micro_step = micro_step == grad_accum_steps - 1 + if defer_grad_sync: + model.set_requires_gradient_sync(is_last_micro_step) + model.set_reshard_after_backward(is_last_micro_step) + with record_function("backward"): + begin_backward(gradient_manager, final_backward=is_last_micro_step) + scaled_loss.backward() + finish_backward(gradient_manager) + + if moes: + counts, confidence_sums = take_routing_counts(moes) + expert_counts.append(counts) + routing_confidence_sums.append(confidence_sums) + + # Every CP rank holds the same rows. + sums = { + "loss": loss_sum, + "loss_tokens": loss_tokens, + "nan_count": nan_count, + "tokens": zero + num_tokens / parallel_dims.cp, + "flops": zero + num_flops / parallel_dims.cp, + "data_time": zero + data_time, + } + maxes = {"data_time_max": zero + data_time} + if moes: + counts, confidence_sums = torch.stack(expert_counts), torch.stack(routing_confidence_sums) + sums["expert_counts"] = counts.sum(dim=0) + # Max violation and routing confidence of every micro-batch, over the EP group's tokens. + if ep_group is not None: + packed = torch.cat([counts, confidence_sums.unsqueeze(-1)], dim=-1) + dist.all_reduce(packed, group=ep_group) + counts, confidence_sums = packed[..., :-1], packed[..., -1] + max_vio, routing_confidence = compute_max_vio_and_routing_confidence( + counts, confidence_sums, moes[0].router.top_k + ) + maxes["max_vio"] = max_vio.amax(dim=-1) + sums["routing_confidence"] = routing_confidence.mean(dim=-1) + stats = reduce_stats(sums, maxes, dp_cp_group) + + global_token_count = int(stats["loss_tokens"]) + if gradient_manager is None and global_token_count > 0: + scale_gradients_(None, model, parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count) + + num_ranks = dist.get_world_size(dp_cp_group) + metrics = {"time/data/max": stats["data_time_max"].item(), "time/data/mean": stats["data_time"].item() / num_ranks} + tokens_per_expert = None + if moes: + tokens_per_expert = stats["expert_counts"] + metrics |= expert_load_metrics(tokens_per_expert, moe_layer_indices(model)) | { + "max_vio/mean": stats["max_vio"].mean().item(), + "max_vio/max": stats["max_vio"].max().item(), + "routing_confidence/mean": stats["routing_confidence"].mean().item() / num_ranks, + } + return StepResult( + loss=stats["loss"].item() / global_token_count if global_token_count > 0 else 0.0, + nan_count=int(stats["nan_count"]), + num_tokens=round(stats["tokens"].item()), + num_flops=stats["flops"].item(), + tokens_per_expert=tokens_per_expert, + metrics=metrics, + ) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index eb008e49acb..64aac93649d 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -206,6 +206,18 @@ def is_tt_moe_model(model: nn.Module) -> bool: return hasattr(config, "num_experts") or hasattr(config, "n_routed_experts") +def compute_max_vio_and_routing_confidence( + tokens_per_expert: Tensor, routing_confidence_sum: Tensor, top_k: int +) -> tuple[Tensor, Tensor]: + """Max violation and routing confidence of one MoE layer's expert counts (last dimension) and confidence sum. + Leading dimensions are batch dimensions.""" + num_routed_tokens = tokens_per_expert.sum(dim=-1) / top_k + loads = tokens_per_expert.sort(dim=-1, descending=True).values[..., top_k:] + balanced_load = loads.mean(dim=-1) + max_vio = (loads.amax(dim=-1) - balanced_load) / balanced_load + return max_vio, routing_confidence_sum / num_routed_tokens + + def get_load_balance_stats( model: nn.Module, group: dist.ProcessGroup | None = None, @@ -231,13 +243,10 @@ def get_load_balance_stats( layer_stats = [(stats[:-1], stats[-1]) for stats in packed_stats.split(sizes)] for block_mlp, (tokens_per_expert, routing_confidence_sum) in zip(block_mlps, layer_stats): - num_routed_tokens = tokens_per_expert.sum() / block_mlp.router.top_k - tokens_per_expert = tokens_per_expert.sort(dim=0, descending=True).values[block_mlp.router.top_k :] - balanced_load = tokens_per_expert.mean() - max_vio = (tokens_per_expert.max() - balanced_load) / balanced_load + max_vio, routing_confidence = compute_max_vio_and_routing_confidence( + tokens_per_expert, routing_confidence_sum, block_mlp.router.top_k + ) per_layer_max_vio.append(max_vio.detach()) - - routing_confidence = routing_confidence_sum / num_routed_tokens per_layer_routing_confidence.append(routing_confidence.detach()) block_mlp.tokens_per_expert.zero_() diff --git a/src/prime_rl/trainer/models/base.py b/src/prime_rl/trainer/models/base.py index a9e79685636..2c61d0bdd54 100644 --- a/src/prime_rl/trainer/models/base.py +++ b/src/prime_rl/trainer/models/base.py @@ -108,6 +108,24 @@ def num_active_params(self) -> int: """Parameters a text token multiplies through, LoRA adapters included.""" return count_active_params(self) + def embedding_params(self) -> list[nn.Parameter]: + """The embedding tables (every ``nn.Embedding``, as in `count_active_params`) and the lm_head.""" + params = [p for m in self.modules() if isinstance(m, nn.Embedding) for p in m.parameters(recurse=False)] + return list({id(p): p for p in [*params, *self.lm_head.parameters()]}.values()) + + def num_params(self) -> int: + return sum(p.numel() for p in self.parameters()) + + def num_embedding_params(self) -> int: + return sum(p.numel() for p in self.embedding_params()) + + def num_non_embedding_params(self) -> int: + return self.num_params() - self.num_embedding_params() + + def num_active_non_embedding_params(self) -> int: + embedding = {id(p) for p in self.embedding_params()} + return count_active_params(self, [p for p in self.parameters() if id(p) not in embedding]) + @classmethod def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: """ diff --git a/src/prime_rl/trainer/pretrain/bytes.py b/src/prime_rl/trainer/pretrain/bytes.py deleted file mode 100644 index 98f676695d7..00000000000 --- a/src/prime_rl/trainer/pretrain/bytes.py +++ /dev/null @@ -1,33 +0,0 @@ -import torch -from transformers.tokenization_utils import PreTrainedTokenizer - - -def bytes_to_unicode() -> dict[int, str]: - """The byte -> character map of byte-level BPE tokenizers (GPT-2, Llama 3, Qwen, DeepSeek).""" - printable = ( - list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1)) - ) - chars = printable[:] - n = 0 - for byte in range(256): - if byte not in printable: - printable.append(byte) - chars.append(256 + n) - n += 1 - return dict(zip(printable, map(chr, chars))) - - -def token_byte_lengths(tokenizer: PreTrainedTokenizer) -> torch.Tensor: - """UTF-8 bytes each token id stands for. Special tokens count as 0 bytes. Exact for - byte-level BPE tokenizers, approximate otherwise.""" - byte_chars = set(bytes_to_unicode().values()) - special_ids = set(tokenizer.all_special_ids) | set(tokenizer.added_tokens_decoder) - lengths = torch.zeros(len(tokenizer), dtype=torch.long) - for token_id, token in enumerate(tokenizer.convert_ids_to_tokens(list(range(len(tokenizer))))): - if token is None or token_id in special_ids: - continue - if all(char in byte_chars for char in token): - lengths[token_id] = len(token) - else: - lengths[token_id] = len(tokenizer.convert_tokens_to_string([token]).encode()) - return lengths diff --git a/src/prime_rl/trainer/pretrain/data.py b/src/prime_rl/trainer/pretrain/data.py deleted file mode 100644 index d94751f94cb..00000000000 --- a/src/prime_rl/trainer/pretrain/data.py +++ /dev/null @@ -1,254 +0,0 @@ -from pathlib import Path - -import numpy as np -import torch -from torchdata.stateful_dataloader import StatefulDataLoader - -from prime_rl.configs.pretrain import PretrainDataConfig -from prime_rl.trainer.sft.data import StatefulIterableDataset -from prime_rl.trainer.world import get_world -from prime_rl.utils.logger import get_logger - - -class TokenShards: - """Datatrove tokenized shards in one or more folders, the ``num_tokens`` first tokens of their files (in folder - order, then file name order), read as one cyclic token stream. - - ``.ds`` holds the token ids (little-endian, 2 or 4 bytes as recorded on the first line of - ``.ds.metadata``) and ``.ds.index`` the uint64 end offset of every document. File - boundaries are document boundaries. - - Each epoch reads all ``num_tokens // seq_len`` windows of the stream once, in a random order and from a random - offset drawn from ``seed`` and the epoch, so epochs differ in order and window boundaries, and part of an epoch is - a sample of the whole stream (``window_start``). The last ``num_tokens % seq_len`` tokens of an epoch are - skipped.""" - - def __init__(self, paths: list[Path], seed: int = 0): - self.files, self.first_files = [], [] - for path in paths: - files = sorted(path.glob("*.ds")) - if not files: - raise FileNotFoundError(f"No datatrove .ds files in {path}") - self.first_files.append(len(self.files)) - self.files += files - self.dtypes = [token_dtype(file) for file in self.files] - sizes = [file.stat().st_size // dtype.itemsize for file, dtype in zip(self.files, self.dtypes)] - self.starts = np.cumsum([0, *sizes]) - self.num_tokens = int(self.starts[-1]) - self.seed = seed - self.memmaps: dict[int, tuple[np.memmap, np.memmap]] = {} - # The current epoch, its window order and its offset. - self.epoch_order: tuple[int | None, np.ndarray | None, int] = (None, None, 0) - - def open(self, file_index: int) -> tuple[np.memmap, np.memmap]: - # Opened lazily, so every dataloader worker maps the files itself. - if file_index not in self.memmaps: - file = self.files[file_index] - tokens = np.memmap(file, dtype=self.dtypes[file_index], mode="r") - doc_ends = np.memmap(f"{file}.index", dtype=" tuple[int, int, int]: - """File index, offset in the file and length of the contiguous file run at stream ``position``, - at most ``length`` tokens.""" - position %= self.num_tokens - length = min(length, self.num_tokens - position) - file_index = int(np.searchsorted(self.starts, position, side="right")) - 1 - length = min(length, int(self.starts[file_index + 1]) - position) - return file_index, position - int(self.starts[file_index]), length - - def read(self, start: int, length: int) -> tuple[np.ndarray, list[int]]: - """Tokens ``[start, start + length)`` of the cyclic stream and the offsets in them where a document starts. - The wrap from the end of the stream to its start begins a new document.""" - pieces, doc_starts = [], [] - offset = 0 - while offset < length: - if offset > 0: - doc_starts.append(offset) - file_index, local, size = self.locate(start + offset, length - offset) - tokens, doc_ends = self.open(file_index) - pieces.append(tokens[local : local + size]) - first, last = np.searchsorted(doc_ends, [local, local + size], side="right") - doc_starts.extend(offset + int(end) - local for end in doc_ends[first:last] if end < local + size) - offset += size - return np.concatenate(pieces).astype(np.int64), doc_starts - - def window_start(self, position: int, seq_len: int) -> int: - """Stream offset of the window to read when ``position`` tokens have been read.""" - num_windows = self.num_tokens // seq_len - epoch, window = divmod(position // seq_len, num_windows) - if self.epoch_order[0] != epoch: - rng = np.random.default_rng([self.seed, epoch]) - order = rng.permutation(num_windows).astype(np.uint32) - self.epoch_order = epoch, order, int(rng.integers(self.num_tokens)) - _, order, offset = self.epoch_order - return (offset + int(order[window]) * seq_len) % self.num_tokens - - def first_documents(self) -> list[np.ndarray]: - """The first document of every folder.""" - documents = [] - for file_index in self.first_files: - tokens, doc_ends = self.open(file_index) - documents.append(np.asarray(tokens[: int(doc_ends[0])])) - return documents - - -def token_dtype(file: Path) -> np.dtype: - first_line = Path(f"{file}.metadata").read_text().splitlines()[0] - return np.dtype(f" tuple[np.ndarray, list[int]]: - return np.random.default_rng([self.seed, start]).integers(0, self.vocab_size, length), [] - - def window_start(self, position: int, seq_len: int) -> int: - return position - - -class TokenDataset(StatefulIterableDataset): - """``seq_len`` windows of a token source, window ``w`` read from ``source.window_start(w * seq_len)``. - - Every data rank walks the whole window sequence but reads only its own windows; the state is the window counter. - Each window reads one token past its end for the last target. Without ``document_masking``, a window is one - sequence: attention and positions run over the whole window.""" - - def __init__( - self, - source: TokenShards | RandomTokens, - seq_len: int, - non_dp_size: int = 1, - document_masking: bool = True, - ): - super().__init__(non_dp_size) - self.source = source - self.seq_len = seq_len - self.document_masking = document_masking - - def state_dict(self) -> dict: - return {"dataset": {"step": self.step}} - - def load_state_dict(self, state_dict: dict): - self.step = state_dict["dataset"]["step"] - - def __iter__(self): - self._setup_world_info() - while True: - window = self.step - self.step += 1 - if window % self.data_world_size != self.data_rank: - continue - start = self.source.window_start(window * self.seq_len, self.seq_len) - tokens, doc_starts = self.source.read(start, self.seq_len + 1) - doc_starts = [s for s in doc_starts if s < self.seq_len] if self.document_masking else [] - yield {"tokens": tokens, "seq_lens": np.diff([0, *doc_starts, self.seq_len])} - - -def collate_rows(rows: list[dict]) -> dict: - """Concatenate windows into one packed sample of shape ``[1, len(rows) * seq_len]``. Positions restart at - every document; attention never crosses a document or a window.""" - tokens = torch.from_numpy(np.stack([row["tokens"] for row in rows])) - seq_lens = torch.from_numpy(np.concatenate([row["seq_lens"] for row in rows])) - return { - "input_ids": tokens[:, :-1].reshape(1, -1), - "target_ids": tokens[:, 1:].reshape(1, -1), - "loss_mask": torch.ones((1, tokens.shape[0] * (tokens.shape[1] - 1)), dtype=torch.bool), - "position_ids": torch.cat([torch.arange(length) for length in seq_lens.tolist()]).unsqueeze(0), - "seq_lens": seq_lens, - "mm_kwargs": None, - "mm_token_type_ids": None, - } - - -class TrainData: - """The training data: windows of the token pool (or of fake tokens), their dataloader and the count of windows - trained on. Iterating yields this rank's micro-batches; every data rank walks the same window sequence, so the - count is global without communication. ``state_dict`` holds the dataloader position and the count, for - checkpoints.""" - - def __init__( - self, config: PretrainDataConfig, max_steps: int, vocab_size: int, eos_token_id: int, non_dp_size: int - ): - if config.type == "fake": - self.source = RandomTokens(vocab_size, config.seed) - else: - self.source = TokenShards(config.paths, seed=config.seed) - check_eos(self.source, eos_token_id) - epochs = max_steps * config.batch_size * config.seq_len / self.source.num_tokens - get_logger().info(f"Token pool: {self.source.num_tokens:,} tokens, {epochs:.3f} planned epochs") - if epochs > config.max_epochs: - raise ValueError(f"Planned epochs {epochs:.3f} exceed data.max_epochs={config.max_epochs}") - dataset = TokenDataset( - self.source, - config.seq_len, - non_dp_size=non_dp_size, - document_masking=config.type == "fake" or config.document_masking, - ) - self.loader = StatefulDataLoader( - dataset, - batch_size=config.micro_batch_size, - collate_fn=collate_rows, - num_workers=config.num_workers, - pin_memory=True, - ) - self.seq_len = config.seq_len - self.windows_per_micro_batch = config.micro_batch_size * get_world().world_size // non_dp_size - self.windows = 0 - self.iterator = None - - def __iter__(self): - return self - - def __next__(self) -> dict: - # Created at the first micro-batch, after a potential resume: iter() forks the workers with the dataset's - # current state. - if self.iterator is None: - self.iterator = iter(self.loader) - micro_batch = next(self.iterator) - self.windows += self.windows_per_micro_batch - return micro_batch - - def state_dict(self) -> dict: - return {"dataloader": self.loader.state_dict(), "windows": self.windows} - - def load_state_dict(self, state_dict: dict): - self.loader.load_state_dict(state_dict["dataloader"]) - self.windows = state_dict["windows"] - - def metrics(self) -> dict[str, float]: - return { - "progress/num_samples": self.windows, - "progress/epochs": self.windows * self.seq_len / self.source.num_tokens, - } - - -def check_eos(shards: TokenShards, eos_token_id: int) -> None: - """Every document ends with the tokenizer's EOS. A different token there means the shards were made with - another tokenizer. Checked on the first document of every folder.""" - for file_index, document in zip(shards.first_files, shards.first_documents()): - if document[-1] != eos_token_id: - raise ValueError( - f"The first document of {shards.files[file_index].parent} ends with token {document[-1]}, not the " - f"tokenizer's EOS {eos_token_id}. Were the shards tokenized with tokenizer.name?" - ) - - -def starts_with_bos(shards: TokenShards, bos_token_id: int | None) -> bool: - """Whether the documents start with the tokenizer's BOS, read off the first document of every folder. The eval - suite formats its requests the same way.""" - starts = { - shards.files[i].parent.as_posix(): int(doc[0]) == bos_token_id - for i, doc in zip(shards.first_files, shards.first_documents()) - } - if len(set(starts.values())) > 1: - raise ValueError(f"Only some token folders start their documents with BOS: {starts}") - return all(starts.values()) diff --git a/src/prime_rl/trainer/pretrain/step.py b/src/prime_rl/trainer/pretrain/step.py deleted file mode 100644 index f17921de2f4..00000000000 --- a/src/prime_rl/trainer/pretrain/step.py +++ /dev/null @@ -1,123 +0,0 @@ -"""The pretrain loop's forward and backward passes over one step's micro-batches, with the step's statistics -reduced once.""" - -import time -from collections.abc import Callable, Iterator -from dataclasses import dataclass - -import torch -import torch.distributed as dist -from torch import Tensor, nn - -from prime_rl.trainer.model import is_tt_moe_model, iter_moe_blocks -from prime_rl.trainer.parallel_dims import ParallelDims -from prime_rl.trainer.pretrain.metrics import expert_load_metrics, moe_layer_indices, reduce_stats, take_routing_counts -from prime_rl.trainer.utils import scale_gradients_ - - -@dataclass -class StepResult: - loss: float - """Mean loss per unmasked token over the global batch.""" - - nan_count: int - """Micro-batches with a NaN loss over all ranks. They are left out of the backward pass.""" - - num_tokens: int - """Input tokens of the global batch.""" - - tokens_per_expert: Tensor | None - """Global ``[num_moe_layers, num_experts]`` token counts of the step (fp32, on the GPU). None for dense models.""" - - metrics: dict[str, float] - """Statistics of the step over all ranks: data wait (max and mean over ranks), and for MoE - models the expert balance of the global counts (`expert_load_metrics`), the busiest expert of one rank's - micro-batch over the mean (``expert_load/micro_max_mean``, mean and max over ranks and micro-batches) and the - routing confidence.""" - - sums: dict[str, float] - """The caller's ``extra_sums``, summed over the ranks.""" - - -def run_micro_batches( - model: nn.Module, - dataiter: Iterator[dict], - compute_loss: Callable[[dict], tuple[Tensor, Tensor, Tensor]], - grad_accum_steps: int, - parallel_dims: ParallelDims, - *, - extra_sums: dict[str, Tensor] | None = None, -) -> StepResult: - """Forward and backward passes over the ``grad_accum_steps`` micro-batches of one step, leaving the - gradients of the summed token loss divided by the global token count (the mean loss per token, as on one - GPU without accumulation). ``compute_loss`` returns the summed loss and the number of unmasked tokens of a - micro-batch, and an auxiliary loss on the same scale as the summed loss that is trained but not logged. - - The step's statistics stay on the GPU during the micro-batches (no collective, no host sync) and are reduced - with one SUM and one MAX all-reduce at the end, together with ``extra_sums`` (GPU tensors the caller adds to - in ``compute_loss``).""" - dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() - is_moe_model = is_tt_moe_model(model) - - zero = torch.zeros((), device="cuda", dtype=torch.float64) - sums = {"loss": zero.clone(), "loss_tokens": zero.clone(), "nan_count": zero.clone(), "tokens": zero.clone()} - maxes = {} - data_time = 0.0 - for micro_step in range(grad_accum_steps): - data_start_time = time.perf_counter() - micro_batch = next(dataiter) - data_time += time.perf_counter() - data_start_time - sums["tokens"] += micro_batch["input_ids"].shape[1] - micro_loss_sum, micro_token_count, aux_loss_sum = compute_loss(micro_batch) - sums["loss_tokens"] += micro_token_count - # A NaN micro-batch is left out of the loss and the backward, without a host sync. - is_nan = torch.isnan(micro_loss_sum.detach()) - sums["nan_count"] += is_nan - sums["loss"] += torch.where(is_nan, 0.0, micro_loss_sum.detach()) - scaled_loss = torch.where(is_nan, torch.nan_to_num(micro_loss_sum, nan=0.0), micro_loss_sum + aux_loss_sum) - is_last_micro_step = micro_step == grad_accum_steps - 1 - # Reduce-scatter (and reshard) once per step, not per micro-batch: the expert gradients cross nodes. - model.set_requires_gradient_sync(is_last_micro_step) - model.set_reshard_after_backward(is_last_micro_step) - (scaled_loss / grad_accum_steps).backward() - - if is_moe_model: - counts, confidence = take_routing_counts(model) - sums["expert_counts"] = sums.get("expert_counts", zero) + counts - sums["routing_confidence"] = sums.get("routing_confidence", zero) + confidence - micro_max_mean = counts.amax(dim=1) / counts.mean(dim=1) - sums["micro_max_mean"] = sums.get("micro_max_mean", zero) + micro_max_mean.mean() - maxes["micro_max_mean_max"] = torch.maximum(maxes.get("micro_max_mean_max", zero), micro_max_mean.amax()) - - sums["data_time"] = zero + data_time - maxes["data_time_max"] = zero + data_time - extra_sums = extra_sums or {} - stats = reduce_stats(sums | {f"extra/{name}": value for name, value in extra_sums.items()}, maxes, dp_cp_group) - - global_token_count = int(stats["loss_tokens"]) - scale_gradients_(None, model, parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count) - - num_ranks = dist.get_world_size(dp_cp_group) - metrics = { - "time/data/max": stats["data_time_max"].item(), - "time/data/mean": stats["data_time"].item() / num_ranks, - } - tokens_per_expert = None - if is_moe_model: - counts = stats["expert_counts"] - routed_tokens = counts.sum(dim=1) / next(iter_moe_blocks(model)).router.top_k - metrics |= expert_load_metrics(counts, moe_layer_indices(model)) | { - "expert_load/micro_max_mean/mean": stats["micro_max_mean"].item() / (grad_accum_steps * num_ranks), - "expert_load/micro_max_mean/max": stats["micro_max_mean_max"].item(), - "router/confidence": (stats["routing_confidence"] / routed_tokens).mean().item(), - } - tokens_per_expert = counts.to("cuda", torch.float32) - return StepResult( - loss=stats["loss"].item() / global_token_count, - nan_count=int(stats["nan_count"]), - # Every CP rank counts its rows in full. - num_tokens=int(stats["tokens"]) // parallel_dims.cp, - tokens_per_expert=tokens_per_expert, - metrics=metrics, - sums={name: stats[f"extra/{name}"].item() for name in extra_sums}, - ) diff --git a/src/prime_rl/trainer/pretrain/train.py b/src/prime_rl/trainer/pretrain/train.py deleted file mode 100644 index a2f2cc7996e..00000000000 --- a/src/prime_rl/trainer/pretrain/train.py +++ /dev/null @@ -1,430 +0,0 @@ -import prime_rl._compat # noqa: F401 — patch ring_flash_attn compat before import - -import asyncio -import json -import math -import time -from datetime import timedelta - -# ruff: noqa: I001 - -from prime_rl.utils.act_offloading import maybe_activation_offloading -import numpy as np -import torch -import torch.distributed as dist -from torch.distributed.tensor import DTensor - -from prime_rl import monitors -from prime_rl.configs.pretrain import PretrainConfig -from prime_rl.configs.trainer import CheckpointConfig -from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager -from prime_rl.trainer.model import ( - forward, - is_tt_moe_model, - iter_moe_blocks, - reshard_module, - setup_model, - setup_tokenizer, -) -from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX -from prime_rl.trainer.optim import setup_optimizer -from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep -from prime_rl.trainer.perf import get_perf_counter -from prime_rl.trainer.pretrain.balance import ( - histogram_quantile_update_selection_biases, - quantile_update_selection_biases, - setup_moe_balance, - update_selection_biases, -) -from prime_rl.trainer.pretrain.bytes import token_byte_lengths -from prime_rl.trainer.pretrain.optim import param_groups -from prime_rl.trainer.pretrain.data import ( - TrainData, - starts_with_bos, -) -from prime_rl.trainer.pretrain.suite import ( - length_label, - prepare_multi_length_set, - score_sets, - set_metrics, - scoring_head_active, - summarize, -) -from prime_rl.trainer.pretrain.step import run_micro_batches -from prime_rl.trainer.pretrain.metrics import ( - DeadExperts, - SpikeDetector, - grad_norms_by_group, - reduce_stats, - routed_k_metrics, - take_routing_counts, -) -from prime_rl.trainer.sft.data import get_dataset_state -from prime_rl.trainer.scheduler import setup_scheduler -from prime_rl.trainer.utils import ( - count_parameters, - get_grad_accum_steps, - get_dcp_group, - GarbageCollection, - clip_grad_norm_, - get_ckpt_disk_metrics, - setup_torch_distributed, -) -from prime_rl.trainer.world import get_world -from prime_rl.utils.config import cli -from prime_rl.utils.cp import setup_context_parallel, setup_cp_params, shard_for_cp -from prime_rl.utils.heartbeat import Heartbeat -from prime_rl.utils.logger import setup_logger -from prime_rl.utils.pathing import resolve_resume_step -from prime_rl.utils.process import set_proc_title -from prime_rl.utils.utils import clean_exit - - -MOE_TRACE_INTERVAL = 32 - - -@clean_exit -def train(config: PretrainConfig): - world = get_world() - logger = setup_logger(config.log.level, json_logging=config.log.json_logging) - logger.info(f"Starting pretrain trainer in {world} (output_dir={config.run_dir})") - - asyncio.run( - monitors.setup( - producer="trainer", - wandb=config.monitors.wandb, - prime=config.monitors.prime, - file=config.monitors.file, - output_dir=config.run_dir, - run_config=config, - overview_flavor="pretrain", - ) - ) - - heart = Heartbeat(config.heartbeat) if config.heartbeat is not None and world.rank == 0 else None - - setup_torch_distributed(timeout=timedelta(seconds=config.dist_timeout_seconds)) - torch.set_float32_matmul_precision(config.matmul_precision) - - resolve_ep(config.model) - parallel_dims = get_parallel_dims(config.model) - - grad_accum_steps = get_grad_accum_steps( - config.data.batch_size, config.data.micro_batch_size, config.model.cp, world.world_size - ) - - ckpt_manager = setup_ckpt_manager(config.run_dir, config.ckpt, resume=config.resume) - checkpoint_step = resolve_resume_step(config.resume, ckpt_manager.ckpt_dir) - - logger.info(f"Initializing model ({config.model})") - # Weights come from the checkpoint when resuming and from the random init otherwise. - model = setup_model( - config.model, - parallel_dims, - loading_from_checkpoint_later=checkpoint_step is not None, - init=config.model.init if checkpoint_step is None else None, - ) - if parallel_dims.cp_enabled: - setup_context_parallel(model, config.model, parallel_dims) - num_params = count_parameters(model) - logger.info( - f"Parameters: {num_params['active_non_embedding'] / 1e6:,.1f}M active non-embedding, " - f"{num_params['non_embedding'] / 1e6:,.1f}M non-embedding, {num_params['embedding'] / 1e6:,.1f}M embedding, " - f"{num_params['total'] / 1e6:,.1f}M total" - ) - if checkpoint_step is None: - asyncio.run(monitors.log({f"params/{name}": count for name, count in num_params.items()} | {"step": 0}, step=0)) - - tokenizer = setup_tokenizer(config.tokenizer) - - optimizer, _ = setup_optimizer( - config.optim, - list(model.named_parameters()), - parallel_dims, - offload=config.model.offload, - model=model, - param_groups=param_groups(model, config.optim, parallel_dims), - ) - - # With ckpt.skip_scheduler, the schedule is rebuilt over the remaining steps. - skip_scheduler = (config.ckpt or CheckpointConfig()).skip_scheduler - scheduler_steps = config.max_steps - checkpoint_step if skip_scheduler and checkpoint_step else config.max_steps - scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) - - data = TrainData(config.data, config.max_steps, len(tokenizer), tokenizer.eos_token_id, config.model.cp) - - progress = Progress() - restore_dataloader = checkpoint_step is not None and not (config.ckpt or CheckpointConfig()).skip_dataloader - if checkpoint_step is not None: - ckpt_manager.load( - checkpoint_step, - model, - [optimizer], - None if skip_scheduler else scheduler, - progress, - dataloader=data if restore_dataloader else None, - path=config.resume.dir / "trainer" if config.resume.dir is not None else None, - ) - progress.step += 1 - if skip_scheduler: - # Loading the optimizer state replaces the param groups the scheduler was built on. - scheduler = setup_scheduler(optimizer, config.scheduler, scheduler_steps, config.optim.lr) - logger.info(f"Resuming from step {checkpoint_step} (dataset_state={get_dataset_state(data.loader)})") - - cp_enabled = parallel_dims.cp_enabled - cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() if cp_enabled else 0 - cp_size = parallel_dims.cp - dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() - is_moe_model = is_tt_moe_model(model) - balance = config.model.moe.balance - moes_with_seq_balance = setup_moe_balance(model, balance, config.data.micro_batch_size) - bias_freeze_step = config.max_steps - round(balance.bias_freeze * config.max_steps) - # Sum and count of the per-row balance of every MoE layer over the current step's micro-batches, reduced with - # the step's statistics. - seq_balance_sums: dict[str, torch.Tensor] = {} - - def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Forward pass returning (loss_sum, token_count) over unmasked tokens and the MoE - sequence-wise balance loss summed over tokens (each token carries its row's loss, so after the division - by the token count it is the mean over rows).""" - input_ids = micro_batch["input_ids"].to("cuda", non_blocking=True) - position_ids = micro_batch["position_ids"].to("cuda", non_blocking=True) - target_ids = micro_batch["target_ids"].to("cuda", non_blocking=True) - loss_mask = micro_batch["loss_mask"].to("cuda", non_blocking=True) - seq_lens = micro_batch["seq_lens"].to("cuda", non_blocking=True) - if cp_enabled: - input_ids, position_ids = setup_cp_params(input_ids, position_ids, cp_rank, cp_size, seq_lens=seq_lens) - target_ids = shard_for_cp(target_ids, cp_rank=cp_rank, cp_world_size=cp_size) - loss_mask = shard_for_cp(loss_mask, cp_rank=cp_rank, cp_world_size=cp_size) - labels = target_ids.masked_fill(~loss_mask, IGNORE_INDEX) - with maybe_activation_offloading(config.model.ac_offloading): - out = forward( - model, input_ids, position_ids, seq_lens=seq_lens, labels=labels, seq_lens_are_pre_shard=cp_enabled - ) - aux_loss_sum = torch.zeros((), device="cuda") - row_len = input_ids.shape[1] // config.data.micro_batch_size - for moe in moes_with_seq_balance: - aux_loss_sum += balance.seq_aux_coeff * moe.seq_balance.sum() * row_len - seq_balance_sums["balance"] = seq_balance_sums.get("balance", 0) + moe.seq_balance.detach().sum() - seq_balance_sums["rows"] = seq_balance_sums.get("rows", 0) + moe.seq_balance.new_tensor( - moe.seq_balance.numel() - ) - return out["loss"], loss_mask.sum(dtype=torch.int64), aux_loss_sum - - def finish_no_grad_forwards() -> None: - """No-grad forwards leave the [lm_head, norm] FSDP group unsharded, and their routing must not count toward - the next step's MoE statistics.""" - reshard_module(model) - if is_moe_model: - take_routing_counts(model) - - eval_sets = [] - if config.evals is not None: - start_time = time.perf_counter() - bos = config.data.type == "tokens" and starts_with_bos(data.source, tokenizer.bos_token_id) - has_token_sets = any(not isinstance(eval_set, str) for eval_set in config.evals.sets) - byte_lengths = token_byte_lengths(tokenizer).numpy() if has_token_sets else None - eval_sets = [ - prepare_multi_length_set( - eval_set, config.evals.data_dir, tokenizer, config.evals.seq_lens, bos, byte_lengths - ) - for eval_set in config.evals.sets - ] - logger.info(f"Prepared {len(eval_sets)} eval sets (BOS: {bos}) in {time.perf_counter() - start_time:.0f}s") - - @torch.no_grad() - def run_suite(step: int) -> None: - """The eval suite on the weights after ``step`` optimizer updates: the packed rows of all sets are dealt - over the ranks, the per-token results gathered on rank 0, which summarizes and logs them.""" - start_time = time.perf_counter() - metrics = {} - # Eager, also for compiled layers: the packed rows vary in length and would recompile at every forward. - with scoring_head_active(model), torch.compiler.set_stance("force_eager"): - scored = score_sets(model, eval_sets, config.evals.micro_batch_tokens, get_dcp_group()) - for eval_set, (logprobs, hits) in zip(eval_sets, scored or []): - for seq_len, version in eval_set.versions.items(): - index = eval_set.index[seq_len] - summary, arrays = summarize(version, [logprobs[i] for i in index], [hits[i] for i in index]) - label = length_label(seq_len) - metrics |= set_metrics(eval_set.name, seq_len, summary) - if config.evals.save_arrays: - out = config.run_dir / "evals" / f"step_{step}" - out.mkdir(parents=True, exist_ok=True) - np.savez(out / f"{eval_set.name}@{label}.npz", **arrays) - (out / f"{eval_set.name}@{label}.json").write_text(json.dumps(summary, default=float)) - finish_no_grad_forwards() - metrics["time/eval_suite"] = time.perf_counter() - start_time - logger.success(f"Eval suite | Step {step} | {metrics['time/eval_suite']:.0f}s | {len(eval_sets)} sets") - asyncio.run(monitors.log({**metrics, "step": step}, step=step)) - - if config.evals is not None and config.evals.eval_on_start: - run_suite(step=checkpoint_step or 0) - if checkpoint_step == config.max_steps: - # The last step's checkpoint is written before its evals: resuming from it only runs those. - logger.info(f"Resuming from the last step ({checkpoint_step}): running its evals only") - if config.evals is not None and not config.evals.eval_on_start: - run_suite(checkpoint_step) - logger.success("Pretrain trainer finished") - asyncio.run(monitors.finalize()) - return - - gc_handler = GarbageCollection(config.gc.interval) if config.gc else None - max_memory = torch.cuda.mem_get_info()[1] / 1024**3 - dead_experts = DeadExperts() - loss_spikes, grad_norm_spikes = SpikeDetector(), SpikeDetector() - last_bias_moves = None - # Per-expert selection bias, router-row norm and load of every MoE layer, every MOE_TRACE_INTERVAL steps. - moe_trace: dict[str, list] = {"step": [], "bias": [], "router_norm": [], "load": []} - logger.info(f"Starting training loop (max_steps={config.max_steps})") - while True: - torch.cuda.reset_peak_memory_stats() - if gc_handler is not None: - gc_handler.run(progress.step) - is_last_step = progress.step >= config.max_steps - step_start_time = time.perf_counter() - - seq_balance_sums.clear() - result = run_micro_batches( - model, - data, - compute_loss, - grad_accum_steps, - parallel_dims, - extra_sums=seq_balance_sums, - ) - forward_backward_time = time.perf_counter() - step_start_time - moe_metrics = {} - if is_moe_model: - moe_metrics["expert_load/dead"] = dead_experts.update(result.tokens_per_expert.cpu()) - if result.sums: - moe_metrics["loss/seq_balance"] = result.sums["balance"] / result.sums["rows"] - biases = [ - moe.router.selection_bias for moe in iter_moe_blocks(model) if moe.router.selection_bias is not None - ] - updates_bias = balance.bias_update == "quantile" or balance.bias_update_rate > 0 - if updates_bias and progress.step <= bias_freeze_step: - previous_biases = [bias.clone() for bias in biases] - if balance.bias_update == "quantile" and balance.quantile_estimator == "histogram": - moe_metrics.update(histogram_quantile_update_selection_biases(model, dp_cp_group)) - elif balance.bias_update == "quantile": - quantile_update_selection_biases(model, dp_cp_group) - else: - update_selection_biases(model, result.tokens_per_expert, balance.bias_update_rate) - if biases: - # Share of experts whose bias moved in the opposite direction to the previous step's move. - moves = torch.stack([bias - previous for bias, previous in zip(biases, previous_biases)]).sign() - if last_bias_moves is not None: - moe_metrics["router/bias/flip_frac"] = (moves * last_bias_moves < 0).float().mean().item() - last_bias_moves = moves - if biases: - stacked = torch.stack(biases) - moe_metrics["router/bias/absmax"] = stacked.abs().max().item() - moe_metrics["router/bias/std"] = stacked.std(dim=1).mean().item() - if progress.step % MOE_TRACE_INTERVAL == 0: - gates = [ - moe.router.gate.weight - for moe in iter_moe_blocks(model) - if moe.router.selection_bias is not None - ] - gates = [gate.full_tensor() if isinstance(gate, DTensor) else gate for gate in gates] - moe_trace["step"].append(progress.step) - moe_trace["bias"].append(stacked.float().cpu().numpy()) - moe_trace["router_norm"].append( - torch.stack([gate.detach().float().norm(dim=1) for gate in gates]).cpu().numpy() - ) - moe_trace["load"].append(result.tokens_per_expert.float().cpu().numpy()) - moe_metrics |= routed_k_metrics(model, dp_cp_group) - - group_grad_norms = grad_norms_by_group(model) - grad_norm = clip_grad_norm_(None, model, config.optim.max_norm) if config.optim.max_norm else None - optimizer.step() - optimizer.zero_grad() - current_lr = optimizer.param_groups[0]["lr"] - scheduler.step() - # Peak memory of the step's forward, backward and optimizer passes, over all ranks. - local_peak_memory = torch.tensor(torch.cuda.max_memory_reserved() / 1024**3, device="cuda") - peak_memory = reduce_stats({"sum": local_peak_memory}, {"max": local_peak_memory}, dp_cp_group) - peak_memory_metrics = { - "perf/peak_memory/max": peak_memory["max"].item(), - "perf/peak_memory/mean": peak_memory["sum"].item() / dist.get_world_size(dp_cp_group), - } - - num_tokens = result.num_tokens - progress.total_tokens += num_tokens - progress.total_samples = data.windows - - # The last step's checkpoint is written before its evals, so a timeout there does not lose it. - save_ckpt_time = 0.0 - disk_metrics = {} - if config.ckpt and (is_last_step or (config.ckpt.interval and progress.step % config.ckpt.interval == 0)): - save_ckpt_start_time = time.perf_counter() - ckpt_manager.save(progress.step, model, [optimizer], scheduler, progress, dataloader=data) - ckpt_manager.maybe_clean() - save_ckpt_time = time.perf_counter() - save_ckpt_start_time - disk_metrics = get_ckpt_disk_metrics(config.run_dir) - - eval_start_time = time.perf_counter() - if config.evals is not None and (is_last_step or progress.step % config.evals.interval == 0): - run_suite(progress.step) - eval_time = time.perf_counter() - eval_start_time - - step_time = time.perf_counter() - step_start_time - eval_time - perf_counter = get_perf_counter(model, config.data.seq_len) - perf_counter.count_tokens(num_tokens) - throughput = perf_counter.get_step_tokens_per_second(num_tokens, step_time) - mfu = perf_counter.get_step_mfu(num_tokens, step_time) - step_message = f"Step {progress.step} | {step_time:.2f}s | Loss {result.loss:.4f}" - if grad_norm is not None: - step_message += f" | Grad. Norm {grad_norm:.4f}" - step_message += ( - f" | LR {current_lr:.2e} | Throughput {throughput:.0f} tokens/s | MFU {mfu:.1f}%" - f" | Peak Mem. {peak_memory_metrics['perf/peak_memory/max']:.1f}/{max_memory:.1f} GiB" - ) - logger.success(step_message) - - metrics = { - "loss/mean": result.loss, - "loss/perplexity": math.exp(min(result.loss, 20)), - "loss/nan_count": result.nan_count, - "loss/spike": loss_spikes.update(result.loss), - "optim/lr": current_lr, - **{f"optim/grad_norm/{group}": norm.item() for group, norm in group_grad_norms.items()}, - **data.metrics(), - "progress/num_tokens": progress.total_tokens, - "perf/throughput": throughput, - "perf/throughput_per_gpu": throughput / world.world_size, - "perf/mfu": mfu, - "time/step": step_time, - "time/forward_backward": forward_backward_time, - "time/save_ckpt": save_ckpt_time, - "time/evals": eval_time, - **disk_metrics, - **peak_memory_metrics, - **result.metrics, - **moe_metrics, - "step": progress.step, - } - if grad_norm is not None: - metrics["optim/grad_norm"] = grad_norm.item() - metrics["optim/grad_norm_spike"] = grad_norm_spikes.update(grad_norm.item()) - asyncio.run(monitors.log(metrics, step=progress.step)) - - if heart is not None: - heart.beat() - if moe_trace["step"] and world.rank == 0 and (is_last_step or progress.step % (16 * MOE_TRACE_INTERVAL) == 0): - np.savez(config.run_dir / "moe_trace.npz", **{name: np.array(values) for name, values in moe_trace.items()}) - if is_last_step: - break - progress.step += 1 - - logger.success("Pretrain trainer finished") - asyncio.run(monitors.finalize()) - - -def main(): - set_proc_title("PretrainTrainer") - train(cli(PretrainConfig)) - - -if __name__ == "__main__": - main() diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 313caca9739..2676ffdf4d4 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -75,7 +75,7 @@ from prime_rl.control.server import ControlServer from prime_rl.utils.process import set_proc_title from prime_rl.utils.worker_pool import WorkerPool -from prime_rl.utils.pathing import resolve_resume_step, write_resolved_config +from prime_rl.utils.pathing import resolve_latest_ckpt_step, write_resolved_config from prime_rl.utils.utils import clean_exit @@ -127,10 +127,17 @@ def train(config: TrainerConfig): parallel_dims = get_parallel_dims(config.model) # Check for checkpoint to resume from + checkpoint_step = None logger.info(f"Initializing checkpoint manager ({config.ckpt})") ckpt_manager = setup_ckpt_manager(config.output_dir, config.ckpt, resume=config.resume) - checkpoint_step = resolve_resume_step(config.resume, ckpt_manager.ckpt_dir) + if config.resume is not None: + if config.resume.dir is not None: + checkpoint_step = config.resume.dir_step + else: + checkpoint_step = config.resume.step + if checkpoint_step is None: + checkpoint_step = resolve_latest_ckpt_step(ckpt_manager.ckpt_dir) # Initialize the model logger.info(f"Initializing model ({config.model})") diff --git a/src/prime_rl/trainer/sft/data.py b/src/prime_rl/trainer/sft/data.py index d7eb056014c..e83b0d3beb5 100644 --- a/src/prime_rl/trainer/sft/data.py +++ b/src/prime_rl/trainer/sft/data.py @@ -480,6 +480,133 @@ def __iter__(self): yield processed_example +class TokenShards: + """Datatrove tokenized shards in one or more folders, the ``num_tokens`` first tokens of their files (in folder + order, then file name order), read as one cyclic token stream. + + ``.ds`` holds the token ids (little-endian, 2 or 4 bytes as recorded on the first line of + ``.ds.metadata``) and ``.ds.index`` the uint64 end offset of every document. File + boundaries are document boundaries. + + Each epoch reads all ``num_tokens // seq_len`` windows of the stream once, in a random order and from a random + offset drawn from ``seed`` and the epoch, so epochs differ in order and window boundaries, and part of an epoch is + a sample of the whole stream (``window_start``). The last ``num_tokens % seq_len`` tokens of an epoch are + skipped.""" + + def __init__(self, paths: list[Path], seed: int = 0): + self.files, self.first_files = [], [] + for path in paths: + files = sorted(path.glob("*.ds")) + if not files: + raise FileNotFoundError(f"No datatrove .ds files in {path}") + self.first_files.append(len(self.files)) + self.files += files + self.dtypes = [token_dtype(file) for file in self.files] + sizes = [file.stat().st_size // dtype.itemsize for file, dtype in zip(self.files, self.dtypes)] + self.starts = np.cumsum([0, *sizes]) + self.num_tokens = int(self.starts[-1]) + self.seed = seed + self.memmaps: dict[int, tuple[np.memmap, np.memmap]] = {} + # The current epoch, its window order and its offset. + self.epoch_order: tuple[int | None, np.ndarray | None, int] = (None, None, 0) + + def open(self, file_index: int) -> tuple[np.memmap, np.memmap]: + # Opened lazily, so every dataloader worker maps the files itself. + if file_index not in self.memmaps: + file = self.files[file_index] + tokens = np.memmap(file, dtype=self.dtypes[file_index], mode="r") + doc_ends = np.memmap(f"{file}.index", dtype=" tuple[int, int, int]: + """File index, offset in the file and length of the contiguous file run at stream ``position``, + at most ``length`` tokens.""" + position %= self.num_tokens + length = min(length, self.num_tokens - position) + file_index = int(np.searchsorted(self.starts, position, side="right")) - 1 + length = min(length, int(self.starts[file_index + 1]) - position) + return file_index, position - int(self.starts[file_index]), length + + def read(self, start: int, length: int) -> tuple[np.ndarray, list[int]]: + """Tokens ``[start, start + length)`` of the cyclic stream and the offsets in them where a document starts. + The wrap from the end of the stream to its start begins a new document.""" + pieces, doc_starts = [], [] + offset = 0 + while offset < length: + if offset > 0: + doc_starts.append(offset) + file_index, local, size = self.locate(start + offset, length - offset) + tokens, doc_ends = self.open(file_index) + pieces.append(tokens[local : local + size]) + first, last = np.searchsorted(doc_ends, [local, local + size], side="right") + doc_starts.extend(offset + int(end) - local for end in doc_ends[first:last] if end < local + size) + offset += size + return np.concatenate(pieces).astype(np.int64), doc_starts + + def window_start(self, position: int, seq_len: int) -> int: + """Stream offset of the window to read when ``position`` tokens have been read.""" + num_windows = self.num_tokens // seq_len + epoch, window = divmod(position // seq_len, num_windows) + if self.epoch_order[0] != epoch: + rng = np.random.default_rng([self.seed, epoch]) + order = rng.permutation(num_windows).astype(np.uint32) + self.epoch_order = epoch, order, int(rng.integers(self.num_tokens)) + _, order, offset = self.epoch_order + return (offset + int(order[window]) * seq_len) % self.num_tokens + + def first_documents(self) -> list[np.ndarray]: + """The first document of every folder.""" + documents = [] + for file_index in self.first_files: + tokens, doc_ends = self.open(file_index) + documents.append(np.asarray(tokens[: int(doc_ends[0])])) + return documents + + +def token_dtype(file: Path) -> np.dtype: + first_line = Path(f"{file}.metadata").read_text().splitlines()[0] + return np.dtype(f" dict: "num_samples": dict(num_samples), "num_tokens": dict(num_tokens), } + + +def check_eos(shards: TokenShards, eos_token_id: int) -> None: + """Every document ends with the tokenizer's EOS. A different token there means the shards were made with + another tokenizer. Checked on the first document of every folder.""" + for file_index, document in zip(shards.first_files, shards.first_documents()): + if document[-1] != eos_token_id: + raise ValueError( + f"The first document of {shards.files[file_index].parent} ends with token {document[-1]}, not the " + f"tokenizer's EOS {eos_token_id}. Were the shards tokenized with tokenizer.name?" + ) + + +def starts_with_bos(shards: TokenShards, bos_token_id: int | None) -> bool: + """Whether the documents start with the tokenizer's BOS, read off the first document of every folder. The eval + suite formats its requests the same way.""" + starts = { + shards.files[i].parent.as_posix(): int(doc[0]) == bos_token_id + for i, doc in zip(shards.first_files, shards.first_documents()) + } + if len(set(starts.values())) > 1: + raise ValueError(f"Only some token folders start their documents with BOS: {starts}") + return all(starts.values()) diff --git a/src/prime_rl/trainer/pretrain/suite.py b/src/prime_rl/trainer/sft/ll_eval.py similarity index 80% rename from src/prime_rl/trainer/pretrain/suite.py rename to src/prime_rl/trainer/sft/ll_eval.py index 7692e91ee3e..e5e7c422875 100644 --- a/src/prime_rl/trainer/pretrain/suite.py +++ b/src/prime_rl/trainer/sft/ll_eval.py @@ -1,5 +1,5 @@ -"""The pretraining eval suite: log-likelihood scoring of text eval sets, as in the offline studies -(`pretrain-evals` offline suite). A set lives in ``//`` as ``test.parquet`` plus ``meta.json`` +"""Log-likelihood evals inside the trainer (``[ll_eval]``): scoring of text eval sets, as in the +`pretrain-evals` offline suite. A set lives in ``//`` as ``test.parquet`` plus ``meta.json`` (``kind``, ``parquet_sha256``). Kinds: Requests start with BOS only if the training documents do (``bos``), so evals see the training format. @@ -16,10 +16,11 @@ as stored, cut into windows the same way; bpb counts the bytes of the scored tokens. A set is tokenized once (`prepare_set`); every round scores its requests (`score_requests`, the model's LM head -replaced by `scoring_head`) and summarizes them (`summarize`).""" +replaced by `scoring_head`) and summarizes them (`summarize`). `LogLikelihoodEval` runs the rounds.""" import json import math +import time import types from collections.abc import Iterator from contextlib import contextmanager @@ -34,10 +35,12 @@ from torch import nn from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.configs.pretrain import TokenEvalSetConfig +from prime_rl.configs.sft import LogLikelihoodEvalConfig, TokenEvalSetConfig from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX -from prime_rl.trainer.pretrain.data import TokenShards, check_eos +from prime_rl.trainer.sft.data import TokenShards, check_eos +from prime_rl.trainer.utils import get_dcp_group from prime_rl.trainer.world import get_world +from prime_rl.utils.logger import get_logger POSITION_BINS = (0, 128, 512, 2048, 4096, 16384) """ppl loss is also reported per bin of the target's position in its window.""" @@ -378,3 +381,76 @@ def summarize_pairs(num_good: int, logprobs) -> tuple[dict, dict]: correct = good_logprob > bad_logprob summary = {"acc": correct.mean(), "items": num_good} return summary, {"good_logprob": good_logprob, "bad_logprob": bad_logprob, "acc": correct} + + +def bytes_to_unicode() -> dict[int, str]: + """The byte -> character map of byte-level BPE tokenizers (GPT-2, Llama 3, Qwen, DeepSeek).""" + printable = ( + list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1)) + ) + chars = printable[:] + n = 0 + for byte in range(256): + if byte not in printable: + printable.append(byte) + chars.append(256 + n) + n += 1 + return dict(zip(printable, map(chr, chars))) + + +def token_byte_lengths(tokenizer: PreTrainedTokenizer) -> torch.Tensor: + """UTF-8 bytes each token id stands for. Special tokens count as 0 bytes. Exact for + byte-level BPE tokenizers, approximate otherwise.""" + byte_chars = set(bytes_to_unicode().values()) + special_ids = set(tokenizer.all_special_ids) | set(tokenizer.added_tokens_decoder) + lengths = torch.zeros(len(tokenizer), dtype=torch.long) + for token_id, token in enumerate(tokenizer.convert_ids_to_tokens(list(range(len(tokenizer))))): + if token is None or token_id in special_ids: + continue + if all(char in byte_chars for char in token): + lengths[token_id] = len(token) + else: + lengths[token_id] = len(tokenizer.convert_tokens_to_string([token]).encode()) + return lengths + + +class LogLikelihoodEval: + """The ``[ll_eval]`` sets, tokenized once at startup and scored on the training model. Requests start with BOS if + ``bos`` (the training documents do).""" + + def __init__(self, config: LogLikelihoodEvalConfig, tokenizer: PreTrainedTokenizer, bos: bool, run_dir: Path): + start_time = time.perf_counter() + self.config = config + self.run_dir = run_dir + has_token_sets = any(not isinstance(eval_set, str) for eval_set in config.sets) + byte_lengths = token_byte_lengths(tokenizer).numpy() if has_token_sets else None + self.sets = [ + prepare_multi_length_set(eval_set, config.data_dir, tokenizer, config.seq_lens, bos, byte_lengths) + for eval_set in config.sets + ] + get_logger().info( + f"Prepared {len(self.sets)} log-likelihood eval sets (BOS: {bos}) in {time.perf_counter() - start_time:.0f}s" + ) + + @torch.no_grad() + def run(self, model: nn.Module, step: int) -> dict[str, float]: + """Score every set on the weights after ``step`` optimizer updates. The metrics are on rank 0 only.""" + start_time = time.perf_counter() + metrics = {} + # Eager, also for compiled layers: the packed rows vary in length and would recompile at every forward. + with scoring_head_active(model), torch.compiler.set_stance("force_eager"): + scored = score_sets(model, self.sets, self.config.micro_batch_tokens, get_dcp_group()) + for eval_set, (logprobs, hits) in zip(self.sets, scored or []): + for seq_len, version in eval_set.versions.items(): + index = eval_set.index[seq_len] + summary, arrays = summarize(version, [logprobs[i] for i in index], [hits[i] for i in index]) + metrics |= set_metrics(eval_set.name, seq_len, summary) + if self.config.save_arrays: + out = self.run_dir / "ll_eval" / f"step_{step}" + out.mkdir(parents=True, exist_ok=True) + label = f"{eval_set.name}@{length_label(seq_len)}" + np.savez(out / f"{label}.npz", **arrays) + (out / f"{label}.json").write_text(json.dumps(summary, default=float)) + metrics["time/ll_eval"] = time.perf_counter() - start_time + get_logger().success(f"Log-likelihood evals | Step {step} | {metrics['time/ll_eval']:.0f}s") + return metrics diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 4c1b4abddf4..857d5810e2b 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -3,7 +3,6 @@ import math import time import asyncio -from contextlib import nullcontext from datetime import timedelta # Import environment before any other imports @@ -11,9 +10,9 @@ from prime_rl.utils.act_offloading import maybe_activation_offloading import torch -from torch.profiler import profile, ProfilerActivity, record_function +from torch.profiler import profile, ProfilerActivity from prime_rl.trainer.ckpt import Progress, setup_ckpt_manager -from prime_rl.utils.pathing import resolve_resume_step +from prime_rl.utils.pathing import resolve_latest_ckpt_step from prime_rl.configs.sft import SFTConfig from prime_rl.configs.trainer import CheckpointConfig from prime_rl.transports.weights import prune_broadcasts_beyond, setup_weight_sender @@ -26,10 +25,9 @@ from prime_rl.trainer.model import ( forward, get_full_offload_dtype_policy, - get_expert_load_stats, - get_global_moe_stats, get_load_balance_stats, is_tt_moe_model, + iter_moe_blocks, reshard_module, setup_processor, setup_tokenizer, @@ -37,24 +35,25 @@ ) from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.perf import get_perf_counter +from prime_rl.trainer.metrics import DeadExperts, SpikeDetector, grad_norms_by_group, take_routing_counts +from prime_rl.trainer.micro_batches import run_micro_batches from prime_rl.trainer.sft.data import ( + TokenDataset, get_dataset_progress, get_dataset_state, load_sft_dataset, setup_dataloader, setup_dataset, + starts_with_bos, ) +from prime_rl.trainer.sft.ll_eval import LogLikelihoodEval from prime_rl.trainer.utils import ( - get_grad_accum_steps, GarbageCollection, MemoryProfiler, - begin_backward, clip_grad_norm_, - finish_backward, get_ckpt_disk_metrics, prepare_gradient_offload, print_sample, - scale_gradients_, setup_full_cpu_optimizer_offload, setup_torch_distributed, ) @@ -117,9 +116,13 @@ def train(config: SFTConfig): # Initialize parallel dimensions parallel_dims = get_parallel_dims(config.model) - grad_accum_steps = get_grad_accum_steps( - config.data.batch_size, config.data.micro_batch_size, config.model.cp, world.world_size + total_micro_batches = config.data.batch_size * config.model.cp + micro_batches_per_step = world.world_size * config.data.micro_batch_size + assert total_micro_batches % micro_batches_per_step == 0, ( + f"batch_size * cp ({total_micro_batches}) must be divisible by " + f"world_size * micro_batch_size ({micro_batches_per_step})" ) + grad_accum_steps = total_micro_batches // micro_batches_per_step if parallel_dims.cp_enabled: assert config.data.seq_len % parallel_dims.cp == 0, "Sequence length must be divisible by CP degree" @@ -128,12 +131,37 @@ def train(config: SFTConfig): logger.info(f"Initializing checkpoint manager ({config.ckpt})") ckpt_manager = setup_ckpt_manager(config.run_dir, config.ckpt, resume=config.resume) - checkpoint_step = resolve_resume_step(config.resume, ckpt_manager.ckpt_dir) + checkpoint_step = None + if config.resume is not None: + if config.resume.dir is not None: + checkpoint_step = config.resume.dir_step + else: + checkpoint_step = config.resume.step + if checkpoint_step is None: + checkpoint_step = resolve_latest_ckpt_step(ckpt_manager.ckpt_dir) # Initialize the model and tokenizer logger.info(f"Initializing model ({config.model})") loading_from_ckpt_later = checkpoint_step is not None - model = setup_model(config.model, parallel_dims, loading_from_ckpt_later) + model = setup_model( + config.model, + parallel_dims, + loading_from_ckpt_later, + init=config.model.init if not loading_from_ckpt_later else None, + ) + logger.info( + f"Parameters: {model.num_active_non_embedding_params() / 1e6:,.1f}M active non-embedding, " + f"{model.num_non_embedding_params() / 1e6:,.1f}M non-embedding, " + f"{model.num_embedding_params() / 1e6:,.1f}M embedding, {model.num_params() / 1e6:,.1f}M total" + ) + if not loading_from_ckpt_later: + param_metrics = { + "params/total": model.num_params(), + "params/embedding": model.num_embedding_params(), + "params/non_embedding": model.num_non_embedding_params(), + "params/active_non_embedding": model.num_active_non_embedding_params(), + } + asyncio.run(monitors.log({**param_metrics, "step": 0}, step=0)) if parallel_dims.cp_enabled: setup_context_parallel(model, config.model, parallel_dims) @@ -223,8 +251,8 @@ def train(config: SFTConfig): cp_enabled = parallel_dims.cp_enabled cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() if cp_enabled else 0 dp_cp_group = parallel_dims.get_mesh("dp_cp").get_group() - ep_group = parallel_dims.get_mesh("ep").get_group() if parallel_dims.ep_enabled else None cp_size = parallel_dims.cp + moes = list(iter_moe_blocks(model)) def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: """Forward pass returning (loss_sum, token_count) over unmasked tokens.""" @@ -274,8 +302,6 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: del out return loss_sum, token_count - maybe_record_function = nullcontext - def run_eval_loop(data_iter): """Validation forward loop. Returns token-weighted global mean loss.""" total_loss_sum = torch.tensor(0.0, device="cuda") @@ -339,7 +365,33 @@ def run_validation(step: int) -> None: ) ) + ll_eval = None + if config.ll_eval is not None: + bos = isinstance(dataset, TokenDataset) and starts_with_bos(dataset.shards, tokenizer.bos_token_id) + ll_eval = LogLikelihoodEval(config.ll_eval, tokenizer, bos, config.run_dir) + + def run_ll_eval(step: int) -> None: + metrics = ll_eval.run(model, step) + # The no-grad forwards leave the [lm_head, norm] FSDP group unsharded, and their routing must not count + # toward the next step's statistics. + reshard_module(model) + if moes: + take_routing_counts(moes) + asyncio.run(monitors.log({**metrics, "step": step}, step=step)) + + def logged_micro_batches(micro_batches): + for micro_batch in micro_batches: + print_sample( + micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer + ) + yield micro_batch + + if ll_eval is not None and config.ll_eval.eval_on_start: + run_ll_eval(checkpoint_step or 0) + gc_handler = GarbageCollection(config.gc.interval) if config.gc else None + dead_experts = DeadExperts() + loss_spikes, grad_norm_spikes = SpikeDetector(), SpikeDetector() # A broadcast must land at every step an online eval env is due. The schedule is # deterministic, so all ranks agree when to enter the transport collective. @@ -373,7 +425,6 @@ def is_online_eval_step(step: int) -> bool: if config.trace_path: logger.info(f"Tracing to {config.trace_path}") prof = profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True).__enter__() - maybe_record_function = record_function # noqa: F841 – captured by run_forward_loop closure max_peak_memory = 0.0 controller = Controller(requires_step=False) control_server = ControlServer(controller, config.control if world.is_master else None) @@ -396,25 +447,12 @@ def is_online_eval_step(step: int) -> bool: step_start_time = time.perf_counter() forward_backward_start_time = time.perf_counter() - step_loss_sum = torch.tensor(0.0, device="cuda") - nan_loss_count = torch.tensor(0, device="cuda") - is_moe_model = is_tt_moe_model(model) - moe_stats = ( - { - "max_vio/mean": torch.tensor(0.0), - "max_vio/max": torch.tensor(0.0), - "routing_confidence/mean": torch.tensor(0.0), - } - if is_moe_model - else {} - ) run_validation_this_step = config.val is not None and ( (is_first_step and config.val.eval_on_start) or (not is_first_step and progress.step % config.val.interval == 0) ) if gradient_manager is None: micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) - step_local_token_count = torch.tensor(0, dtype=torch.int64, device="cuda") else: micro_batches = [next(dataiter) for _ in range(grad_accum_steps)] local_token_count = sum(int(micro_batch["loss_mask"].sum()) for micro_batch in micro_batches) @@ -431,58 +469,23 @@ def is_online_eval_step(step: int) -> bool: grad_scale, overlap_optimizer=not run_validation_this_step, ) + if config.log.log_data: + micro_batches = logged_micro_batches(micro_batches) - step_tokens_per_expert = 0 perf_counter = get_perf_counter(model) - step_local_num_tokens = 0 - step_local_flops = 0 - for micro_step, micro_batch in enumerate(micro_batches): - step_local_num_tokens += micro_batch["input_ids"].shape[1] - step_local_flops += perf_counter.flop_counter(micro_batch["seq_lens"].tolist()) - if config.log.log_data: - print_sample( - micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer - ) - - with maybe_record_function("forward"): - local_loss_sum, batch_token_count = compute_loss(micro_batch) - - if gradient_manager is None: - step_local_token_count += batch_token_count - - if torch.isnan(local_loss_sum.detach()): - nan_loss_count += 1 - logger.warning("Local loss is nan, excluding this micro step from backward") - scaled_loss = torch.nan_to_num(local_loss_sum, nan=0.0) / grad_accum_steps - else: - step_loss_sum += local_loss_sum.detach() - scaled_loss = local_loss_sum / grad_accum_steps - - with maybe_record_function("backward"): - begin_backward(gradient_manager, final_backward=micro_step == grad_accum_steps - 1) - scaled_loss.backward() - finish_backward(gradient_manager) - - if is_moe_model: - micro_moe_stats, tokens_per_expert = get_global_moe_stats(model, ep_group, dp_cp_group) - for name, value in micro_moe_stats.items(): - moe_stats[f"{name}/mean"] += value / grad_accum_steps - if name == "max_vio": - moe_stats["max_vio/max"] = torch.maximum(moe_stats["max_vio/max"], value) - step_tokens_per_expert += tokens_per_expert - - # Wait for the queued backward kernels so this times GPU work, not kernel launches. - torch.cuda.synchronize() + result = run_micro_batches( + model, + iter(micro_batches), + compute_loss, + grad_accum_steps, + parallel_dims, + perf_counter, + gradient_manager=gradient_manager, + defer_grad_sync=config.model.defer_grad_sync, + ) forward_backward_time = time.perf_counter() - forward_backward_start_time - expert_load_stats = get_expert_load_stats(step_tokens_per_expert, dp_cp_group) if is_moe_model else {} - - if gradient_manager is None: - global_step_token_count = step_local_token_count.clone() - dist.all_reduce(global_step_token_count, op=dist.ReduceOp.SUM, group=dp_cp_group) - global_token_count_val = global_step_token_count.item() - if global_token_count_val > 0: - grad_scale = parallel_dims.fsdp_gradient_divide_factor * grad_accum_steps / global_token_count_val - scale_gradients_(None, model, grad_scale) + if result.nan_count > 0: + logger.warning(f"{result.nan_count} micro-batches had a NaN loss and were excluded from the backward pass") # Run validation after forward-backward (so torch.compile sees training graph first) but before # optimizer step (so eval_on_start evaluates untrained weights) @@ -493,15 +496,7 @@ def is_online_eval_step(step: int) -> bool: # unsharded parameters. Reshard so clipping sees every gradient. reshard_module(model) - # Compute the global mean loss for logging. - dist.all_reduce(step_loss_sum, op=dist.ReduceOp.SUM, group=dp_cp_group) - dist.all_reduce(nan_loss_count, op=dist.ReduceOp.SUM) - if global_token_count_val > 0: - batch_loss = (step_loss_sum / global_token_count_val).item() - else: - batch_loss = 0.0 - nan_loss_count = nan_loss_count.item() - + group_grad_norms = grad_norms_by_group(model) if gradient_manager is None else {} grad_norm: torch.Tensor | None = None if config.optim.max_norm is not None: logger.debug(f"Clipping gradients with max norm {config.optim.max_norm}") @@ -534,20 +529,23 @@ def is_online_eval_step(step: int) -> bool: weight_sender.broadcast(model, step=progress.step) broadcast_weights_time = time.perf_counter() - broadcast_start_time + ll_eval_time = 0.0 + if ll_eval is not None and (is_last_step or progress.step % config.ll_eval.interval == 0): + ll_eval_start_time = time.perf_counter() + run_ll_eval(progress.step) + ll_eval_time = time.perf_counter() - ll_eval_start_time + # Optionally, dump memory snapshot if memory_profiler is not None: memory_profiler.step() - # The dp mesh excludes cp, whose ranks hold the same rows. - global_counts = torch.tensor([step_local_num_tokens, step_local_flops], dtype=torch.float64, device="cuda") - dist.all_reduce(global_counts, op=dist.ReduceOp.SUM, group=parallel_dims.get_mesh("dp").get_group()) - num_tokens, num_flops = int(global_counts[0].item()), global_counts[1].item() + num_tokens, num_flops = result.num_tokens, result.num_flops progress.total_tokens += num_tokens dataset_progress = get_dataset_progress(dataloader) progress.total_samples = dataset_progress["step"] # Throughput / MFU per step over the full step wall time, as torchtitan reports them with # log_freq=1 (tokens since last log / elapsed time), instead of a smoothed sliding window. - step_time = time.perf_counter() - step_start_time + step_time = time.perf_counter() - step_start_time - ll_eval_time perf_counter.count_tokens(num_tokens) throughput = perf_counter.get_step_tokens_per_second(num_tokens, step_time) mfu = perf_counter.get_step_mfu(num_flops, step_time) @@ -556,13 +554,14 @@ def is_online_eval_step(step: int) -> bool: max_peak_memory = max(max_peak_memory, peak_memory) # Log step metrics + batch_loss = result.loss step_message = f"Step {progress.step} | {format_time(step_time):>7} | Loss {batch_loss:.4f}" if grad_norm is not None: step_message += f" | Grad. Norm {grad_norm:.4f}" step_message += f" | LR {current_lr:.2e} | Throughput {throughput:.0f} tokens/s | MFU {mfu:.1f}% | MFU 6DN {mfu_6dn:.1f}% | Peak Mem. {peak_memory:.1f}/{max_memory:.1f} GiB ({peak_memory / max_memory * 100:.1f}%)" - if is_moe_model: + if result.tokens_per_expert is not None: for name, label in (("max_vio/mean", "Max Vio"), ("routing_confidence/mean", "Routing Conf.")): - value = moe_stats[name].item() + value = result.metrics[name] if value > 0: step_message += f" | {label} {value:.4f}" logger.success(step_message) @@ -608,14 +607,17 @@ def is_online_eval_step(step: int) -> bool: "optim/lr": current_lr, "step": progress.step, } + optim_metrics |= {f"optim/grad_norm/{group}": norm.item() for group, norm in group_grad_norms.items()} if grad_norm is not None: optim_metrics["optim/grad_norm"] = grad_norm.item() + optim_metrics["optim/grad_norm_spike"] = grad_norm_spikes.update(grad_norm.item()) asyncio.run(monitors.log(optim_metrics, step=progress.step)) loss_log_metrics = { "loss/mean": batch_loss, "loss/perplexity": math.exp(min(batch_loss, 20)), - "loss/nan_count": nan_loss_count, + "loss/nan_count": result.nan_count, + "loss/spike": loss_spikes.update(batch_loss), "step": progress.step, } # Log tensor stats @@ -636,9 +638,10 @@ def is_online_eval_step(step: int) -> bool: disk_metrics["step"] = progress.step asyncio.run(monitors.log(disk_metrics, step=progress.step)) - moe_log_metrics = {name: value.item() for name, value in moe_stats.items()} | expert_load_stats - if moe_log_metrics: - asyncio.run(monitors.log({**moe_log_metrics, "step": progress.step}, step=progress.step)) + step_metrics = {**result.metrics, "step": progress.step} + if result.tokens_per_expert is not None: + step_metrics["expert_load/dead"] = dead_experts.update(result.tokens_per_expert) + asyncio.run(monitors.log(step_metrics, step=progress.step)) is_first_step = False diff --git a/src/prime_rl/trainer/utils.py b/src/prime_rl/trainer/utils.py index 20caa0295a0..3ae52354f7d 100644 --- a/src/prime_rl/trainer/utils.py +++ b/src/prime_rl/trainer/utils.py @@ -17,7 +17,6 @@ from torch.distributed.tensor import DTensor from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.trainer.models.base import count_active_params from prime_rl.trainer.world import get_world from prime_rl.utils.logger import format_time, get_logger from prime_rl.utils.pathing import get_ckpt_dir @@ -364,30 +363,3 @@ def step(self): f"Finished dumping memory snapshot in {time.monotonic() - begin:.2f} seconds, load {file_path} at https://docs.pytorch.org/memory_viz to visualize the memory usage" ) self.step_num += 1 - - -def get_grad_accum_steps(batch_size: int, micro_batch_size: int, cp: int, world_size: int) -> int: - """Micro-batches each rank runs per optimizer step: ``batch_size`` samples, each split over ``cp`` ranks, in - micro-batches of ``micro_batch_size`` on ``world_size`` ranks.""" - total_micro_batches = batch_size * cp - micro_batches_per_step = world_size * micro_batch_size - assert total_micro_batches % micro_batches_per_step == 0, ( - f"batch_size * cp ({total_micro_batches}) must be divisible by " - f"world_size * micro_batch_size ({micro_batches_per_step})" - ) - return total_micro_batches // micro_batches_per_step - - -def count_parameters(model: nn.Module) -> dict[str, int]: - """Total, embedding (input embedding and LM head), non-embedding and active non-embedding parameter counts.""" - embedding_ids = {id(p) for p in (*model.get_input_embeddings().parameters(), *model.lm_head.parameters())} - params = list(model.parameters()) - non_embedding = [p for p in params if id(p) not in embedding_ids] - total = sum(p.numel() for p in params) - num_non_embedding = sum(p.numel() for p in non_embedding) - return { - "total": total, - "embedding": total - num_non_embedding, - "non_embedding": num_non_embedding, - "active_non_embedding": count_active_params(model, non_embedding), - } diff --git a/src/prime_rl/utils/pathing.py b/src/prime_rl/utils/pathing.py index 4fb063740f2..2052530cffb 100644 --- a/src/prime_rl/utils/pathing.py +++ b/src/prime_rl/utils/pathing.py @@ -12,7 +12,7 @@ from pydantic import BaseModel from prime_rl.configs.orchestrator import EnvConfig -from prime_rl.configs.shared import LogConfig, ResumeConfig +from prime_rl.configs.shared import LogConfig from prime_rl.utils.config import BaseConfig, dump_resolved_config from prime_rl.utils.logger import get_logger @@ -336,18 +336,6 @@ def resolve_latest_ckpt_step(ckpt_dir: Path) -> int | None: return latest_step -def resolve_resume_step(resume: ResumeConfig | None, ckpt_dir: Path) -> int | None: - """The checkpoint step a run resumes from: ``resume.dir``'s step, ``resume.step``, or the latest checkpoint in - ``ckpt_dir``. None without ``resume``.""" - if resume is None: - return None - if resume.dir is not None: - return resume.dir_step - if resume.step is not None: - return resume.step - return resolve_latest_ckpt_step(ckpt_dir) - - def has_checkpoints(output_dir: Path) -> bool: """Check if the output directory contains any checkpoints.""" ckpt_dir = get_ckpt_dir(output_dir) diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 4fa03841a5b..adf2230e768 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -24,7 +24,6 @@ from prime_rl.configs.orchestrator import ( ModelConfig as OrchestratorModelConfig, ) -from prime_rl.configs.pretrain import PretrainConfig from prime_rl.configs.rl import RLConfig from prime_rl.configs.sft import SFTConfig from prime_rl.configs.trainer import ModelConfig as TrainerModelConfig @@ -36,7 +35,6 @@ RLConfig, TrainerConfig, SFTConfig, - PretrainConfig, OrchestratorConfig, InferenceConfig, EnvServerConfig, diff --git a/tests/unit/train/pretrain/test_bytes.py b/tests/unit/train/pretrain/test_bytes.py deleted file mode 100644 index d2403e4d2e7..00000000000 --- a/tests/unit/train/pretrain/test_bytes.py +++ /dev/null @@ -1,12 +0,0 @@ -from transformers import AutoTokenizer - -from prime_rl.trainer.pretrain.bytes import token_byte_lengths - - -def test_token_byte_lengths_count_utf8_bytes(): - tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") - lengths = token_byte_lengths(tokenizer) - text = "Grüße aus Zürich — 東京 🚀 naïve café" - token_ids = tokenizer(text, add_special_tokens=False)["input_ids"] - assert int(lengths[token_ids].sum()) == len(text.encode()) - assert lengths[tokenizer.eos_token_id] == 0 diff --git a/tests/unit/train/pretrain/test_data.py b/tests/unit/train/pretrain/test_data.py deleted file mode 100644 index 20aeb21586c..00000000000 --- a/tests/unit/train/pretrain/test_data.py +++ /dev/null @@ -1,79 +0,0 @@ -from itertools import islice -from pathlib import Path - -import numpy as np - -from prime_rl.trainer.pretrain.data import TokenDataset, TokenShards, collate_rows - - -def write_shard(path: Path, documents: list[list[int]]) -> None: - """A datatrove DocumentTokenizer shard: uint32 tokens, uint64 cumulative document ends, metadata.""" - path.parent.mkdir(parents=True, exist_ok=True) - np.array([token for document in documents for token in document], dtype=" 1 for seq_lens in masked) - unmasked = islice(TokenDataset(shards, seq_len=4, document_masking=False), 8) - assert [window["seq_lens"].tolist() for window in unmasked] == [[4]] * 8 - - -def test_collate_rows_restarts_positions_at_documents(): - rows = [ - {"tokens": np.arange(5), "seq_lens": np.array([3, 1])}, - {"tokens": np.arange(10, 15), "seq_lens": np.array([4])}, - ] - batch = collate_rows(rows) - assert batch["input_ids"].tolist() == [[0, 1, 2, 3, 10, 11, 12, 13]] - assert batch["target_ids"].tolist() == [[1, 2, 3, 4, 11, 12, 13, 14]] - assert batch["position_ids"].tolist() == [[0, 1, 2, 0, 0, 1, 2, 3]] - assert batch["seq_lens"].tolist() == [3, 1, 4] diff --git a/tests/unit/train/pretrain/test_launcher.py b/tests/unit/train/pretrain/test_launcher.py deleted file mode 100644 index eee8ad647a2..00000000000 --- a/tests/unit/train/pretrain/test_launcher.py +++ /dev/null @@ -1,33 +0,0 @@ -import re - -from prime_rl.configs.pretrain import PretrainConfig -from prime_rl.entrypoints.pretrain import pretrain -from prime_rl.utils.config import cli - - -def test_multi_node_dry_run_writes_trainer_sbatch(tmp_path, monkeypatch): - monkeypatch.setenv("PRL_RUN_ID", "test") - monkeypatch.setenv("PRL_RUN_NAME", "test") - config = PretrainConfig.model_validate( - { - "output_dir": tmp_path, - "model": {"name": "configs/debug/fake/pretrain-model"}, - "tokenizer": {"name": "PrimeIntellect/Qwen3-0.6B"}, - "data": {"type": "fake", "seq_len": 8, "batch_size": 2, "micro_batch_size": 1}, - "optim": {"lr": 1e-3}, - "max_steps": 1, - "deployment": {"type": "multi_node", "num_train_nodes": 2}, - "slurm": {}, - "dry_run": True, - } - ) - pretrain(config) - - script = (config.run_dir / "launcher" / "pretrain.sbatch").read_text() - assert "#SBATCH --nodes=2" in script - assert "-m prime_rl.trainer.pretrain.train" in script - # The trainers on every node parse the written config, which has no [deployment] or [slurm]. - config_path = re.search(r"^export TRAIN_CONFIG_PATH=(.+)$", script, re.M).group(1) - trainer_config = cli(PretrainConfig, args=["@", config_path]) - assert trainer_config.slurm is None - assert trainer_config.model == config.model diff --git a/tests/unit/train/pretrain/test_pretrain_metrics.py b/tests/unit/train/pretrain/test_pretrain_metrics.py deleted file mode 100644 index eaf12314414..00000000000 --- a/tests/unit/train/pretrain/test_pretrain_metrics.py +++ /dev/null @@ -1,34 +0,0 @@ -import torch -from torch import nn - -from prime_rl.trainer.pretrain.metrics import expert_load_metrics, grad_norms_by_group, param_group - - -def test_expert_load_metrics_per_layer_summaries(): - counts = torch.tensor([[1.0, 1.0, 1.0, 1.0], [4.0, 0.0, 0.0, 0.0]], dtype=torch.float64) - metrics = expert_load_metrics(counts, layer_indices=[1, 2]) - assert metrics["expert_load/min_mean/min"] == 0.0 - assert metrics["expert_load/entropy/mean"] == 0.5 # uniform layer 1, one-hot layer 0 - assert metrics["expert_load/layer_2/max_mean"] == 4.0 - assert metrics["expert_load/layer_1/top_share"] == 0.25 - assert metrics["expert_load/layer_2/top_share"] == 1.0 - - -def test_grad_norms_by_group_add_up_to_the_total_norm(): - model = nn.ModuleDict( - { - "embed_tokens": nn.Embedding(8, 4), - "self_attn": nn.Linear(4, 4), - "input_layernorm": nn.LayerNorm(4), - "lm_head": nn.Linear(4, 8), - } - ) - for param in model.parameters(): - param.grad = torch.randn_like(param) - norms = grad_norms_by_group(model) - assert set(norms) == {"embed", "attn", "norm", "lm_head"} - total = torch.nn.utils.get_total_norm([param.grad for param in model.parameters()]) - assert torch.allclose(torch.stack(list(norms.values())).square().sum().sqrt().float(), total) - assert param_group("model.layers.3.mlp.router.gate.weight") == "router" - assert param_group("model.layers.3.mlp.shared_expert.w1") == "shared_experts" - assert param_group("model.layers.3.mlp.experts.w1") == "experts" diff --git a/tests/unit/train/pretrain/test_suite.py b/tests/unit/train/pretrain/test_suite.py deleted file mode 100644 index b82ff94d5d4..00000000000 --- a/tests/unit/train/pretrain/test_suite.py +++ /dev/null @@ -1,161 +0,0 @@ -import json -import math - -import numpy as np -import pyarrow as pa -import pyarrow.parquet as pq -import pytest -import torch -import torch.distributed as dist -from torch import nn - -from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX -from prime_rl.trainer.pretrain.data import TokenShards -from prime_rl.trainer.pretrain.suite import ( - head_logprobs, - pack_requests, - prepare_multi_length_set, - prepare_set, - prepare_token_set, - score_requests, - score_sets, - set_metrics, - summarize, - summarize_ctxcont, -) - - -def test_pack_requests_covers_every_request_within_the_budget(): - requests = [([0] * n, 1) for n in (5, 9, 3, 9, 2)] - batches = pack_requests(requests, token_budget=10) - assert sorted(i for batch in batches for i in batch) == list(range(5)) - assert all(sum(len(requests[i][0]) - 1 for i in batch) <= 10 for batch in batches) - - -def test_summarize_ctxcont_scores_the_gold_choice(): - rows = [ - {"context": "Q", "continuation": "ab", "choices": ["ab", "cd"], "gold": 0}, - {"context": "Q", "continuation": "e", "choices": ["x", "e"], "gold": 1}, - ] - options = [[("Q", c) for c in row["choices"]] for row in rows] - logprobs = [np.array([-1.0, -1.0]), np.array([-3.0]), np.array([-0.5]), np.array([-2.0])] - hits = [np.array([True, False]), np.array([False]), np.array([True]), np.array([False])] - summary, _ = summarize_ctxcont(rows, options, logprobs, hits) - assert summary["acc"] == 0.5 # item 0: gold wins; item 1: the other choice wins - assert math.isclose(summary["bpb"], (2.0 + 2.0) / (3 * math.log(2))) # gold nll / gold continuation bytes - assert summary["greedy_em"] == 0.0 - metrics = set_metrics("mc", 2048, summary) - assert metrics["eval/mc/acc@2k"] == 0.5 and all(type(value) is float for value in metrics.values()) - - -class CharTokenizer: - bos_token_id, eos_token_id = 1, 2 - - def __call__(self, texts, add_special_tokens=False, return_offsets_mapping=False): - if isinstance(texts, str): - return { - "input_ids": [ord(char) for char in texts], - "offset_mapping": [(i, i + 1) for i in range(len(texts))], - } - return { - "input_ids": [[ord(char) for char in text] for text in texts], - "offset_mapping": [[(i, i + 1) for i in range(len(text))] for text in texts], - } - - -def test_multi_length_set_scores_shared_requests_once(tmp_path): - (tmp_path / "docs").mkdir() - pq.write_table(pa.table({"text": ["short", "a much longer document"]}), tmp_path / "docs" / "test.parquet") - (tmp_path / "docs" / "meta.json").write_text(json.dumps({"kind": "ppl", "parquet_sha256": "x"})) - multi = prepare_multi_length_set("docs", tmp_path, CharTokenizer(), [8, 32], bos=True) - # "short" (5 chars + BOS + EOS) is one window at both lengths; the long document is 3 windows at 8, 1 at 32. - assert len(multi.versions[8].requests) == 4 and len(multi.versions[32].requests) == 2 - assert len(multi.requests) == 5 - assert multi.index[8][0] == multi.index[32][0] - - -def test_token_set_scores_documents_like_a_ppl_set_without_bos(tmp_path): - texts = ["short", "a much longer document"] - (tmp_path / "docs").mkdir() - pq.write_table(pa.table({"text": texts}), tmp_path / "docs" / "test.parquet") - (tmp_path / "docs" / "meta.json").write_text(json.dumps({"kind": "ppl", "parquet_sha256": "x"})) - documents = [[ord(char) for char in text] + [CharTokenizer.eos_token_id] for text in texts] - (tmp_path / "shards").mkdir() - np.array([token for document in documents for token in document], dtype=" None: + """A datatrove DocumentTokenizer shard: uint32 tokens, uint64 cumulative document ends, metadata.""" + path.parent.mkdir(parents=True, exist_ok=True) + np.array([token for document in documents for token in document], dtype=" 2 + assert row["position_ids"] == [p for length in row["seq_lens"] for p in range(length)] + assert row["input_ids"] == windows[0]["input_ids"] + windows[1]["input_ids"] From c97975427e4776ce6eef21a4db018d83b7b6fad4 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 22 Oct 2026 23:11:56 +0200 Subject: [PATCH 126/132] test(pretrain): dense pretraining integration test (start and resume), in the SFT CI group --- .github/workflows/tests.yml | 7 +- .../ci/integration/pretrain/model/config.json | 28 +++++ configs/ci/integration/pretrain/resume.toml | 4 + configs/ci/integration/pretrain/start.toml | 23 ++++ tests/integration/test_pretrain.py | 105 ++++++++++++++++++ 5 files changed, 164 insertions(+), 3 deletions(-) create mode 100644 configs/ci/integration/pretrain/model/config.json create mode 100644 configs/ci/integration/pretrain/resume.toml create mode 100644 configs/ci/integration/pretrain/start.toml create mode 100644 tests/integration/test_pretrain.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 6e90c6aa2db..0d8fa61c583 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -155,8 +155,8 @@ jobs: - 'src/prime_rl/entrypoints/{sft,dashboard}.py' - 'src/prime_rl/monitors/**' - 'src/prime_rl/eval/{online,runner}.py' - - 'tests/integration/test_reverse_text_sft{,_lora}.py' - - 'configs/ci/integration/reverse-text-sft{,-lora}/**' + - 'tests/integration/test_{reverse_text_sft,reverse_text_sft_lora,pretrain}.py' + - 'configs/ci/integration/{reverse-text-sft,reverse-text-sft-lora,pretrain}/**' rl: - *runtime - 'tools/convert_{dcp_to_bf16,dcp_to_fp8,bf16_to_fp8,fp8_to_bf16}.py' @@ -188,7 +188,8 @@ jobs: tests=$(jq -cn --argjson all "$RUN_ALL" --argjson sft "$SFT" --argjson rl "$RL" '[ if $all or $sft then {test: "reverse_text_sft", name: "Reverse Text SFT"}, - {test: "reverse_text_sft_lora", name: "Reverse Text SFT LoRA"} + {test: "reverse_text_sft_lora", name: "Reverse Text SFT LoRA"}, + {test: "pretrain", name: "Pretrain"} else empty end, if $all or $rl then {test: "reverse_text", name: "Reverse Text"}, diff --git a/configs/ci/integration/pretrain/model/config.json b/configs/ci/integration/pretrain/model/config.json new file mode 100644 index 00000000000..e991ed03b69 --- /dev/null +++ b/configs/ci/integration/pretrain/model/config.json @@ -0,0 +1,28 @@ +{ + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": null, + "eos_token_id": 151645, + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 256, + "initializer_range": 0.02, + "intermediate_size": 1024, + "max_position_embeddings": 2048, + "mlp_bias": false, + "model_type": "llama", + "num_attention_heads": 4, + "num_hidden_layers": 2, + "num_key_value_heads": 2, + "pad_token_id": 151643, + "pretraining_tp": 1, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "rope_theta": 10000.0, + "rope_type": "default" + }, + "tie_word_embeddings": false, + "transformers_version": "5.6.2", + "use_cache": true, + "vocab_size": 151680 +} diff --git a/configs/ci/integration/pretrain/resume.toml b/configs/ci/integration/pretrain/resume.toml new file mode 100644 index 00000000000..160d07d310e --- /dev/null +++ b/configs/ci/integration/pretrain/resume.toml @@ -0,0 +1,4 @@ +max_steps = 30 + +[resume] +step = 20 diff --git a/configs/ci/integration/pretrain/start.toml b/configs/ci/integration/pretrain/start.toml new file mode 100644 index 00000000000..d9852df46ca --- /dev/null +++ b/configs/ci/integration/pretrain/start.toml @@ -0,0 +1,23 @@ +# A small dense model trained from scratch on token shards. The test writes the shards (Reverse-Text-SFT, Qwen3 +# tokenizer) and passes their folder as --data.paths. +max_steps = 20 + +[model] +name = "configs/ci/integration/pretrain/model" + +[model.init] + +[tokenizer] +name = "PrimeIntellect/Qwen3-0.6B" + +[data] +type = "tokens" +paths = ["set-by-the-test"] +seq_len = 1024 +batch_size = 8 +micro_batch_size = 2 + +[optim] +lr = 3e-3 + +[ckpt] diff --git a/tests/integration/test_pretrain.py b/tests/integration/test_pretrain.py new file mode 100644 index 00000000000..39cc71c76bc --- /dev/null +++ b/tests/integration/test_pretrain.py @@ -0,0 +1,105 @@ +from pathlib import Path +from typing import Callable + +import numpy as np +import pytest +from datasets import load_dataset +from transformers import AutoTokenizer + +from tests.conftest import ProcessResult +from tests.utils import check_loss_goes_down, strip_escape_codes + +pytestmark = [pytest.mark.slow, pytest.mark.gpu] + +RUN_NAME = "pretrain" +TIMEOUT = 300 # 5 minutes + + +@pytest.fixture(scope="module") +def run_dir(output_dir: Path) -> Path: + return output_dir / RUN_NAME + + +@pytest.fixture(scope="module") +def wandb_name(branch_name: str) -> str: + return f"test-pretrain:{branch_name}" + + +@pytest.fixture(scope="module") +def token_dir(output_dir: Path) -> Path: + """Reverse-Text-SFT as one datatrove token shard: every conversation one document, ending with EOS.""" + tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/Qwen3-0.6B") + documents = [ + tokenizer.encode("\n".join(message["content"] for message in row["prompt"] + row["completion"])) + + [tokenizer.eos_token_id] + for row in load_dataset("PrimeIntellect/Reverse-Text-SFT", split="train") + ] + path = output_dir / "pretrain-tokens" + path.mkdir(parents=True, exist_ok=True) + np.array([token for document in documents for token in document], dtype=" ProcessResult: + cmd = pretrain_command(["start.toml"], token_dir, wandb_project, wandb_name, output_dir) + return run_process([*cmd, "--clean"], timeout=TIMEOUT) + + +@pytest.fixture(scope="module") +def pretrain_resume_process( + pretrain_process: ProcessResult, + run_process: Callable[..., ProcessResult], + token_dir: Path, + wandb_project: str, + wandb_name: str, + output_dir: Path, +) -> ProcessResult: + cmd = pretrain_command(["start.toml", "resume.toml"], token_dir, wandb_project, f"{wandb_name}-resume", output_dir) + return run_process(cmd, timeout=TIMEOUT) + + +def trainer_log(run_dir: Path) -> list[str]: + with open(run_dir / "logs" / "latest" / "trainer.log") as f: + return strip_escape_codes(f.read()).splitlines() + + +def test_no_error(pretrain_process: ProcessResult): + assert pretrain_process.returncode == 0, f"Process has non-zero return code ({pretrain_process})" + + +def test_loss_goes_down(pretrain_process: ProcessResult, run_dir: Path): + check_loss_goes_down(trainer_log(run_dir)) + + +def test_no_error_resume(pretrain_resume_process: ProcessResult): + assert pretrain_resume_process.returncode == 0, f"Process has non-zero return code ({pretrain_resume_process})" + + +def test_loss_goes_down_resume(pretrain_resume_process: ProcessResult, run_dir: Path): + check_loss_goes_down(trainer_log(run_dir)) From a348ebe5747b1649eae1410e453880b762a5266e Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 22 Oct 2026 23:12:35 +0200 Subject: [PATCH 127/132] docs(pretrain): a short pretraining section in training.md instead of docs/pretraining.md; the new metrics --- docs/mint.json | 1 - docs/overview.md | 1 - docs/pretraining.md | 149 -------------------------------------------- docs/training.md | 69 ++++++++++++++++++-- 4 files changed, 65 insertions(+), 155 deletions(-) delete mode 100644 docs/pretraining.md diff --git a/docs/mint.json b/docs/mint.json index b75d50f8137..d0f201b4e42 100644 --- a/docs/mint.json +++ b/docs/mint.json @@ -7,7 +7,6 @@ "overview", "configuration", "training", - "pretraining", "eval", "inference", "scaling", diff --git a/docs/overview.md b/docs/overview.md index 78334758fd1..7df75dac171 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -38,7 +38,6 @@ The `rl` entrypoint reads `examples/basic/reverse-text/rl.toml`, splits it into - **[Configuration](configuration.md)** — TOML composition, CLI overrides, dry-run. - **[Training](training.md)** — Launch and observe RL and SFT runs. -- **[Pretraining](pretraining.md)** — Train a model from scratch on token shards. - **[Inference](inference.md)** — vLLM-backed server (or fleet) holding the current policy. - **[Scaling](scaling.md)** — Single-GPU through multi-node clusters via FSDP / EP / CP and SLURM. - **[Algorithms](algorithms.md)** — Async semantics, loss / advantage / filter plugins, trajectory merging. diff --git a/docs/pretraining.md b/docs/pretraining.md deleted file mode 100644 index 37aad5a4f25..00000000000 --- a/docs/pretraining.md +++ /dev/null @@ -1,149 +0,0 @@ -# Pretraining - -`uv run pretrain` trains a model from scratch on pre-tokenized token shards. It has its own config (`PretrainConfig`) and training loop, built from the same model setup, FSDP/EP/CP, optimizer, checkpointing and monitors as the SFT trainer. - -## Table of Contents - -- [Launch](#launch) -- [Model](#model) -- [Initialization](#initialization) -- [Optimizer Groups](#optimizer-groups) -- [Learning-Rate Schedule](#learning-rate-schedule) -- [MoE Load Balancing](#moe-load-balancing) -- [Data](#data) -- [Monitoring](#monitoring) -- [Evals](#evals) -- [Multi-Node](#multi-node) - -## Launch - -```bash -uv run pretrain @ examples/pretrain/prime_moe.toml # MoE example on token shards -uv run pretrain @ configs/debug/fake/pretrain.toml # 5-step smoke test on random tokens -``` - -It runs locally on one node or, with `[slurm]`, as a SLURM job on one or more nodes (see [Multi-Node](#multi-node)). - -## Model - -The model is set up as in SFT and RL: `model.name` is an HF model id or a local model directory with a `config.json`. Its weights are not loaded. The model starts from a random init (below) or from a checkpoint. A new architecture is a directory with only a `config.json` for a model type in `prime_rl.trainer.models`; [`examples/pretrain/moe-d512`](../examples/pretrain/moe-d512/config.json) is the one in the example. `vocab_size` and the BOS, EOS and pad ids in the config must fit the tokenizer. The checked-in configs round `len(tokenizer)` up to a multiple of 128 and use the tokenizer's special ids, with EOS as the pad id if the tokenizer has none. - -```toml -[model] -name = "examples/pretrain/moe-d512" - -[tokenizer] -name = "tokenizers/deepseek-v4-flash-base" # defaults to model.name -``` - -MoE models use `prime_moe`, pretraining's own model (`prime_rl.trainer.models.prime_moe`), so architecture work does not touch the shared MoE code that SFT and RL use. Its base is a decoder of grouped-query attention (optional QK norm, partial RoPE) and MoE layers with sigmoid-scored token-choice routed experts, a selection bias and shared experts, after leading dense layers. Its module names and state-dict layout match `glm4_moe`'s, so checkpoints convert the same way. On top of the base: -- Gated Residual connections (Qwen3.8-Next): `hc_count` residual streams (above 1) with read gates of rank `hc_lowrank` replace the pre-norms and the final norm. Their projections train with AdamW. -- Null experts (AdaMoE): `n_null_experts` zero-output experts the router can select next to the routed ones, so tokens use between 0 and `num_experts_per_tok` routed experts. Only the (token, routed expert) pairs are dispatched. The metrics are `router/k/*`. -- Exclusive self attention (XSA): `exclusive_attention` removes from each head's output its component along the token's own value vector. -- Gated attention (Qwen3-Next): `attn_output_gate` multiplies the attention output elementwise by `sigmoid(g_proj(x))`, computed from the block input, before `o_proj` (after XSA). -- LatentMoE: with `latent_dim`, the routed experts work in that width (projection down and RMS norm before the dispatch, projection up after the combine). The router and the shared experts stay full width. With `latent_dim = d/2`, doubling the routed experts, the null experts and `num_experts_per_tok` (and `routed_scaling_factor`, which keeps each expert's gate) keeps the expert parameters and routed FLOPs per token (Kimi K3, NVIDIA LatentMoE). `latent_norm` puts its RMS norm on the expert input (`"input"`, Marin) or on the combined expert output before the up-projection (`"output"`, Kimi K3's Stable LatentMoE). -- `expert_activation = "situ_glu"`: Kimi K3's SiTU-GLU in the routed and shared experts, SwiGLU with both factors soft-capped (`4 tanh(g/4) sigmoid(g) * 25 tanh(u/25)`). -- `router_bias_on_logits` (default true) adds the selection bias to the router logits instead of the sigmoid scores (for quantile balancing). - -The trainer and the HF export (`tools/convert_dcp_to_bf16.py`) read the same model directory and tokenizer; vLLM does not support `prime_moe` yet. The trainer logs the active non-embedding (routed experts counted at `top_k / num_experts`), non-embedding, embedding (input embedding and LM head) and total parameter counts at startup, and as `params/*` at step 0. - -## Initialization - -Weights start from a random init, unless the run resumes from a checkpoint. The init lives with the models (`prime_rl.trainer.models.init`) and runs in the shared `setup_model` when it gets an `init` config. SFT and RL do not pass one, so their behavior is unchanged. - -| Field | Meaning | -|---|---| -| `model.init.std` | Every weight matrix and embedding (including the LM head, the router, expert and Gated Residual projections) starts from N(0, std) (default 0.02), with no width or depth scaling. Norm weights start at 1, biases at 0; zero-centred norm gains and the MoE selection biases start at 0. | -| `model.init.seed` | Seed. Every rank initializes only its own shard, and the values depend only on the seed and the parameter name, not on the number of GPUs. | - -## Optimizer - -Pretraining uses the optimizer setup of SFT and RL. With `optim.type = "muon"` (the default), every weight matrix except the embedding and the LM head is trained by Muon (`rms_norm` learning-rate adjustment; expert matrices per expert), and the embedding, the LM head, norms, biases and the Gated Residual projections (`hyper_connection`) by AdamW with `optim.eps`. `optim.type = "adamw"` trains everything with AdamW. `model.offload = "optimizer"` keeps the optimizer states in CPU RAM between steps; `"full"` is not supported. - -## Learning-Rate Schedule - -Pretraining uses the shared `[scheduler]`. Its default here is the `linear` scheduler with a warmup over 10% of the steps and a linear decay to `min_lr = 0` over the remaining 90%. Warmup and decay take a number of steps or a fraction of `max_steps` (rounded down), and `decay_shape` is `linear`, `cosine` or `1-sqrt`: - -```toml -[scheduler] -type = "linear" -warmup_steps = 0.05 -decay_steps = 0.2 -decay_shape = "1-sqrt" -``` - -The schedule scales every param group's own learning rate, so ratios between groups hold. AdamW and Muon apply weight decay as `lr * weight_decay`, so the schedule also scales the decay. The scheduler state is checkpointed; resuming with `ckpt.skip_scheduler = true` rebuilds the schedule over the remaining steps. - -## MoE Load Balancing - -`[model.moe.balance]` balances the routed experts of `prime_moe` as in DeepSeek-V3: - -| Field | Meaning | -|---|---| -| `bias_update` | `quantile` (default): quantile balancing. Every step, each layer's bias is set so that every expert would have been selected by `num_experts_per_tok / num_experts` of the step's tokens (null experts included), estimated as `quantile_estimator` sets; no rate. Best with `router_bias_on_logits` (on by default). `sign`: the update below. The bias starts at 0. | -| `quantile_estimator` | How `quantile` estimates the thresholds. `histogram` (default; Kimi K3 report, §2.3.3): one histogram of the margins per expert over the whole step, summed over the ranks with one all-reduce, read with in-bin interpolation; the quantile of the pooled step, independent of how the tokens are split. `micro_batch_mean`: each micro-batch's exact quantile, averaged over the micro-batches and ranks. | -| `quantile_bins` | Bins of the `histogram` estimator (default 1024). Its grid is the previous step's global margin range, `[-1, 1]` on the first step after a start or resume. | -| `bias_update_rate` | With `sign`: every step, each layer's selection bias moves by this step size (gamma) towards underloaded experts, using the step's global expert counts. The update is mean-centred. Default 3e-3; 0 turns it off. | -| `bias_freeze` | Fraction of `max_steps` at the end during which the biases stop updating. Default 0.1. | -| `seq_aux_coeff` | Weight (alpha) of the sequence-wise balance loss, computed per packed row of `data.seq_len` tokens and added to the training loss. Default 1e-3; 0 turns it off. | - -The metrics are `expert_load/*` (global balance of the step), `loss/seq_balance` (1 is balanced) and `router/bias/*`. Every 32 steps the trainer also records each layer's per-expert selection bias, router-row norm and load in `moe_trace.npz` in the run directory. - -## Data - -`data.type = "tokens"` reads [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: folders of `*.ds` token files with their `.ds.index` (document ends) and `.ds.metadata`. The files of all folders in `data.paths` (in folder order, then file-name order) are read as one token pool. Each epoch reads every row of the pool once, in a random order over the whole pool and from a random start offset, both drawn from `data.seed` and the epoch (the last `seq_len` tokens or fewer of an epoch are skipped). So every epoch has a new order and new row boundaries, and part of an epoch is a sample of the whole pool, whatever the order of the files. Every row is a random read (about 33 ms each from a cold cache on CephFS); at the screening shapes, `data.num_workers = 2` with the default prefetch hides this (`time/data` stays at 0). Check `time/data` on other storage. Documents continue across rows, and attention and positions restart at every document and every row; with `data.document_masking = false`, a row is one sequence (attention spans the whole row and positions run over it, as in GPT-style training without document masking). Every document must end with the tokenizer's EOS; the trainer checks the first document of every folder, which catches shards made with another tokenizer. - -Every folder contributes rows in proportion to its tokens, so mixtures are made offline, as folders whose token counts follow the mixture weights. At startup the trainer computes how many passes over the pool the run makes and fails above `data.max_epochs`. - -```toml -[data] -type = "tokens" -seq_len = 8192 # tokens per row -batch_size = 32 # rows per optimizer step, over all GPUs -micro_batch_size = 8 # rows per forward pass on one GPU -paths = ["data/web", "data/code"] -``` - -Batch sizes are in rows, as in SFT: a step trains on `batch_size * seq_len` tokens, so changing `seq_len` changes the tokens per step. Resume restores the data position on the same GPU layout and `data.num_workers`; resuming on a different layout is not supported yet. `data.type = "fake"` trains on random tokens, for smoke and throughput tests. - -## Monitoring - -Every step logs the loss, learning rate, gradient norms (total and per parameter group, `optim/grad_norm/`), loss and grad-norm spike flags, the rows, tokens and epochs trained on (`progress/num_samples`, `progress/num_tokens`, `progress/epochs`), throughput, MFU, data wait and peak memory (max and mean over ranks). MoE models add the expert balance of the step's global counts (`expert_load/*`, also per layer), `router/*` and `loss/seq_balance`. The step statistics are reduced once per step. - -## Evals - -`[evals]`: every `interval` steps, at the last step and (with `eval_on_start`, default) on the weights the run starts from, the trainer scores its eval sets with forward passes without gradients on the training model itself (same sharding, dtype and fusions). Evals run after the optimizer update, so the metrics at step N describe the same weights as the step-N checkpoint. A set is either - -- the name of a log-likelihood set of the pretrain-evals offline suite (`ppl`, `ctxcont` or `pairs`, in `//{test.parquet,meta.json}`), or -- a held-out folder of datatrove token shards (`{ name = ..., path = ... }`), scored like a `ppl` set: every document of the folder as stored, in windows of each of `seq_lens`, every token after its first scored (loss in nats per token and bits per UTF-8 byte, comparable across tokenizers). Each eval scores the whole folder at every length, so keep held-out folders small: at d512 on one H200, a token set costs about 1 s per 200K tokens per length (5M tokens at `seq_lens = [2048, 8192]`: 40 s, against 280 s for the 58-set suite). Because every document is scored from its start, these losses are lower than those of the former `[val]` (by 0.014-0.027 nats at d512), which started rows mid-document and predicted each document's first token across the document boundary: don't compare old `val/*` numbers with `eval/*` token-set numbers. - -Requests are packed into document-masked rows of `evals.micro_batch_tokens` tokens (default: the training micro-batch, `data.micro_batch_size * data.seq_len`, or the longest of `seq_lens`, whichever is larger) that are dealt over the ranks; the LM head returns fp32 per-token log-probs for the duration of the evals. Every set is scored at each context length in `seq_lens` (default: `data.seq_len`) and logged as `eval/{set}/{metric}@{length}` (e.g. `eval/fwedu-holdout/bpb@4k`, plus accuracies and greedy exact match for multiple-choice and gold-answer sets); requests that are the same at several lengths (e.g. documents shorter than the smallest window) are scored once. Requests of the pretrain-evals sets start with BOS only if the training documents do (read off the first document of every training folder at startup), so the evals see the training format; without BOS, the first token of a `ppl` document or `pairs` sentence is only context (`ppl` bpb leaves out its bytes). Requires `model.cp = 1`. Evals run without gradients, so longer rows can fit: in our tests, eval memory stayed at most 0.6x the training peak, and at d512 (8k training micro-batches), rows of 32k instead of 16k tokens cut the eval time by 19%. - -```toml -[evals] -interval = 1000 -data_dir = "/shared/evals-data" -sets = ["fwedu-holdout", "lambada", "math", { name = "heldout", path = "data/dolma3-50b-dsv4/mix-val-5M" }] -seq_lens = [2048, 4096, 16384] -``` - - -## Multi-Node - -`uv run pretrain` takes the same `[deployment]` and `[slurm]` blocks as `uv run sft` (see [Scaling § SLURM](scaling.md#slurm)) and submits the same bundled templates. The launcher writes the model dir on the submitting node, then submits one job in which every node runs `torchrun` on the pretrain trainer: - -```toml -output_dir = "/shared/outputs/pretrain" - -[deployment] -type = "multi_node" -num_train_nodes = 4 -gpus_per_node = 8 -nodes_per_fsdp_group = 1 # optional: HSDP, shard within a node and replicate across nodes - -[slurm] -partition = "all" -time = "04:00:00" -``` - -Row `w` of the token stream goes to data stream `w % num_streams`, with `num_streams = data_parallel_ranks * data.num_workers`, and step `s` (counted from 1) trains on rows `(s - 1) * rows_per_step` to `s * rows_per_step - 1`. So runs with different numbers of GPUs train on the same rows at every step, only split differently into micro-batches. The loss and gradient are normalized by the token count of the whole global batch, as on one GPU with gradient accumulation. diff --git a/docs/training.md b/docs/training.md index 9649264c6f2..9f92acce368 100644 --- a/docs/training.md +++ b/docs/training.md @@ -17,6 +17,7 @@ This page covers everything you need to launch, observe, checkpoint, and recover - [Launch](#launch-1) - [SFT-Specific Knobs](#sft-specific-knobs) - [Important Metrics](#important-metrics-1) +- [Pretraining](#pretraining) - [Checkpointing](#checkpointing) - [Enabling Checkpoints](#enabling-checkpoints) - [Resuming a Run](#resuming-a-run) @@ -34,8 +35,7 @@ This page covers everything you need to launch, observe, checkpoint, and recover | Command | Purpose | Notes | |---|---|---| | `uv run rl` | Wraps the trainer, orchestrator, and inference server in one launch from a merged TOML. | The default for any RL run. Runs locally for single-node experiments; submits to SLURM for single- or multi-node when `[slurm]` is set (see [Scaling § SLURM](scaling.md#slurm)). | -| `uv run sft` | Supervised fine-tuning on a HF dataset. | Launches torchrun internally; never call torchrun directly. | -| `uv run pretrain` | Pretraining from scratch on token shards. | See [Pretraining](pretraining.md). | +| `uv run sft` | Supervised fine-tuning on a HF dataset, or [pretraining](#pretraining) on token shards. | Launches torchrun internally; never call torchrun directly. | | `uv run inference` | vLLM server. | Always use this entrypoint over `vllm serve` — it adds `/update_weights`, `/load_lora_adapter`, and `/init_broadcaster`. | | `uv run orchestrator` | Standalone orchestrator process. | Pair with a separately-launched inference server and one `env-server` per source. | | `uv run eval` | Multi-env evals against a live inference server. | One epoch per source, pinned (or adaptive) concurrency, cursor checkpoints + `--resume`, dashboard + optional platform upload; see [Eval](eval.md). | @@ -284,10 +284,11 @@ Pulled from the console log and mirrored to W&B. **Stability and optimization:** -- `optim/grad_norm` — spikes precede divergence. +- `optim/grad_norm` — spikes precede divergence. `optim/grad_norm/` splits it by parameter group (embed, attn, mlp, experts, router, norm, lm_head, ...). +- `loss/spike`, `optim/grad_norm_spike` — 1 when the value is more than 6 standard deviations above the mean of the previous 128 steps. - `optim/lr` — LR schedule. - For MoE: `max_vio/mean`, `max_vio/max` — mean and max over the step's microsteps of the largest load-balancing violation across layers and EP groups, computed from expert token counts summed across each EP group. `routing_confidence/mean` — mean routing confidence. -- For MoE: `expert_load/{cv,max_mean,cold_frac}/{mean,max}` — expert load per layer over the whole step, from token counts summed across all ranks, as the mean and max over MoE layers. `cv` is std/mean of the per-expert token counts, `max_mean` the busiest expert's count over the mean, and `cold_frac` the fraction of experts receiving under 0.1x the mean. A trainable router that drifts shows up as a rising `cv`/`cold_frac` (MiMo-V2.6 §5.4). +- For MoE: `expert_load/{cv,max_mean,cold_frac}/{mean,max}` — expert load per layer over the whole step, from token counts summed across all ranks, as the mean and max over MoE layers. `cv` is std/mean of the per-expert token counts, `max_mean` the busiest expert's count over the mean, and `cold_frac` the fraction of experts receiving under 0.1x the mean. A trainable router that drifts shows up as a rising `cv`/`cold_frac` (MiMo-V2.6 §5.4). `expert_load/{min_mean,entropy}/{mean,min}` add the least-loaded expert over the mean and the normalized entropy of the expert shares (1 is uniform), `expert_load/layer_/{max_mean,top_share}` the same per decoder layer, and `expert_load/dead` the experts that got no token in the last 100 steps. **Performance:** @@ -298,6 +299,66 @@ Pulled from the console log and mirrored to W&B. | `perf/mfu_6dn` | MFU from the plain 6 x active params x tokens estimate (no attention FLOPs) | | `perf/peak_memory` | peak GPU memory (GiB) | | `time/step`, `time/forward_backward`, `time/save_ckpt` | step breakdown | +| `time/data/{max,mean}` | time the step waited for data, max and mean over ranks | +| `params/*` | total, embedding, non-embedding and active non-embedding parameter counts, at step 0 | + +## Pretraining + +Pretraining is SFT on token shards: set `data.type = "tokens"` and `[model.init]`, and launch with `uv run sft` as above. `model.name` only supplies the model config (an HF model id or a directory with a `config.json`); the weights start from the init, or from the checkpoint when resuming. + +```toml +max_steps = 10000 + +[model] +name = "path/to/model-dir" # a config.json; vocab size and special ids must fit the tokenizer +defer_grad_sync = true # optional, see below + +[model.init] +std = 0.02 # every weight matrix and embedding from N(0, std); norms 1, biases 0 +seed = 0 # the values do not depend on the number of GPUs + +[tokenizer] +name = "path/to/tokenizer" + +[data] +type = "tokens" +paths = ["data/web", "data/code"] +seq_len = 8192 # tokens per window +batch_size = 256 # windows per optimizer step, over all GPUs +micro_batch_size = 4 # windows per forward pass on one GPU + +[optim] +type = "muon" +lr = 3e-3 + +[scheduler] +type = "linear" +warmup_steps = 1000 +decay_steps = 2000 + +[ckpt] +interval = 1000 +``` + +**Token shards.** `paths` are folders of [datatrove](https://github.com/huggingface/datatrove) `DocumentTokenizer` output: `*.ds` token files, each with its `.ds.index` (document ends) and `.ds.metadata`. The files of all folders (in folder order, then file name order) are read as one token pool, so mixtures are made offline, as folders sized to the mixture weights. Every document must end with the tokenizer's EOS; the trainer checks the first document of every folder. + +**Windows.** Each epoch reads every `seq_len` window of the pool once, in a random order and from a random offset drawn from `data.seed` and the epoch. Every token is trained on. Attention and positions restart at every document (`data.document_masking = false` makes a window one sequence) and at every window. Resume restores the data position on the same GPU layout and `data.num_workers`. `progress/epoch` is the fraction of the pool read. + +**Gradient sync.** `model.defer_grad_sync = true` reduce-scatters the gradients once per optimizer step instead of after every micro-batch. It saves communication with gradient accumulation (most with experts across nodes), but keeps the unsharded parameters and gradients in memory across the micro-batches, so a config that fits without it can run out of memory. + +**Log-likelihood evals.** `[ll_eval]` scores eval sets with forward passes without gradients on the training model, every `interval` steps, at the last step and (with `eval_on_start`, the default) on the weights the run starts from. Evals run after the optimizer update, so step N's metrics describe the step-N checkpoint. A set is the name of a pretrain-evals set (`ppl`, `ctxcont` or `pairs`, in `//{test.parquet,meta.json}`) or a held-out folder of token shards, scored like a `ppl` set: every document on its own, every token after its first. Every set is scored at each of `seq_lens` (default `data.seq_len`). Requests start with BOS only if the training documents do. Requires `model.cp = 1`. + +```toml +[ll_eval] +interval = 1000 +data_dir = "/shared/evals-data" +sets = ["lambada", { name = "heldout", path = "data/heldout" }] +seq_lens = [2048, 8192] +``` + +The metrics are `eval//@`, e.g. `eval/heldout/bpb@8k`: `loss` (nats per token), `bpb` (bits per UTF-8 byte, comparable across tokenizers), and `acc`, `acc_norm` and `greedy_em` for multiple-choice and gold-answer sets. `time/ll_eval` is the time per round. Training metrics are the SFT ones above. + +For multi-node runs, use `[deployment]` and `[slurm]` as for SFT (see [Scaling § SLURM](scaling.md#slurm)). ## Checkpointing From 2fd8b9b08ce239df095e0c53e661bff72842b331 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 22 Oct 2026 23:30:47 +0200 Subject: [PATCH 128/132] fix(pretrain): convert_dcp_to_bf16 reads SFT runs' model config as SFTModelConfig; the test passes data.paths as a list --- tests/integration/test_pretrain.py | 2 +- tools/convert_dcp_to_bf16.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/integration/test_pretrain.py b/tests/integration/test_pretrain.py index 39cc71c76bc..1bc3c24ab0c 100644 --- a/tests/integration/test_pretrain.py +++ b/tests/integration/test_pretrain.py @@ -49,7 +49,7 @@ def pretrain_command(configs: list[str], token_dir: Path, wandb_project: str, wa "sft", *[arg for config in configs for arg in ("@", f"configs/ci/integration/pretrain/{config}")], "--data.paths", - token_dir.as_posix(), + f'["{token_dir.as_posix()}"]', "--deployment.num-train-gpus", "2", "--monitors.wandb.project", diff --git a/tools/convert_dcp_to_bf16.py b/tools/convert_dcp_to_bf16.py index a42bbf77029..1e0f6609a85 100644 --- a/tools/convert_dcp_to_bf16.py +++ b/tools/convert_dcp_to_bf16.py @@ -31,6 +31,7 @@ from torch.distributed.checkpoint import FileSystemReader from torch.distributed.checkpoint.state_dict_loader import load as dcp_load +from prime_rl.configs.sft import SFTModelConfig from prime_rl.configs.trainer import ModelConfig, MoERuntimeConfig, TokenizerConfig from prime_rl.trainer.ckpt import AppState from prime_rl.trainer.model import setup_model, setup_processor, setup_tokenizer @@ -57,7 +58,7 @@ "attn": "flash_attention_2", } -RUN_CONFIG_NAMES = ("trainer.json", "sft.json") +RUN_MODEL_CONFIGS = {"trainer.json": ModelConfig, "sft.json": SFTModelConfig} def resolve_dcp_dir(ckpt_dir: Path) -> Path: @@ -73,7 +74,7 @@ def resolve_run_configs(step_dir: Path) -> tuple[ModelConfig, TokenizerConfig]: """Model/tokenizer configs from the run's resolved config.""" logger = get_logger() config_dir = get_config_dir(step_dir.parent.parent) - for name in RUN_CONFIG_NAMES: + for name in RUN_MODEL_CONFIGS: path = config_dir / name if path.exists(): logger.info(f"Reading model config from {path}") @@ -81,11 +82,11 @@ def resolve_run_configs(step_dir: Path) -> tuple[ModelConfig, TokenizerConfig]: break else: raise FileNotFoundError( - f"No resolved run config ({' or '.join(RUN_CONFIG_NAMES)}) found under {config_dir} " + f"No resolved run config ({' or '.join(RUN_MODEL_CONFIGS)}) found under {config_dir} " "- the checkpoint must live in its run directory" ) - model = ModelConfig(**run_config["model"]) + model = RUN_MODEL_CONFIGS[name](**run_config["model"]) tokenizer = TokenizerConfig(**run_config["tokenizer"]) return model.model_copy(update=CONVERSION_OVERRIDES), tokenizer From e2c3465543a7de2ce53967c88434980a81ed8c91 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 23 Oct 2026 00:08:56 +0200 Subject: [PATCH 129/132] refactor(pretrain): the log-likelihood sets are part of [val] instead of [ll_eval] val.sets (pretrain-evals sets or held-out token folders), val.data_dir, val.seq_lens, val.micro_batch_tokens and val.save_arrays; val.data becomes optional. They run where val.data does, logged as val//@, and time/val times the validation round. SFT's val.data path is unchanged. --- configs/ci/integration/pretrain/start.toml | 8 +- docs/training.md | 10 +- .../src/prime_rl/configs/sft.py | 88 ++++++++--------- src/prime_rl/trainer/sft/train.py | 96 +++++++++---------- .../trainer/sft/{ll_eval.py => val_sets.py} | 25 +++-- tests/integration/test_pretrain.py | 2 + .../sft/{test_ll_eval.py => test_val_sets.py} | 4 +- 7 files changed, 109 insertions(+), 124 deletions(-) rename src/prime_rl/trainer/sft/{ll_eval.py => val_sets.py} (94%) rename tests/unit/train/sft/{test_ll_eval.py => test_val_sets.py} (93%) diff --git a/configs/ci/integration/pretrain/start.toml b/configs/ci/integration/pretrain/start.toml index d9852df46ca..fa8f62edce9 100644 --- a/configs/ci/integration/pretrain/start.toml +++ b/configs/ci/integration/pretrain/start.toml @@ -1,5 +1,5 @@ -# A small dense model trained from scratch on token shards. The test writes the shards (Reverse-Text-SFT, Qwen3 -# tokenizer) and passes their folder as --data.paths. +# A small dense model trained from scratch on token shards, validated on them as a held-out token set. The test +# writes the shards (Reverse-Text-SFT, Qwen3 tokenizer) and passes their folder as --data.paths and --val.sets. max_steps = 20 [model] @@ -20,4 +20,8 @@ micro_batch_size = 2 [optim] lr = 3e-3 +[val] +interval = 10 +sets = [{ name = "heldout", path = "set-by-the-test" }] + [ckpt] diff --git a/docs/training.md b/docs/training.md index 9f92acce368..cc24c8668d6 100644 --- a/docs/training.md +++ b/docs/training.md @@ -277,7 +277,7 @@ Pulled from the console log and mirrored to W&B. **Progress and loss:** - `loss/mean`, `loss/perplexity` — main signal. Should decrease through the run. -- `val/loss`, `val/perplexity` — validation metrics when `[val]` is set, logged every `val.interval` steps. +- `val/loss`, `val/perplexity` — validation metrics when `val.data` is set, logged every `val.interval` steps; `val//...` for `val.sets` (see [Pretraining](#pretraining)). - `eval/{env}/...` — online eval metrics when `[eval]` is set, logged at each evaluated checkpoint step. - `progress/epoch`, `progress/num_samples`, `progress/num_tokens` — dataset progress. - `progress//ratio_{samples,tokens}` — when training on multiple HF subsets/splits, the realized mixing ratio. @@ -300,6 +300,7 @@ Pulled from the console log and mirrored to W&B. | `perf/peak_memory` | peak GPU memory (GiB) | | `time/step`, `time/forward_backward`, `time/save_ckpt` | step breakdown | | `time/data/{max,mean}` | time the step waited for data, max and mean over ranks | +| `time/val` | time of a validation round | | `params/*` | total, embedding, non-embedding and active non-embedding parameter counts, at step 0 | ## Pretraining @@ -346,17 +347,18 @@ interval = 1000 **Gradient sync.** `model.defer_grad_sync = true` reduce-scatters the gradients once per optimizer step instead of after every micro-batch. It saves communication with gradient accumulation (most with experts across nodes), but keeps the unsharded parameters and gradients in memory across the micro-batches, so a config that fits without it can run out of memory. -**Log-likelihood evals.** `[ll_eval]` scores eval sets with forward passes without gradients on the training model, every `interval` steps, at the last step and (with `eval_on_start`, the default) on the weights the run starts from. Evals run after the optimizer update, so step N's metrics describe the step-N checkpoint. A set is the name of a pretrain-evals set (`ppl`, `ctxcont` or `pairs`, in `//{test.parquet,meta.json}`) or a held-out folder of token shards, scored like a `ppl` set: every document on its own, every token after its first. Every set is scored at each of `seq_lens` (default `data.seq_len`). Requests start with BOS only if the training documents do. Requires `model.cp = 1`. +**Validation sets.** `[val]` also scores log-likelihood sets on the training model (`val.sets`), with or without `val.data`. A set is the name of a pretrain-evals set (`ppl`, `ctxcont` or `pairs`, in `//{test.parquet,meta.json}`) or a held-out folder of token shards, scored like a `ppl` set: every document on its own, every token after its first. Every set is scored at each of `seq_lens` (default `data.seq_len`). Requests start with BOS only if the training documents do. As for `val.data`, step N scores the weights before step N's update. Requires `model.cp = 1`. ```toml -[ll_eval] +[val] interval = 1000 +eval_on_start = true data_dir = "/shared/evals-data" sets = ["lambada", { name = "heldout", path = "data/heldout" }] seq_lens = [2048, 8192] ``` -The metrics are `eval//@`, e.g. `eval/heldout/bpb@8k`: `loss` (nats per token), `bpb` (bits per UTF-8 byte, comparable across tokenizers), and `acc`, `acc_norm` and `greedy_em` for multiple-choice and gold-answer sets. `time/ll_eval` is the time per round. Training metrics are the SFT ones above. +The metrics are `val//@`, e.g. `val/heldout/bpb@8k`: `loss` (nats per token), `bpb` (bits per UTF-8 byte, comparable across tokenizers), and `acc`, `acc_norm` and `greedy_em` for multiple-choice and gold-answer sets. `time/val` is the time of a validation round. Training metrics are the SFT ones above. For multi-node runs, use `[deployment]` and `[slurm]` as for SFT (see [Scaling § SLURM](scaling.md#slurm)). diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index c58468a1d0b..8f382324ba4 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -174,67 +174,58 @@ class TokensDataConfig(BaseDataConfig): whole window and positions run over it.""" -class SFTValConfig(BaseConfig): - interval: int = Field(50, ge=1) - """Run validation every N training steps.""" - - eval_on_start: bool = False - """Run validation before the first training step.""" - - data: SFTDataConfig - - -DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig | TokensDataConfig, Field(discriminator="type")] - - class TokenEvalSetConfig(BaseConfig): """A held-out folder of datatrove tokenized shards, scored like a ``ppl`` set: every document on its own, in - windows of each of ``seq_lens``.""" + windows of each of ``val.seq_lens``.""" name: str - """Names the metrics, ``eval//@``.""" + """Names the metrics, ``val//@``.""" path: Path """Folder of ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``, as for ``tokens`` data.""" -class LogLikelihoodEvalConfig(BaseConfig): - """Log-likelihood evals scored inside the trainer: forward passes without gradients on the training model, every - ``interval`` steps. A set is a pretrain-evals set (``ppl``, ``ctxcont`` or ``pairs``) or a folder of held-out - token shards.""" +class SFTValConfig(BaseConfig): + interval: int = Field(50, ge=1) + """Run validation every N training steps.""" + + eval_on_start: bool = False + """Run validation before the first training step.""" - interval: int = Field(ge=1) - """Score the sets every N steps (after that step's optimizer update), and at the last step.""" + data: SFTDataConfig | None = None + """Validation dataset. Its mean loss is logged as ``val/loss`` and ``val/perplexity``.""" - eval_on_start: bool = True - """Also score the weights the run starts from (logged at step 0, or at the step a resumed run starts from).""" + sets: list[str | TokenEvalSetConfig] = [] + """Log-likelihood sets to score: the name of a pretrain-evals set (``ppl``, ``ctxcont`` or ``pairs``) in + ``data_dir``, or a held-out token folder. Logged as ``val//@``.""" data_dir: Path | None = None """Directory of the pretrain-evals sets, one directory per set with ``test.parquet`` and ``meta.json``.""" - sets: list[str | TokenEvalSetConfig] = Field(min_length=1) - """The sets to score: the name of a pretrain-evals set in ``data_dir``, or a held-out token folder.""" - seq_lens: list[int] | None = None """Context lengths to score every set at: the window length for ``ppl`` and token sets and the context limit for - ``ctxcont`` sets. Defaults to ``data.seq_len``. Metrics are named ``eval//@`` (e.g. - ``bpb@4k``).""" + ``ctxcont`` sets. Defaults to ``data.seq_len``.""" micro_batch_tokens: int | None = Field(None, ge=1) - """Tokens per packed eval row (one forward pass on one GPU). Defaults to the training micro-batch + """Tokens per packed row of the sets (one forward pass on one GPU). Defaults to the training micro-batch (``data.micro_batch_size * data.seq_len``) or the longest of ``seq_lens``, whichever is larger; must be at least the longest of ``seq_lens``.""" save_arrays: bool = False - """Write per-document arrays (npz) and summaries to ``/ll_eval/step_N/``.""" + """Write the sets' per-document arrays (npz) and summaries to ``/val/step_N/``.""" @model_validator(mode="after") - def validate_data_dir(self): + def validate_sources(self): + if self.data is None and not self.sets: + raise ValueError("[val] needs val.data or val.sets") if self.data_dir is None and any(isinstance(eval_set, str) for eval_set in self.sets): - raise ValueError("ll_eval.data_dir is required for pretrain-evals sets") + raise ValueError("val.data_dir is required for pretrain-evals sets") return self +DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig | TokensDataConfig, Field(discriminator="type")] + + class SFTModelConfig(ModelConfig): init: InitConfig | None = None """Initialize the weights from scratch (pretraining) instead of loading them from ``name``, which then only @@ -300,9 +291,6 @@ class SFTConfig(BaseTrainerConfig): val: SFTValConfig | None = None """Validation configuration. If None, no validation runs.""" - ll_eval: LogLikelihoodEvalConfig | None = None - """Log-likelihood evals scored inside the trainer. If None, they never run.""" - eval: SFTOnlineEvalConfig | None = None """Online evaluation configuration: rollout-based evals against a live inference server that receives the trainer's weight broadcasts. If None, no online evals run.""" @@ -536,7 +524,7 @@ def auto_setup_online_eval(self): @model_validator(mode="after") def validate_typed_renderer(self): """Require a typed renderer whenever SFT renders real samples.""" - if self.data.type in ("fake", "tokens") and self.val is None: + if self.data.type in ("fake", "tokens") and (self.val is None or self.val.data is None): return self model_id = self.tokenizer.name or self.model.name @@ -560,7 +548,7 @@ def validate_cp_seq_len(self): if self.model.cp > 1: if self.data.seq_len % self.model.cp != 0: raise ValueError("Sequence length must be divisible by CP degree") - if self.val is not None and self.val.data.seq_len % self.model.cp != 0: + if self.val is not None and self.val.data is not None and self.val.data.seq_len % self.model.cp != 0: raise ValueError("Validation sequence length must be divisible by CP degree") return self @@ -568,7 +556,7 @@ def validate_cp_seq_len(self): def resolve_mega_moe_max_tokens_per_rank(self): dispatch = self.model.moe.dispatch if dispatch.type == "mega_moe": - data_configs = [self.data] + ([self.val.data] if self.val is not None else []) + data_configs = [self.data] + ([self.val.data] if self.val is not None and self.val.data is not None else []) dispatch.max_tokens_per_rank = mega_moe_max_tokens_per_rank( dispatch.max_tokens_per_rank, max(data.micro_batch_size * data.seq_len for data in data_configs) ) @@ -579,7 +567,7 @@ def validate_cp_micro_batch_size(self): if self.model.cp > 1: if self.data.micro_batch_size != 1: raise ValueError("Micro batch size must be 1 when CP is enabled") - if self.val is not None and self.val.data.micro_batch_size != 1: + if self.val is not None and self.val.data is not None and self.val.data.micro_batch_size != 1: raise ValueError("Validation micro batch size must be 1 when CP is enabled") return self @@ -590,18 +578,18 @@ def validate_defer_grad_sync(self): return self @model_validator(mode="after") - def auto_setup_ll_eval(self): - if self.ll_eval is None: + def auto_setup_val_sets(self): + if self.val is None or not self.val.sets: return self if self.model.cp > 1: - raise ValueError("[ll_eval] does not support context parallelism yet") - if self.ll_eval.seq_lens is None: - self.ll_eval.seq_lens = [self.data.seq_len] - longest = max(self.ll_eval.seq_lens) - if self.ll_eval.micro_batch_tokens is None: - self.ll_eval.micro_batch_tokens = max(self.data.micro_batch_size * self.data.seq_len, longest) - elif self.ll_eval.micro_batch_tokens < longest: - raise ValueError(f"ll_eval.micro_batch_tokens must be at least the longest of ll_eval.seq_lens ({longest})") + raise ValueError("val.sets does not support context parallelism yet") + if self.val.seq_lens is None: + self.val.seq_lens = [self.data.seq_len] + longest = max(self.val.seq_lens) + if self.val.micro_batch_tokens is None: + self.val.micro_batch_tokens = max(self.data.micro_batch_size * self.data.seq_len, longest) + elif self.val.micro_batch_tokens < longest: + raise ValueError(f"val.micro_batch_tokens must be at least the longest of val.seq_lens ({longest})") return self @model_validator(mode="after") @@ -614,7 +602,7 @@ def validate_vlm_constraints(self): ) if self.data.micro_batch_size != 1: raise ValueError("VLM SFT requires data.micro_batch_size = 1.") - if self.val is not None and self.val.data.micro_batch_size != 1: + if self.val is not None and self.val.data is not None and self.val.data.micro_batch_size != 1: raise ValueError("VLM SFT requires val.data.micro_batch_size = 1.") return self diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 857d5810e2b..0305e9bcb2c 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -46,7 +46,7 @@ setup_dataset, starts_with_bos, ) -from prime_rl.trainer.sft.ll_eval import LogLikelihoodEval +from prime_rl.trainer.sft.val_sets import ValSets from prime_rl.trainer.utils import ( GarbageCollection, MemoryProfiler, @@ -211,7 +211,7 @@ def train(config: SFTConfig): dataloader = setup_dataloader(dataset, config.data, config.model.cp) val_raw_dataset = None - if config.val is not None: + if config.val is not None and config.val.data is not None: logger.info(f"Loading validation dataset ({config.val.data})") val_raw_dataset = load_sft_dataset(config.val.data) @@ -335,48 +335,49 @@ def run_eval_loop(data_iter): mean_loss = (total_loss_sum / total_token_count).item() if total_token_count.item() > 0 else float("nan") return mean_loss, nan_count.item() + val_sets = None + if config.val is not None and config.val.sets: + bos = isinstance(dataset, TokenDataset) and starts_with_bos(dataset.shards, tokenizer.bos_token_id) + val_sets = ValSets(config.val, tokenizer, bos, config.run_dir) + def run_validation(step: int) -> None: - val_dataset = setup_dataset( - tokenizer, - config.val.data, - config.model.cp, - max_epochs=1, - raw_dataset=val_raw_dataset, - renderer_config=config.renderer, - processor=processor, - multimodal=multimodal, - ) - val_dataloader = setup_dataloader(val_dataset, config.val.data, config.model.cp) - - # No train/eval switch: no dropout in these models, and toggling would trigger torch.compile recompilation - mean_loss, nan_count = run_eval_loop(val_dataloader) - if is_tt_moe_model(model): - # Keep validation routing out of the next training step's statistics. - get_load_balance_stats(model) - if nan_count > 0: - logger.warning(f"Validation at step {step}: {nan_count} batches had NaN loss") - if mean_loss != mean_loss: - logger.warning(f"Validation at step {step} had no valid tokens") - else: - logger.success(f"Validation | Step {step} | Loss {mean_loss:.4f}") - asyncio.run( - monitors.log( - {"val/loss": mean_loss, "val/perplexity": math.exp(min(mean_loss, 20)), "step": step}, step=step + start_time = time.perf_counter() + if config.val.data is not None: + val_dataset = setup_dataset( + tokenizer, + config.val.data, + config.model.cp, + max_epochs=1, + raw_dataset=val_raw_dataset, + renderer_config=config.renderer, + processor=processor, + multimodal=multimodal, ) - ) - - ll_eval = None - if config.ll_eval is not None: - bos = isinstance(dataset, TokenDataset) and starts_with_bos(dataset.shards, tokenizer.bos_token_id) - ll_eval = LogLikelihoodEval(config.ll_eval, tokenizer, bos, config.run_dir) - - def run_ll_eval(step: int) -> None: - metrics = ll_eval.run(model, step) - # The no-grad forwards leave the [lm_head, norm] FSDP group unsharded, and their routing must not count - # toward the next step's statistics. - reshard_module(model) - if moes: - take_routing_counts(moes) + val_dataloader = setup_dataloader(val_dataset, config.val.data, config.model.cp) + + # No train/eval switch: no dropout in these models, and toggling would trigger torch.compile recompilation + mean_loss, nan_count = run_eval_loop(val_dataloader) + if is_tt_moe_model(model): + # Keep validation routing out of the next training step's statistics. + get_load_balance_stats(model) + if nan_count > 0: + logger.warning(f"Validation at step {step}: {nan_count} batches had NaN loss") + if mean_loss != mean_loss: + logger.warning(f"Validation at step {step} had no valid tokens") + else: + logger.success(f"Validation | Step {step} | Loss {mean_loss:.4f}") + asyncio.run( + monitors.log( + {"val/loss": mean_loss, "val/perplexity": math.exp(min(mean_loss, 20)), "step": step}, step=step + ) + ) + metrics = {} + if val_sets is not None: + metrics = val_sets.run(model, step) + if moes: + take_routing_counts(moes) + logger.success(f"Validation sets | Step {step} | {len(val_sets.sets)} sets") + metrics["time/val"] = time.perf_counter() - start_time asyncio.run(monitors.log({**metrics, "step": step}, step=step)) def logged_micro_batches(micro_batches): @@ -386,9 +387,6 @@ def logged_micro_batches(micro_batches): ) yield micro_batch - if ll_eval is not None and config.ll_eval.eval_on_start: - run_ll_eval(checkpoint_step or 0) - gc_handler = GarbageCollection(config.gc.interval) if config.gc else None dead_experts = DeadExperts() loss_spikes, grad_norm_spikes = SpikeDetector(), SpikeDetector() @@ -529,12 +527,6 @@ def is_online_eval_step(step: int) -> bool: weight_sender.broadcast(model, step=progress.step) broadcast_weights_time = time.perf_counter() - broadcast_start_time - ll_eval_time = 0.0 - if ll_eval is not None and (is_last_step or progress.step % config.ll_eval.interval == 0): - ll_eval_start_time = time.perf_counter() - run_ll_eval(progress.step) - ll_eval_time = time.perf_counter() - ll_eval_start_time - # Optionally, dump memory snapshot if memory_profiler is not None: memory_profiler.step() @@ -545,7 +537,7 @@ def is_online_eval_step(step: int) -> bool: progress.total_samples = dataset_progress["step"] # Throughput / MFU per step over the full step wall time, as torchtitan reports them with # log_freq=1 (tokens since last log / elapsed time), instead of a smoothed sliding window. - step_time = time.perf_counter() - step_start_time - ll_eval_time + step_time = time.perf_counter() - step_start_time perf_counter.count_tokens(num_tokens) throughput = perf_counter.get_step_tokens_per_second(num_tokens, step_time) mfu = perf_counter.get_step_mfu(num_flops, step_time) diff --git a/src/prime_rl/trainer/sft/ll_eval.py b/src/prime_rl/trainer/sft/val_sets.py similarity index 94% rename from src/prime_rl/trainer/sft/ll_eval.py rename to src/prime_rl/trainer/sft/val_sets.py index e5e7c422875..abfc7c91080 100644 --- a/src/prime_rl/trainer/sft/ll_eval.py +++ b/src/prime_rl/trainer/sft/val_sets.py @@ -1,4 +1,4 @@ -"""Log-likelihood evals inside the trainer (``[ll_eval]``): scoring of text eval sets, as in the +"""The log-likelihood sets of ``[val]`` (``val.sets``): scoring of text eval sets inside the trainer, as in the `pretrain-evals` offline suite. A set lives in ``//`` as ``test.parquet`` plus ``meta.json`` (``kind``, ``parquet_sha256``). Kinds: @@ -16,7 +16,7 @@ as stored, cut into windows the same way; bpb counts the bytes of the scored tokens. A set is tokenized once (`prepare_set`); every round scores its requests (`score_requests`, the model's LM head -replaced by `scoring_head`) and summarizes them (`summarize`). `LogLikelihoodEval` runs the rounds.""" +replaced by `scoring_head`) and summarizes them (`summarize`). `ValSets` runs the rounds.""" import json import math @@ -35,7 +35,7 @@ from torch import nn from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.configs.sft import LogLikelihoodEvalConfig, TokenEvalSetConfig +from prime_rl.configs.sft import SFTValConfig, TokenEvalSetConfig from prime_rl.trainer.models.layers.lm_head import IGNORE_INDEX from prime_rl.trainer.sft.data import TokenShards, check_eos from prime_rl.trainer.utils import get_dcp_group @@ -293,9 +293,9 @@ def prepare_multi_length_set( def set_metrics(name: str, seq_len: int, summary: dict) -> dict[str, float]: - """A set's summary as metrics ``eval//@``, as Python floats (the summaries hold NumPy + """A set's summary as metrics ``val//@``, as Python floats (the summaries hold NumPy scalars).""" - return {f"eval/{name}/{key}@{length_label(seq_len)}": float(value) for key, value in summary.items()} + return {f"val/{name}/{key}@{length_label(seq_len)}": float(value) for key, value in summary.items()} def length_label(seq_len: int) -> str: @@ -414,11 +414,11 @@ def token_byte_lengths(tokenizer: PreTrainedTokenizer) -> torch.Tensor: return lengths -class LogLikelihoodEval: - """The ``[ll_eval]`` sets, tokenized once at startup and scored on the training model. Requests start with BOS if +class ValSets: + """The ``val.sets``, tokenized once at startup and scored on the training model. Requests start with BOS if ``bos`` (the training documents do).""" - def __init__(self, config: LogLikelihoodEvalConfig, tokenizer: PreTrainedTokenizer, bos: bool, run_dir: Path): + def __init__(self, config: SFTValConfig, tokenizer: PreTrainedTokenizer, bos: bool, run_dir: Path): start_time = time.perf_counter() self.config = config self.run_dir = run_dir @@ -429,13 +429,12 @@ def __init__(self, config: LogLikelihoodEvalConfig, tokenizer: PreTrainedTokeniz for eval_set in config.sets ] get_logger().info( - f"Prepared {len(self.sets)} log-likelihood eval sets (BOS: {bos}) in {time.perf_counter() - start_time:.0f}s" + f"Prepared {len(self.sets)} validation sets (BOS: {bos}) in {time.perf_counter() - start_time:.0f}s" ) @torch.no_grad() def run(self, model: nn.Module, step: int) -> dict[str, float]: - """Score every set on the weights after ``step`` optimizer updates. The metrics are on rank 0 only.""" - start_time = time.perf_counter() + """Score every set at step ``step``. The metrics are on rank 0 only.""" metrics = {} # Eager, also for compiled layers: the packed rows vary in length and would recompile at every forward. with scoring_head_active(model), torch.compiler.set_stance("force_eager"): @@ -446,11 +445,9 @@ def run(self, model: nn.Module, step: int) -> dict[str, float]: summary, arrays = summarize(version, [logprobs[i] for i in index], [hits[i] for i in index]) metrics |= set_metrics(eval_set.name, seq_len, summary) if self.config.save_arrays: - out = self.run_dir / "ll_eval" / f"step_{step}" + out = self.run_dir / "val" / f"step_{step}" out.mkdir(parents=True, exist_ok=True) label = f"{eval_set.name}@{length_label(seq_len)}" np.savez(out / f"{label}.npz", **arrays) (out / f"{label}.json").write_text(json.dumps(summary, default=float)) - metrics["time/ll_eval"] = time.perf_counter() - start_time - get_logger().success(f"Log-likelihood evals | Step {step} | {metrics['time/ll_eval']:.0f}s") return metrics diff --git a/tests/integration/test_pretrain.py b/tests/integration/test_pretrain.py index 1bc3c24ab0c..d24220be0d6 100644 --- a/tests/integration/test_pretrain.py +++ b/tests/integration/test_pretrain.py @@ -50,6 +50,8 @@ def pretrain_command(configs: list[str], token_dir: Path, wandb_project: str, wa *[arg for config in configs for arg in ("@", f"configs/ci/integration/pretrain/{config}")], "--data.paths", f'["{token_dir.as_posix()}"]', + "--val.sets", + f'[{{"name": "heldout", "path": "{token_dir.as_posix()}"}}]', "--deployment.num-train-gpus", "2", "--monitors.wandb.project", diff --git a/tests/unit/train/sft/test_ll_eval.py b/tests/unit/train/sft/test_val_sets.py similarity index 93% rename from tests/unit/train/sft/test_ll_eval.py rename to tests/unit/train/sft/test_val_sets.py index 0b28ae2723f..71d0c704c78 100644 --- a/tests/unit/train/sft/test_ll_eval.py +++ b/tests/unit/train/sft/test_val_sets.py @@ -5,7 +5,7 @@ import pyarrow.parquet as pq import torch -from prime_rl.trainer.sft.ll_eval import head_logprobs, prepare_set, set_metrics, summarize +from prime_rl.trainer.sft.val_sets import head_logprobs, prepare_set, set_metrics, summarize class CharTokenizer: @@ -41,7 +41,7 @@ def test_pairs_requests_fit_the_window(tmp_path): assert all(len(ids) - 1 <= 8 for ids, _ in eval_set.requests) logprobs = [np.full(len(ids) - start, -1.0) for ids, start in eval_set.requests] summary, _ = summarize(eval_set, logprobs, [np.ones_like(lp, dtype=bool) for lp in logprobs]) - assert set_metrics("pairs", 8, summary)["eval/pairs/acc@8"] == 0.0 # equal scores: good never beats bad + assert set_metrics("pairs", 8, summary)["val/pairs/acc@8"] == 0.0 # equal scores: good never beats bad def test_requests_without_bos_start_at_the_text(tmp_path): From 44be70fac9be6bc11eb0f106e83388aa57d31b0b Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 23 Oct 2026 00:18:55 +0200 Subject: [PATCH 130/132] feat(sft): val.eval_on_end scores the final weights after the final checkpoint --- configs/ci/integration/pretrain/start.toml | 2 ++ docs/training.md | 3 ++- packages/prime-rl-configs/src/prime_rl/configs/sft.py | 3 +++ src/prime_rl/trainer/sft/train.py | 4 ++++ 4 files changed, 11 insertions(+), 1 deletion(-) diff --git a/configs/ci/integration/pretrain/start.toml b/configs/ci/integration/pretrain/start.toml index fa8f62edce9..8520ce21deb 100644 --- a/configs/ci/integration/pretrain/start.toml +++ b/configs/ci/integration/pretrain/start.toml @@ -22,6 +22,8 @@ lr = 3e-3 [val] interval = 10 +eval_on_start = true +eval_on_end = true sets = [{ name = "heldout", path = "set-by-the-test" }] [ckpt] diff --git a/docs/training.md b/docs/training.md index cc24c8668d6..19c96306922 100644 --- a/docs/training.md +++ b/docs/training.md @@ -347,12 +347,13 @@ interval = 1000 **Gradient sync.** `model.defer_grad_sync = true` reduce-scatters the gradients once per optimizer step instead of after every micro-batch. It saves communication with gradient accumulation (most with experts across nodes), but keeps the unsharded parameters and gradients in memory across the micro-batches, so a config that fits without it can run out of memory. -**Validation sets.** `[val]` also scores log-likelihood sets on the training model (`val.sets`), with or without `val.data`. A set is the name of a pretrain-evals set (`ppl`, `ctxcont` or `pairs`, in `//{test.parquet,meta.json}`) or a held-out folder of token shards, scored like a `ppl` set: every document on its own, every token after its first. Every set is scored at each of `seq_lens` (default `data.seq_len`). Requests start with BOS only if the training documents do. As for `val.data`, step N scores the weights before step N's update. Requires `model.cp = 1`. +**Validation sets.** `[val]` also scores log-likelihood sets on the training model (`val.sets`), with or without `val.data`. A set is the name of a pretrain-evals set (`ppl`, `ctxcont` or `pairs`, in `//{test.parquet,meta.json}`) or a held-out folder of token shards, scored like a `ppl` set: every document on its own, every token after its first. Every set is scored at each of `seq_lens` (default `data.seq_len`). Requests start with BOS only if the training documents do. As for `val.data`, step N scores the weights before step N's update; `eval_on_end` also scores the final weights, after the final checkpoint, logged at the last step. Requires `model.cp = 1`. ```toml [val] interval = 1000 eval_on_start = true +eval_on_end = true data_dir = "/shared/evals-data" sets = ["lambada", { name = "heldout", path = "data/heldout" }] seq_lens = [2048, 8192] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 8f382324ba4..44192c43dd1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -192,6 +192,9 @@ class SFTValConfig(BaseConfig): eval_on_start: bool = False """Run validation before the first training step.""" + eval_on_end: bool = False + """Run validation on the final weights, after the last optimizer step and the final checkpoint.""" + data: SFTDataConfig | None = None """Validation dataset. Its mean loss is logged as ``val/loss`` and ``val/perplexity``.""" diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 0305e9bcb2c..200a4ba7575 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -662,6 +662,10 @@ def is_online_eval_step(step: int) -> bool: else: ckpt_manager.maybe_clean() + # After the final checkpoint, so a failure here cannot lose it. + if config.val is not None and config.val.eval_on_end and not paused: + run_validation(progress.step) + # Broadcast the final weights so the evals process can run its forced final epoch. A paused # run has no final epoch. if weight_sender is not None and not paused: From 847eab4205090fd072f5d0abb8776742c98e56e7 Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 23 Oct 2026 00:22:35 +0200 Subject: [PATCH 131/132] fix(sft): with val.eval_on_end, skip the in-loop validation at the last step --- packages/prime-rl-configs/src/prime_rl/configs/sft.py | 3 ++- src/prime_rl/trainer/sft/train.py | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 44192c43dd1..05ad5038014 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -193,7 +193,8 @@ class SFTValConfig(BaseConfig): """Run validation before the first training step.""" eval_on_end: bool = False - """Run validation on the final weights, after the last optimizer step and the final checkpoint.""" + """Run validation on the final weights, after the last optimizer step and the final checkpoint, instead of before + the last step's update.""" data: SFTDataConfig | None = None """Validation dataset. Its mean loss is logged as ``val/loss`` and ``val/perplexity``.""" diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 200a4ba7575..40e561f6beb 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -449,6 +449,9 @@ def is_online_eval_step(step: int) -> bool: (is_first_step and config.val.eval_on_start) or (not is_first_step and progress.step % config.val.interval == 0) ) + # The end round scores the final weights instead. + if config.val is not None and config.val.eval_on_end and progress.step == config.max_steps: + run_validation_this_step = False if gradient_manager is None: micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) else: From b816470c7117fd20d331585514f33ee82e75407c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 23 Oct 2026 00:22:51 +0200 Subject: [PATCH 132/132] fix(sft): eval_on_end keeps the eval_on_start round of a one-step run --- src/prime_rl/trainer/sft/train.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index 40e561f6beb..bd621b5fe34 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -445,13 +445,15 @@ def is_online_eval_step(step: int) -> bool: step_start_time = time.perf_counter() forward_backward_start_time = time.perf_counter() + # With eval_on_end, the end round scores the last step instead. run_validation_this_step = config.val is not None and ( (is_first_step and config.val.eval_on_start) - or (not is_first_step and progress.step % config.val.interval == 0) + or ( + not is_first_step + and progress.step % config.val.interval == 0 + and not (config.val.eval_on_end and progress.step == config.max_steps) + ) ) - # The end round scores the final weights instead. - if config.val is not None and config.val.eval_on_end and progress.step == config.max_steps: - run_validation_this_step = False if gradient_manager is None: micro_batches = (next(dataiter) for _ in range(grad_accum_steps)) else: