Skip to content

Set 3B pre-training budget to 60B tokens (stage1 + stage2) per model - #52

Merged
iamtatsuki05 merged 1 commit into
developfrom
config/en-3b-stage1-0.65-stage2-1epoch
Sep 14, 2026
Merged

iamtatsuki05 merged 1 commit into
developfrom
config/en-3b-stage1-0.65-stage2-1epoch

Conversation

@iamtatsuki05

@iamtatsuki05 iamtatsuki05 commented Sep 14, 2026 •

Copy link
Copy Markdown
Owner

Why

The 3B pre-training budgets were epoch-based and uneven across models (EN 47B / 40B, JP 65B / 89B tokens). This PR fixes the budget to 60B tokens per model, stage1 + stage2 combined, mirroring how the 1B stage2 runs were started from a stage1 checkpoint that brought the total to ~20B.

Change

config before after stage1 tokens stage2 total
Llama-EN-3B/stage1.json 0.45 epoch 0.5473 epoch (13,700 steps) 57.5B 1 epoch, 2.5B 60.0B
Llama-JP-3B/stage1.json 0.8 epoch 0.7127 epoch (13,800 steps) 57.9B 6 epoch, 2.2B 60.1B
ModernBERT-EN-3B/stage1.json 0.45 epoch 0.65 epoch (13,893 steps) 57.3B 1 epoch, 2.45B 59.8B
ModernBERT-JP-3B/stage1.json 0.8 epoch 0.5214 epoch (13,800 steps) 57.7B 6 epoch, 2.2B 59.9B
Llama-EN-3B/stage2.json, ModernBERT-EN-3B/stage2.json 3 epochs 1 epoch (same as 1B)

Step counts assume total batch 4096 × seq 1024 in stage1; per-step token actuals (4.19M CLM, 4.13–4.18M MLM with padding) come from trainer_state.num_input_tokens_seen. Stage2 token counts are the 1B measurements (same data and total batch 512 × 8192).

JP stage2 (6 epochs) and the -bs16384 / -bs2048 legacy variants are untouched.

Notes

  • Optimizer is schedule-free RAdam with constant LR, so running stage1 jobs resume under the new budget without any schedule discontinuity.
  • Server-side job configs (cc25 mirei-1b3b-redo-20260725 use explicit max_steps 13,700 / 13,800; GCP mben3-s1 uses 0.65 epoch) are updated separately from this PR. ModernBERT-JP-3B stage1 already passed 13,800; its stage2 starts from checkpoint-13800.

Align the 3B English pre-training budget with the 1B design (stage1 +
stage2 ~= 20x non-embedding parameters, same epoch for Llama and
ModernBERT within a language):

- stage1 (FineWeb-Edu, seq 1024): 0.45 -> 0.65 epoch
  Llama-EN-3B 16,270 steps / 68.2B tokens, ModernBERT-EN-3B 13,892 steps / 58.3B tokens
- stage2 (wiki40b, seq 8192): 3 -> 1 epoch (same as 1B)
  ~2.5B tokens each

Totals: Llama-EN-3B 70.7B (25x of 2.83B), ModernBERT-EN-3B 60.8B (19.9x of 3.05B).
JP configs are unchanged (0.8 / 6 epochs).
@iamtatsuki05 iamtatsuki05 self-assigned this Sep 14, 2026
@iamtatsuki05
iamtatsuki05 merged commit d606dcd into develop Sep 14, 2026
2 checks passed
@iamtatsuki05 iamtatsuki05 changed the title Raise EN 3B stage1 to 0.65 epoch and cut EN 3B stage2 to 1 epoch Set 3B pre-training budget to 60B tokens (stage1 + stage2) per model Sep 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant