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..8520ce21deb --- /dev/null +++ b/configs/ci/integration/pretrain/start.toml @@ -0,0 +1,29 @@ +# 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] +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 + +[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 4117e6f7d65..19c96306922 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,7 +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 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). | @@ -276,17 +277,18 @@ 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. **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:** @@ -297,6 +299,69 @@ 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 | +| `time/val` | time of a validation round | +| `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. + +**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] +``` + +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)). ## Checkpointing 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 7e578e5f0e0..05ad5038014 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -20,7 +20,7 @@ SlurmConfig, WeightBroadcastConfig, ) -from prime_rl.configs.trainer import BaseTrainerConfig, mega_moe_max_tokens_per_rank +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 @@ -155,6 +155,36 @@ 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 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 ``val.seq_lens``.""" + + name: str + """Names the metrics, ``val//@``.""" + + path: Path + """Folder of ``*.ds`` token files, each with its ``.ds.index`` and ``.ds.metadata``, as for ``tokens`` data.""" + + class SFTValConfig(BaseConfig): interval: int = Field(50, ge=1) """Run validation every N training steps.""" @@ -162,10 +192,52 @@ class SFTValConfig(BaseConfig): eval_on_start: bool = False """Run validation before the first training step.""" - data: SFTDataConfig + eval_on_end: bool = False + """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``.""" + + 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``.""" + + 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``.""" + + micro_batch_tokens: int | None = Field(None, ge=1) + """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 the sets' per-document arrays (npz) and summaries to ``/val/step_N/``.""" -DataConfig: TypeAlias = Annotated[FakeDataConfig | SFTDataConfig, Field(discriminator="type")] + @model_validator(mode="after") + 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("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 + 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): @@ -213,6 +285,8 @@ class SFTConfig(BaseTrainerConfig): # 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.""" @@ -454,7 +528,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 or self.val.data is None): return self model_id = self.tokenizer.name or self.model.name @@ -478,7 +552,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 @@ -486,7 +560,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) ) @@ -497,10 +571,31 @@ 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 + @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_val_sets(self): + if self.val is None or not self.val.sets: + return self + if self.model.cp > 1: + 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") def validate_vlm_constraints(self): if self.model.vlm is None: @@ -511,7 +606,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/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 222aca108f7..01b01e6749a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -261,6 +261,17 @@ class MoERuntimeConfig(BaseConfig): dispatch: MoEDispatchConfig = TorchMoEDispatchConfig() +class InitConfig(BaseConfig): + """Random weight initialization: every weight matrix and embedding from N(0, std). Norm weights start at 1, + biases and other 1-D parameters at 0.""" + + std: float = Field(0.02, gt=0) + """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.""" diff --git a/src/prime_rl/trainer/metrics.py b/src/prime_rl/trainer/metrics.py new file mode 100644 index 00000000000..4c5a3ce9451 --- /dev/null +++ b/src/prime_rl/trainer/metrics.py @@ -0,0 +1,146 @@ +import math +from collections import defaultdict, deque +from collections.abc import Iterable + +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 compute_expert_load_stats +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]: + """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(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.""" + 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 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: `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) + 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}/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 + + +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 any("hyper_connection" in part for part in parts): + return "hyper_connection" + 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 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, 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(tensors)).double().square().sum() for tensors 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()} + + +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/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 39ac90e7cbf..64aac93649d 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -31,6 +31,7 @@ ActivationCheckpointConfig, CompileConfig, FP8Config, + InitConfig, ModelConfig, MXFP8Config, TokenizerConfig, @@ -54,6 +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 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 @@ -204,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, @@ -229,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_() @@ -872,7 +883,10 @@ def setup_model( config: ModelConfig, parallel_dims: ParallelDims, loading_from_checkpoint_later: bool = False, + init: InitConfig | 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 instead.""" resolve_auto_attn(config) if config.attn == "flash_attention_3" and not is_flash_attn_3_available(): @@ -936,9 +950,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) @@ -948,6 +964,8 @@ def setup_model( load_dcp_from_hf(model, config) _reset_runtime_moe_buffers(model) + if init is not None: + init_weights(model, init) return model 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/models/init.py b/src/prime_rl/trainer/models/init.py new file mode 100644 index 00000000000..b7ed7ff8313 --- /dev/null +++ b/src/prime_rl/trainer/models/init.py @@ -0,0 +1,88 @@ +import hashlib +import math +import re + +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.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 +# blocks that overlap its shard and the values do not depend on how the tensor is sharded. +BLOCK_NUMEL = 1 << 22 + + +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 and the MoE selection biases in place: weight matrices and embeddings from + 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) + local = param.to_local() if isinstance(param, DTensor) else param + if param.ndim == 1: + if name.endswith("bias"): + local.zero_() + elif "norm" in name: + local.fill_(1.0) + else: + local.zero_() + zeroed.append(name) + 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_(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. + 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/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/sft/train.py b/src/prime_rl/trainer/sft/train.py index 3124e936eac..bd621b5fe34 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,7 +10,7 @@ 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_latest_ckpt_step from prime_rl.configs.sft import SFTConfig @@ -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,23 +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.val_sets import ValSets 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, ) @@ -143,7 +143,25 @@ def train(config: SFTConfig): # 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) @@ -193,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) @@ -233,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.""" @@ -284,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") @@ -319,37 +335,61 @@ 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, ) - ) + 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): + for micro_batch in micro_batches: + print_sample( + micro_batch["input_ids"].flatten().tolist(), micro_batch["loss_mask"].flatten().tolist(), tokenizer + ) + yield micro_batch 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. @@ -383,7 +423,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) @@ -406,25 +445,17 @@ 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 {} - ) + # 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) + ) ) 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) @@ -441,58 +472,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) @@ -503,15 +499,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}") @@ -548,10 +536,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_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"] @@ -566,13 +551,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) @@ -618,14 +604,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 @@ -646,9 +635,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 @@ -677,6 +667,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: diff --git a/src/prime_rl/trainer/sft/val_sets.py b/src/prime_rl/trainer/sft/val_sets.py new file mode 100644 index 00000000000..abfc7c91080 --- /dev/null +++ b/src/prime_rl/trainer/sft/val_sets.py @@ -0,0 +1,453 @@ +"""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: + +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 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`). `ValSets` runs the rounds.""" + +import json +import math +import time +import types +from collections.abc import Iterator +from contextlib import contextmanager +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.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 +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.""" + +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 | None + + +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.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, 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]) + 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} + + +@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]]: + """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 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, 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 [*prefix, *encoding["input_ids"]], first + len(prefix) + + +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, 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() + 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) + # 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, 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 = [([*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})") + 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 + 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( + eval_set: str | TokenEvalSetConfig, + data_dir: Path | None, + tokenizer: PreTrainedTokenizer, + seq_lens: list[int], + bos: bool, + byte_lengths: np.ndarray | None = None, +) -> MultiLengthSet: + """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] + for seq_len, version in versions.items() + } + 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 ``val//@``, as Python floats (the summaries hold NumPy + scalars).""" + return {f"val/{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) + + +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["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(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(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() + summary = { + "loss": nll.sum() / tokens.sum(), + "bpb": nll.sum() / (num_bytes.sum() * math.log(2)), + "docs": len(num_bytes), + "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} + + +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 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: SFTValConfig, 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)} 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 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"): + 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 / "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)) + return metrics diff --git a/tests/integration/test_pretrain.py b/tests/integration/test_pretrain.py new file mode 100644 index 00000000000..d24220be0d6 --- /dev/null +++ b/tests/integration/test_pretrain.py @@ -0,0 +1,107 @@ +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)) 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/sft/test_token_dataset.py b/tests/unit/train/sft/test_token_dataset.py new file mode 100644 index 00000000000..ed3b7202e7f --- /dev/null +++ b/tests/unit/train/sft/test_token_dataset.py @@ -0,0 +1,38 @@ +from itertools import islice +from pathlib import Path + +import numpy as np + +from prime_rl.trainer.sft.data import CatDataset, TokenDataset, TokenShards + + +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=" 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"] diff --git a/tests/unit/train/sft/test_val_sets.py b/tests/unit/train/sft/test_val_sets.py new file mode 100644 index 00000000000..71d0c704c78 --- /dev/null +++ b/tests/unit/train/sft/test_val_sets.py @@ -0,0 +1,63 @@ +import json + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import torch + +from prime_rl.trainer.sft.val_sets import head_logprobs, prepare_set, set_metrics, summarize + + +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_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) + + +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, 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)["val/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) 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