Repository navigation
Conversation
…zing weights model.debug.random_init keeps its meaning (don't download or load weights). The new model.debug.random_init_mode decides what the weights hold instead: 'skip' leaves them uninitialized (the previous behavior, fastest for benchmarks), 'init' runs model.init. 'auto' resolves per entrypoint: skip for sft and rl, init for pretrain. The init draws each parameter's local shard from per-block seeds, so the values depend only on the seed and the canonical parameter name, not on the world size or sharding.
A pretraining entrypoint with its own config, launcher and training loop, built from the shared model setup, optimizer, scheduler, checkpointing and monitors. The SFT path is unchanged. - [model.arch] overrides the baseline model's HF config; the launcher writes the result to <run_dir>/model and trains from there. - Text data: documents are packed back to back into rows of seq_len tokens and continue into the next row (no truncation or padding), with optional document masking and a holdout. - [val]: loss, perplexity and bits per byte on fixed, pre-packed held-out sets. - Evals run after the optimizer update, so step N matches the step-N checkpoint.
model.init.scheme = "completep" scales init, learning rate, Adam eps and weight decay per parameter role relative to a base shape (muP for width, depth with alpha = 1). The residual-branch and logit multipliers are folded into o_proj, down_proj and the LM head, so the model and its exports stay plain Llama. model.init.truncate draws from the normal truncated at 3 std. AdamW gets an eps field (default unchanged).
setup_optimizer accepts custom parameter groups for AdamW and Muon, so the CompleteP groups go through the shared optimizer setup (fused-qkv partitions, meshes, Muon warm-up). With Muon, the hidden matrices and residual output projections are Muon groups whose learning rate is scaled by 1/sqrt(m_N) to offset the sqrt(m_N) growth of the rms_norm update. MuonConfig gets an eps field for its AdamW groups (default unchanged).
snimu
changed the title
Update each MoE layer's selection bias once per step with the DeepSeek sign rule (gamma = load_balance_coeff, now a GLM-4.5 config field) from the step's globally reduced tokens per expert; zero the bias on random init; count shared experts in MFU.
…final-stage freeze, sequence-wise balance loss
…lurm config and templates
…ad norms per group, batched perplexity evals Per-step statistics are accumulated on the GPU and reduced with one SUM and one MAX all-reduce after the micro-batches, replacing the per-micro-step MoE all-reduces and host copies. Expert balance comes from the step's global counts (max/mean, min/mean, cv, cold, entropy, dead experts) and per-rank micro-batches (mean and max over ranks); router confidence, bias size and flips; gradient norm per parameter group; loss and grad-norm spike flags; data wait and peak memory as max and mean over ranks. Perplexity evals score micro_batch_size rows per forward.
… in tokens Simplifies the pretraining path: - Data: the only real data type is datatrove token shards (the loader of the closed draft #3869, moved out of SFT): per-source folders read as one stream, mixture phases at fractions of max_steps, a startup check of the planned epochs per source, and a resume state (global row counter and tokens read per source) that does not depend on the world size. The first document of every source must end with the tokenizer's EOS, which catches shards made with another tokenizer. The HF text loader, stream packer and holdout split are gone; data.type = "fake" draws random tokens through the same path. - Batch sizes are in tokens (data.batch_tokens, data.micro_batch_tokens), so changing seq_len keeps the tokens per step and the number of steps. - Model: [model.arch] is the whole architecture (model_type plus the fields of its prime-rl config class; unknown fields are rejected). vocab_size (len(tokenizer) rounded up to 128) and the BOS, EOS and pad ids come from the tokenizer. No baseline HF model is downloaded; model.name is <run_dir>/model. - Init: the trainer initializes the weights itself unless it resumes, so model.debug.random_init_mode, its resolvers in the SFT and RL configs and the forced debug flag are gone. model.init moves into the pretrain config. - [val] reads held-out token shards per source and computes loss, perplexity and bits per byte through the training loss code, so it works with context parallelism and needs no in-memory eval sets. - Defaults: max_steps, optim.lr, data and tokenizer.name are required; the optimizer defaults to Muon.
Folds #4071 (per-hyperparameter multiplier schedules) into the pretrain PR. The example uses the default schedule (WSD, warmup 0.1, linear decay 0.9 to 0).
[schedules] (per-hyperparameter wsd/cosine/constant/custom schedules and a clip-norm schedule) becomes [schedule]: one WSD multiplier on every param group's learning rate (warmup, decay, end_ratio, decay_shape). Cosine decay is a decay_shape; weight decay already follows the learning rate. The schedulable-field map on the shared optimizer configs is gone, so the shared configs are unchanged.
Multi-node SLURM launch for pretrain through SFT's deployment, slurm config and templates. Resolved against the token-data and local-architecture changes of feat/pretrain: the launcher no longer downloads data, and the launcher test uses the new required fields.
With [eval_snapshots], the trainer writes a model-only DCP snapshot every interval steps and at the last step, in the background. `python -m prime_rl.trainer.pretrain.eval_worker` (torchrun on dedicated GPUs) builds the run's model once, scores the newest finished snapshot on the eval suite (ppl / ctxcont / pairs sets, ported from the pretrain-evals offline suite: fp32 scoring head, packed document-masked rows, rows dealt over the worker's GPUs), logs eval/{set}/{metric} at the snapshot's step through the run's monitors, deletes scored snapshots and stops after the last step.
The model already reads config.head_dim (falling back to hidden_size // num_attention_heads). As an explicit field it is part of the config class, so [model.arch] accepts it.
…oops The forward/backward loop over the micro-batches of a step (token-normalized gradient scaling, the offload path, NaN micro-steps, MoE routing statistics, the global loss) was copied between trainer/sft/train.py and trainer/pretrain/train.py. It is now trainer/step.py:run_micro_batches.
… roles Pretraining builds its optimizer with the default groups of SFT and RL instead of its own per-role groups, and the init classifies 1-D parameters directly, with the same values.
…the training data Mixing moves offline: data.sources / data.weights become data.paths, a list of shard folders read as one token pool (folder order, then file order), with one seeded random window order per epoch; max_epochs is checked on the pool. The per-row source draw, the per-source epoch planning and token accounting, the weight validator and the per-source progress metrics go away; progress/epochs is the pool's. TrainData (pretrain/data.py) owns the dataset, its dataloader, the resume state and the count of windows trained on: train() builds it, iterates it and checkpoints it. Its checkpoint state is the dataloader state plus the window count.
Main's default Muon groups put the hyper_connection projections on Muon, which trains them worse (d384 with Gated Residual: +0.029 val loss over 4 seeds). Pretraining moves them to the AdamW group.
Since the main merge, the LM-head wrapper passes topk_ids to the head's forward; the scoring head that replaces it for the eval suite did not take it, so every run with [evals] crashed at its first eval.
[evals] scores both kinds of sets on one interval: evals.sets lists pretrain-evals set names (in evals.data_dir, as before) and held-out token folders ({ name, path }). A token folder is scored like a ppl set through the suite's scoring path: every document as stored, in windows of each of seq_lens, every token after its first scored; bpb counts the bytes of the scored tokens. Its metrics are eval/<name>/loss@<length>, bpb@<length>, ...
ValConfig, TokenSourceConfig, run_val and the val windows go, and with them the seq-balance guard for no-grad forwards and the unseeded window order of TokenShards. The W&B overview and the dashboard drop their val/*/bpb panels.
…and why old val numbers don't compare examples/pretrain/prime_moe.toml scores data/dolma3-50b-dsv4/mix-val-5M instead of three whole source val folders. docs/pretraining.md shows it, gives the measured cost at d512 (about 1 s per 200K tokens per length on one H200; 5M tokens at [2048, 8192] in 40 s against 280 s for the suite) and says why token-set losses sit 0.014-0.027 nats below the former [val].
The flag now initializes the weights (N(0, 0.02), norms 1, biases 0) instead of leaving them uninitialized, so it is named for what it does.
snimu
changed the title
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
| from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens | ||
|
|
||
|
|
||
| class PrimeMoeDecoderLayer(GradientCheckpointingLayer): |
There was a problem hiding this comment.
Missing KL mismatch table
Medium Severity
This PR adds a new custom model (prime_moe), but it does not include a table of mean KL mismatch across 20 steps on the math environment with batch_size=64, with every entry below 0.015.
Triggered by project rule: BugBot Instructions
Reviewed by Cursor Bugbot for commit bd856a3. Configure here.
garrett361
left a comment
garrett361
left a comment
There was a problem hiding this comment.
This PR makes various logically independent changes that I'd prefer to live in their own PRs.
And also the headline item of adding pretraining is largely copying over the SFT loop, since SFT and pretraining are so closely related to each other. I think it's simpler to just use the sft loop and logic, maybe making some filenames and functions more generic, than to copy over code for a nearly identical loop. Or do you see benefits to having the two be separate?
In my mind, SFT and pretraining are effectively identical except for the data and initialization, and some of the features you'll be adding (like auxiliary-loss-free expert load balancing) are features we'd likely want to experiment with for SFT as well. Keeping the loops unified would help avoid needing to duplicate features. Let me know your thoughts.
Asks below are written assuming a separate pretraining loop is kept.
Claude generated below, intended to pass to your agent.
Reviewed at head bd856a3 against merge base 477d438. The asks, in order:
- Fix two bugs:
perf/mfuis computed from tokens, andmodel.debug.random_initnow crashes several existing models. - Remove
compile_inside_full_checkpoint, which re-lands the layout #4037 reverted. - Keep config changes purely additive; only small shared helpers are refactored here.
- Split
prime_moe(with its architecture options) and MoE load balancing into a stack of follow-up PRs.
Bugs
1. perf/mfu is computed from tokens instead of FLOPs
pretrain/train.py#L373-L376 calls perf_counter.get_step_mfu(num_tokens, step_time), but since #3840 get_step_mfu takes FLOPs (perf.py#L86-L87), so the logged MFU is too small by a factor of 6N. The same block passes config.data.seq_len as get_perf_counter's window_size argument (perf.py#L133). Expected: compute FLOPs the way SFT does (sft/train.py#L558-L559), log perf/mfu and perf/mfu_6dn, and call get_perf_counter(model). The MFU numbers in the PR description therefore cannot come from this head.
2. model.debug.random_init now crashes on several existing models
model.py#L655 routes the flag through init_weights, which raises ValueError("No init for 1-D parameter ...") on any 1-D parameter whose name is not a bias or a norm (init.py#L61-L67). On main the flag runs on every model; with this PR it crashes on:
- gpt_oss:
sinks(gpt_oss/attention.py#L50) - every GatedDeltaNet model (qwen3_5, qwen3_8_flash_next):
A_log(gated_delta_net.py#L56) - nemotron_h:
A_log,D - deepseek_v4:
sinks,base,scale
This needs to be fixed before merge. A complete init design (every model family, per-role scales) needs its own plan later.
Remove compile_inside_full_checkpoint
Remove compile_inside_full_checkpoint (model.py#L773-L786, #L951-L952). It turns compile(AC(block)) into AC(compile(block)) for every block under full AC. That is the layout of #3925, which #4037 reverted because MoE runs with CP, AC and compile crashed in an early backward (aten.empty_like.default ... not found in storage), with the root cause still unknown. Pretraining supports CP, so it is exposed to the same crash. Re-land it only after the root cause from #4037's follow-up is found, for all trainers at once.
Try to keep this PR additive
Goal: config changes in this PR are purely additive: new pretrain config classes, with no SFT, RL or shared config class changed or moved. Small helper functions that SFT and pretraining both need can be refactored into shared code here, but let's try to keep the diff to existing code small. Larger shared changes can follow in their own PRs once pretraining has landed.
Config: revert to main and add instead.
-
configs/sft.py: leave it as on main.PackedDataConfigexists only to keeppad_to_seq_lenout of pretrain's data config. Instead, haveBaseBatchConfiginherit main'sBaseDataConfigand rejectpad_to_seq_len = true(pretrain rows are always full), or define its own four fields. -
TrainerLaunchConfig(trainer.py#L755) and the move of the deployment configs toconfigs/shared.py: revert.PretrainConfigdeclares its own launch fields and imports the deployment configs fromconfigs/sft.py. This costs about 90 duplicated lines inconfigs/pretrain.py; deduplicating them is a follow-up. -
Fractional
warmup_steps/decay_stepsanddecay_shapechange the sharedLinearSchedulerConfig(trainer.py#L404-L414). Use integer steps here and land the scheduler change in its own PR. -
moe_runtime.py#L89-L90weakens main's check tocompute/dispatchonly becausePretrainMoEConfigaddsbalance. Any field later added toMoERuntimeConfigwould then skip the check silently. Keep main's check and compare the runtime part on the pretrain side:runtime = MoERuntimeConfig(**config.moe.model_dump(include=set(MoERuntimeConfig.model_fields))) if runtime != MoERuntimeConfig(): raise ...
-
Revert changes no pretrain caller needs:
scheduler: LRScheduler | Noneinckpt.py#L181and#L273. Pretraining always saves with a scheduler (pretrain/train.py#L362); only the load takesNone, and main's load already accepts it.
Helpers: sharing these is fine, but they should live in a neutral module.
-
resolve_resume_step(pathing.py#L339) andget_grad_accum_steps(trainer/utils.py#L369) are fine as they are. -
Shared code should not import pretraining.
entrypoints/sft.py#L10importsPretrainConfigfor the type hints ofwrite_configandwrite_trainer_slurm_script, andentrypoints/pretrain.pyimports them fromentrypoints/sft.py. Move the shared launcher helpers to a neutral module (e.g.utils/slurm.py) that both entrypoints import.monitors/prime.py#L17importsPretrainConfigonly for anisinstancecheck; have the caller pass the run flavor instead. -
The
{{ entrypoint }}variable insingle_node_sft.sbatch.j2#L44andmulti_node_sft.sbatch.j2#L74is fine: three lines, SFT output unchanged, and copying the 282-line multi-node template would be worse. But pretraining should not be launched through files named*_sft. Rename both templates to trainer-neutral names (e.g.single_node_trainer.sbatch.j2,multi_node_trainer.sbatch.j2) and point both configs' defaulttemplate_pathat them. -
count_parameters(trainer/utils.py#L381) wraps main'scount_active_params(models/base.py#L20-L42) but defines embeddings differently: only the input embedding and LM head, wherecount_active_paramstreats everynn.Embedding(e.g. qwen3_8_flash_next's n-gram table) as an embedding. For models like that,params/non_embeddingcounts tables thatparams/active_non_embeddingleaves out. Follow thenum_active_params()pattern instead: one method onPreTrainedModelPrimeRLper count, all built on one embedding rule so they agree by construction:def embedding_params(self) -> list[nn.Parameter]: # the rule count_active_params uses return [p for m in self.modules() if isinstance(m, nn.Embedding) for p in m.parameters(recurse=False)] + lm_head params def num_params(self) -> int: ... def num_embedding_params(self) -> int: ... # from embedding_params() def num_non_embedding_params(self) -> int: ... # num_params() - num_embedding_params() def num_active_non_embedding_params(self) -> int: ... # count_active_params over the non-embedding params
Suggestion: split the rest into a stack
Goal: this PR adds pretraining for the models prime-rl already has; the new research model and the load-balancing methods come as stacked PRs, each with its own tests and docs.
- Pretraining core (this PR): token-shard data, the minimal init, the pretrain metrics and in-trainer evals, and the dense integration test. No model-code changes and no load balancing. Existing MoE models can pretrain, without bias updates.
prime_moe: the research MoE model (prime_moe/) with all its architecture options in one PR. With every option off it is GLM-4.5's MoE architecture withglm4_moe's state-dict layout. Its options and the shared changes they need:- Gated Residual, with the
hyper_connectionmove (main has one copy, so the description's "two byte-identical copies" is inaccurate), its init cases ininit.py, and the AdamW group inpretrain/optim.py. Instead ofsetup_optimizer(param_groups=...)(optim/__init__.py#L60: unguarded under full offload, silently falls back on[], ignored by SignSGD, andpretrain/optim.pyimports the private_default_muon_param_groups), add a name predicate that moves parameters from the Muon group to AdamW; - null experts (variable k): replace the
torch.bincounthost sync (prime_moe/moe.py#L253) withscatter_add_, as line 92 of the same file does; - XSA and gated attention, with the
project_outputhook (attn.py#L564-L569); - LatentMoE and SiTU-GLU experts.
- Gated Residual, with the
- MoE load balancing: Main's shared MoE keeps a selection bias and per-expert token counts (
layers/moe.py#L162-L169) but never updates the bias.- The sign-rule bias update (
update_selection_biases) uses only those buffers, so write it against the shared MoE and let every MoE model use it; drop theprime_moe-only check inbalance.py#L14-L15for it. - The quantile bias update (histogram estimator,
router_bias_on_logits) and the sequence-wise balance loss needPrimeRouter/PrimeMoE, so they stayprime_moe-only. The bias freeze applies to both updates. - Fix there: checkpoint
quantile_grid/quantile_rangeso resume is exact (balance.py#L27-L28); the "resume reproduces steps 21-40" result is from an earlier head. - Fix there: stop recording histograms after the freeze (
prime_moe/moe.py#L107,train.py#L306) by settingrouter.quantile_balance = Falseat the freeze. - Fix there: under CP,
sequence_balancetreats each rank's CP chunk of a window as a whole window; reduce across the CP group or document it.
- The sign-rule bias update (
Does pretraining need a separate training loop?
Once the model work moves to the stack, little in this PR needs its own loop.
About 180 of the 379 non-blank lines of pretrain/train.py closely match sft/train.py (setup, resume, compute_loss, clip/step/scheduler/checkpoint). Bug 1 is what that duplication costs: main changed the perf API and only the SFT copy was updated. What remains different after the split:
- Random init: already just
setup_model(init=...). - Token-shard data:
TokenDatasetalready subclasses SFT'sStatefulIterableDatasetand emits SFT's batch keys with an all-ones loss mask. - A micro-batch loop with no host syncs that syncs gradients once per step (
step.py#L42). SFT would benefit from this too. - The in-trainer evals and the extra metrics (and, later in the stack, the MoE bias update): optional per-step hooks.
SFT's features that pretraining doesn't use (LoRA, VLM, online evals, the control server, full offload) are opt-in, so a pretraining config would simply leave them off. One structural difference to handle: SFT's CatDataset records each sample as one attention segment (sft/data.py#L543), while a pretrain window can hold several documents with their own seq_lens. If the loops are unified, the shape would be roughly this (the new data type and init field are additions to SFT's config, not changes to existing fields):
# configs/sft.py
DataConfig = Annotated[FakeDataConfig | SFTDataConfig | TokensDataConfig, Field(discriminator="type")]
# sft/train.py
model = setup_model(config.model, parallel_dims, init=config.model.init if config.from_scratch else None)
result = run_micro_batches(model, micro_batches, ...) # shared, moved out of pretrain/step.py
for hook in step_hooks: # bias update, evals, pretrain metrics
hook.after_optimizer_step(step, result)Smaller issues
-
Gradient sync once per optimizer step is always on:
step.py#L80-L81disablesset_requires_gradient_syncandset_reshard_after_backwardon every micro-batch of gradient accumulation except the last, so FSDP reduce-scatters only once per optimizer step but keeps the unsharded parameters and gradients in memory across the micro-batches. A config that fits under SFT can OOM here. Make it a config option. -
Integration test:
test_pretrain_moe.pyruns a tinyprime_moewith the opt-in options and FP8 on. Mirrortest_reverse_text_sft.pyinstead: one small dense model, start and resume, with no FP8 or other extras. The MoE test goes with theprime_moePR. The CI filter (tests.yml#L173-L179) includes*runtimelike SFT's, so the test runs on every trainer change, not "only when pretraining files change" as the description says; fix the description. -
Unit tests: did any of the tests in
tests/unit/train/pretrain/catch a real failure while you worked on this? If some are needed as regression tests, keep those and put each in its natural home: with the PR that adds the code it tests, and next to the existing tests when the code is generic (e.g.test_grad_norms_by_group_add_up_to_the_total_normtests a function that works on any model). Drop the rest. -
Integration configs:
pretrain-moe/resume.tomlrepeats all 37 lines ofstart.tomlto change three (max_steps = 30,[resume] step = 20). In the new dense test's configs, makeresume.tomlan overlay with only the changed lines and compose it as@ start.toml @ resume.toml. -
Metric names: pretraining logs some of SFT's quantities under different names, so panels and comparisons don't carry across flavors. Keep SFT's names (renaming would break existing runs' history), and add pretrain's extra metrics to SFT as well so both flavors log them. Note
expert_load/max_meanismax_vio + 1:Quantity SFT Pretrain Epochs progress/epochprogress/epochsPeak memory perf/peak_memoryperf/peak_memory/max,/meanMoE imbalance max_vio/mean,max_vio/maxexpert_load/max_mean/*Router confidence routing_confidence/meanrouter/confidence -
docs/pretraining.mdis 2.3k words, against about 1.5k for the whole SFT section ofdocs/training.md. Cut it to what a user needs to launch and monitor a run: launch commands, the token-shard data format and[data]fields, the[evals]block, MoE balancing fields, and the metrics. Remove:- the comparison with the former
[val]numbers (AGENTS.md: docs describe main, not history); - "SFT and RL do not pass one, so their behavior is unchanged", which belongs in the PR description;
- timings and memory ratios from one setup (33 ms per CephFS read, 1 s per 200K eval tokens, 0.6x peak, 19%);
- the
prime_moearchitecture list, which belongs in thePrimeMoeConfigfield docstrings; - the SLURM setup repeated from
scaling.md. Link to it and keep only the paragraph on how rows map to ranks.
- the comparison with the former
…es unknown 1-D parameters
… balancing and the fractional scheduler out of this PR They move to the stacked branches feat/pretrain-prime-moe and feat/pretrain-moe-balance, and the scheduler change to feat/scheduler-fractional. ckpt.py and moe_runtime.py are back to main.
The separate pretrain trainer, config, launcher and entrypoint are gone. Pretraining runs with `uv run sft`:
- data.type = "tokens": datatrove token shards as SFT data (TokenDataset); CatDataset keeps every document of a
window as its own segment
- model.init: initialize the weights from scratch instead of loading them
- [ll_eval]: the log-likelihood evals, scored inside the trainer
- the micro-batch loop moves to trainer/micro_batches.py: no host sync per micro-batch, the step statistics are
reduced once, with the same gradients and metric names as before; model.defer_grad_sync reduce-scatters the
gradients once per optimizer step
- extra metrics for every run: params/*, loss/spike, optim/grad_norm/<group>, optim/grad_norm_spike, time/data/*,
expert_load/{min_mean,entropy,dead,layer_N/*}
- parameter counts are methods of PreTrainedModelPrimeRL, built on one embedding rule
…, in the SFT CI group
… docs/pretraining.md; the new metrics
…TModelConfig; the test passes data.paths as a list
snimu
changed the title
… of [ll_eval] val.sets (pretrain-evals sets or held-out token folders), val.data_dir, val.seq_lens, val.micro_batch_tokens and val.save_arrays; val.data becomes optional. They run where val.data does, logged as val/<set>/<metric>@<length>, and time/val times the validation round. SFT's val.data path is unchanged.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
There are 2 total unresolved issues (including 1 from previous review).
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 44be70f. Configure here.
Adds pretraining from scratch on pre-tokenized token shards. Pretraining runs in the SFT trainer (
uv run sft): it is SFT with a token-shard data type and a from-scratch init. There is no separate pretrain entrypoint, config or training loop.Existing SFT configs train identically: same defaults, same gradients and the same metric names; SFT gains the extra metrics below. RL is unchanged.
What this PR adds
Token-shard data (
data.type = "tokens",TokensDataConfigin SFT's data union):DocumentTokenizershards. The folders indata.pathsare read as one token pool. Each epoch reads everyseq_lenwindow once, in a seeded random order from a random offset.TokenDatasetsubclasses SFT'sStatefulIterableDatasetand yields SFT samples. A window carries oneseq_lensentry per document, andCatDatasetnow keeps a sample's segments (one per sample for SFT data, as before). Attention and positions restart at every document (data.document_masking) and at every window.From-scratch init (
model.init:std,seed;SFTModelConfigadds it to the trainer model config). Every weight matrix and embedding starts from N(0, std), norms at 1, and biases and other 1-D parameters at 0. The values do not depend on the world size or the sharding.model.debug.random_initbehaves as on main.Log-likelihood validation sets in
[val](val.sets,val.data_dir,val.seq_lens,val.micro_batch_tokens,val.save_arrays;val.databecomes optional). They score pretrain-evals sets (ppl / ctxcont / pairs) and held-out token folders on the training model. They run whereval.dataruns, and the metrics areval/<set>/<metric>@<length>.val.eval_on_end(default off) also runs a validation round on the final weights, after the final checkpoint. SFT'sval.datapath is unchanged (val/loss,val/perplexity).time/valtimes a validation round.Shared micro-batch loop (
trainer/micro_batches.py, used by SFT and so by pretraining). There is no host sync per micro-batch, and the step statistics are reduced with one SUM and one MAX all-reduce. The gradients are the same as before. The logged loss is now summed in fp64, so it can differ from main in the last fp32 bits.model.defer_grad_sync(default off): reduce-scatter the gradients once per optimizer step instead of after every micro-batch. It saves communication but keeps the unsharded parameters and gradients in memory across micro-batches. Pretraining configs turn it on.Extra metrics for every SFT run:
params/*at step 0loss/spike,optim/grad_norm_spikeandoptim/grad_norm/<group>time/data/{max,mean}andtime/valexpert_load/{min_mean,entropy}/*,expert_load/layer_<i>/*andexpert_load/deadSFT's existing names (
progress/epoch,perf/peak_memory,max_vio/*,routing_confidence/mean, ...) are kept, and pretraining logs them too. With token data,progress/epochis the fraction of the pool read.Parameter counts are methods of
PreTrainedModelPrimeRL(num_params,num_embedding_params,num_non_embedding_params,num_active_non_embedding_params), built on one embedding rule (embedding_params).Small shared changes:
setup_model(init=...).compute_max_vio_and_routing_confidenceis factored out ofget_load_balance_stats; RL is unchanged.tools/convert_dcp_to_bf16.pyreads an SFT run's model config asSFTModelConfig.Docs: a short pretraining section in
docs/training.md.Test:
tests/integration/test_pretrain.py(small dense model from scratch, start + resume overlay, with a held-out token set in[val]). It runs in the SFT CI group, so on every change to the trainer runtime or the SFT files.Moved out of this PR (branches, no PRs yet)
feat/pretrain-prime-moe(on this PR): theprime_moemodel and the shared changes it needs:hyper_connection.pymove and theproject_outputattention hooktrains_with_adamw)scatter_add_instead ofbincountfeat/pretrain-moe-balance(onprime_moe):[model.moe_balance]:prime_moeonlyfeat/scheduler-fractional(on main): fractionalwarmup_steps/decay_stepsanddecay_shape.feat/pretrain-full-2026-10-22: the previous full state of this PR, kept for reference.compile_inside_full_checkpoint, which re-landed the layout that revert(trainer): compile around the checkpoint again for MoE blocks with an fp32 router #4037 reverted.Validation (H100 cluster, at 2fd8b9b; the
[val]changes at e2c3465)samsja/mini-glm-moe), 2 GPUs dense with 4 micro-batches of gradient accumulation, and 2 GPUs MoE with EP=2.optim/grad_normand loss. The only difference isloss/meanat about 5e-8 relative, from the fp64 loss sum.[val]withval.data, main vs this branch:val/lossandval/perplexityare identical at every validation step. The only new key istime/val.tests/integration/test_reverse_text_sft.py: passes, excepttest_dashboard. That test cannot install playwright's browser dependencies without root on the cluster.tests/integration/test_pretrain.py(dense, start + resume, with a held-out token set in[val]): passes, and logsval/heldout/*@1k.val.sets(pretrain-evals sets and a held-out token folder, two lengths) andmodel.defer_grad_sync = true: runs, and logsval/<set>/<metric>@<length>.test_pretrain_moe.py(start + resume) passes onfeat/pretrain-prime-moeand onfeat/pretrain-moe-balance. Also run:prime_moesmoke with the validation setsquantile_gridMFU numbers from earlier heads are removed; they were not re-measured at this head.
Note
Medium Risk
Touches core SFT training, FSDP grad sync, and validation paths; behavior is intended to match existing SFT aside from new metrics and fp64 loss aggregation, but the surface area is large.
Overview
Pretraining is added to the existing SFT path (
uv run sft):data.type = "tokens"reads datatrove token shards as one pool with seeded window order;[model.init]builds weights from scratch (sharding-independent RNG); optionalmodel.defer_grad_syncdefers FSDP gradient reduce-scatter to the last micro-batch.Validation expands beyond HF
val.data: optionalval.setsscores pretrain-evals parquet sets and held-out token folders on the training model (val/<set>/<metric>@<length>), withval.eval_on_endfor a final round after the last checkpoint.val.datastays optional when sets are configured.SFT training loop is refactored into shared
run_micro_batches/metricshelpers (fewer per-micro-batch host syncs, fp64 loss sum). All SFT runs gain extra observability (params/*, grad-norm by group, loss/grad spikes, data wait times, richer MoE expert-load metrics).CI and docs: integration test
test_pretrain(start + resume, held-out token val), CI matrix entry, and a Pretraining section indocs/training.md. Checkpoint export usesSFTModelConfigwhen readingsft.json.Reviewed by Cursor Bugbot for commit b816470. Bugbot is set up for automated code reviews on this repo. Configure here.