feat(pretrain): pretraining in the SFT trainer (token shards, from-scratch init, log-likelihood evals)

.github/workflows/tests.yml

@@ -170,6 +170,13 @@ jobs:
               - 'tests/integration/test_alphabet_sort.py'
               - 'configs/ci/integration/reverse-text{,-lora,-moe,-rl-opd,-rl-sft}/**'
               - 'configs/ci/integration/alphabet_sort.toml'
+            pretrain:
+              - *runtime
+              - 'packages/prime-rl-configs/src/prime_rl/configs/{pretrain,monitors}.py'
+              - 'src/prime_rl/entrypoints/pretrain.py'
+              - 'src/prime_rl/monitors/**'
+              - 'tests/integration/test_pretrain_moe.py'
+              - 'configs/ci/integration/pretrain-moe/**'
             eval:
               - *runtime
               - 'packages/prime-rl-configs/src/prime_rl/configs/{eval,orchestrator,inference,env_server,algorithm,monitors}.py'
@@ -184,12 +191,14 @@ jobs:
           RUN_ALL: ${{ github.event_name != 'pull_request' }}
           SFT: ${{ steps.filter.outputs.sft || 'false' }}
           RL: ${{ steps.filter.outputs.rl || 'false' }}
+          PRETRAIN: ${{ steps.filter.outputs.pretrain || 'false' }}
         run: |
-          tests=$(jq -cn --argjson all "$RUN_ALL" --argjson sft "$SFT" --argjson rl "$RL" '[
+          tests=$(jq -cn --argjson all "$RUN_ALL" --argjson sft "$SFT" --argjson rl "$RL" --argjson pretrain "$PRETRAIN" '[
               if $all or $sft then
                 {test: "reverse_text_sft", name: "Reverse Text SFT"},
                 {test: "reverse_text_sft_lora", name: "Reverse Text SFT LoRA"}
               else empty end,
+              if $all or $pretrain then {test: "pretrain_moe", name: "Pretrain MoE"} else empty end,
               if $all or $rl then
                 {test: "reverse_text", name: "Reverse Text"},
                 {test: "reverse_text_rl_opd", name: "Reverse Text RL OPD"},

configs/ci/integration/pretrain-moe/model/config.json

@@ -0,0 +1,39 @@
+{
+  "model_type": "prime_moe",
+  "vocab_size": 151680,
+  "bos_token_id": null,
+  "eos_token_id": 151645,
+  "pad_token_id": 151643,
+  "hidden_size": 256,
+  "num_hidden_layers": 4,
+  "num_attention_heads": 4,
+  "num_key_value_heads": 2,
+  "head_dim": 64,
+  "intermediate_size": 1024,
+  "hidden_act": "silu",
+  "max_position_embeddings": 4096,
+  "rms_norm_eps": 1e-06,
+  "partial_rotary_factor": 1.0,
+  "rope_parameters": {
+    "rope_type": "default",
+    "rope_theta": 10000.0,
+    "partial_rotary_factor": 1.0
+  },
+  "tie_word_embeddings": false,
+  "attention_bias": false,
+  "use_qk_norm": true,
+  "first_k_dense_replace": 1,
+  "n_routed_experts": 16,
+  "n_null_experts": 8,
+  "num_experts_per_tok": 6,
+  "moe_intermediate_size": 128,
+  "n_shared_experts": 1,
+  "norm_topk_prob": true,
+  "routed_scaling_factor": 2.5,
+  "hc_count": 2,
+  "hc_lowrank": 32,
+  "attn_output_gate": true,
+  "exclusive_attention": true,
+  "latent_dim": 128,
+  "router_bias_on_logits": true
+}

configs/ci/integration/pretrain-moe/resume.toml

@@ -0,0 +1,40 @@
+# Tiny prime_moe with the base components (Gated Residual, null experts, gated attention, XSA, LatentMoE, quantile
+# balancing) in FP8 on 2 GPUs. The test writes the token shards (Reverse-Text-SFT, Qwen3 tokenizer) and passes their
+# folder as --data.paths. Resumes the start run from its step-20 checkpoint.
+max_steps = 30
+
+[model]
+name = "configs/ci/integration/pretrain-moe/model"
+ac = "None"
+
+[model.quantization]
+type = "fp8"
+
+[model.moe.compute]
+type = "deepgemm_fp8"
+
+[model.moe.balance]
+bias_update = "quantile"
+
+[tokenizer]
+name = "PrimeIntellect/Qwen3-0.6B"
+
+[data]
+type = "tokens"
+seq_len = 2048
+batch_size = 8
+micro_batch_size = 2
+max_epochs = 100
+paths = ["set-by-the-test"]
+
+[optim]
+type = "muon"
+lr = 0.004
+
+[scheduler]
+type = "constant"
+
+[ckpt]
+
+[resume]
+step = 20

configs/ci/integration/pretrain-moe/start.toml

@@ -0,0 +1,37 @@
+# Tiny prime_moe with the base components (Gated Residual, null experts, gated attention, XSA, LatentMoE, quantile
+# balancing) in FP8 on 2 GPUs. The test writes the token shards (Reverse-Text-SFT, Qwen3 tokenizer) and passes their
+# folder as --data.paths.
+max_steps = 20
+
+[model]
+name = "configs/ci/integration/pretrain-moe/model"
+ac = "None"
+
+[model.quantization]
+type = "fp8"
+
+[model.moe.compute]
+type = "deepgemm_fp8"
+
+[model.moe.balance]
+bias_update = "quantile"
+
+[tokenizer]
+name = "PrimeIntellect/Qwen3-0.6B"
+
+[data]
+type = "tokens"
+seq_len = 2048
+batch_size = 8
+micro_batch_size = 2
+max_epochs = 100
+paths = ["set-by-the-test"]
+
+[optim]
+type = "muon"
+lr = 0.004
+
+[scheduler]
+type = "constant"
+
+[ckpt]

configs/debug/fake/pretrain-model/config.json

@@ -0,0 +1,28 @@
+{
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "eos_token_id": 151645,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 2,
+  "num_key_value_heads": 2,
+  "pad_token_id": 151643,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "5.6.2",
+  "use_cache": true,
+  "vocab_size": 151680
+}

configs/debug/fake/pretrain.toml

@@ -0,0 +1,17 @@
+max_steps = 5
+
+[model]
+name = "configs/debug/fake/pretrain-model"
+compile = "None"
+
+[tokenizer]
+name = "PrimeIntellect/Qwen3-0.6B"
+
+[data]
+type = "fake"
+seq_len = 512
+batch_size = 4
+micro_batch_size = 1
+
+[optim]
+lr = 1e-3

docs/mint.json

@@ -7,6 +7,7 @@
                 "overview",
                 "configuration",
                 "training",
+                "pretraining",
                 "eval",
                 "inference",
                 "scaling",

docs/overview.md

@@ -38,6 +38,7 @@ The `rl` entrypoint reads `examples/basic/reverse-text/rl.toml`, splits it into
 
 - **[Configuration](configuration.md)** — TOML composition, CLI overrides, dry-run.
 - **[Training](training.md)** — Launch and observe RL and SFT runs.
+- **[Pretraining](pretraining.md)** — Train a model from scratch on token shards.
 - **[Inference](inference.md)** — vLLM-backed server (or fleet) holding the current policy.
 - **[Scaling](scaling.md)** — Single-GPU through multi-node clusters via FSDP / EP / CP and SLURM.
 - **[Algorithms](algorithms.md)** — Async semantics, loss / advantage / filter plugins, trajectory merging.