Markdown Editor

Model Training Plan Template

Objective and success criteria, data mix and splits, baseline, architecture, 6ND compute budget in KaTeX, ablations, checkpoints, stop rules, gates, risks

AI & DataFullPlanning concepts from Kaplan et al. 2020 and Hoffmann et al. 2022modeltrainingplanmath

How to use: Use it when you are about to reserve compute for a training or large fine-tuning run and need sign-off. Fill the success criteria, the data splits and the compute budget first, then the ablations that de-risk the main run. The main pitfall is treating scaling-law estimates as exact instead of replacing them with pilot measurements.

Preview

Model Training Plan: [Model name]

Plan for a model training run, written and approved before the main run starts: the objective and success criteria, the data and its splits, the architecture and the baseline it must beat, the compute budget, the ablations that de-risk the run, the checkpoint and stop rules, the evaluation gates, and the risks. The compute estimate uses the widely cited approximation of about 6 FLOPs per parameter per training token (Kaplan et al. 2020) and the compute-optimal sizing idea of Hoffmann et al. 2022 as planning concepts, not as guarantees. Replace every [bracketed] value and every row marked Example. Worked examples use a 1-billion-parameter decoder-only language model trained from scratch for code completion.

Document control

FieldValue
Document ID[TRN-PLAN-001]
Model name and target version[Model name] [v1.0]
Run family[Example: code-lm-1b-v1]
Version[0.1]
Status[Draft / In review / Approved / Running / Closed]
Owner[@training-lead]
Approvers / reviewers[@research-lead], [@infra-lead], [@responsible-ai], [@finance-partner]
Last updatedYYYY-MM-DD
Classification[Internal / Confidential]
Code repository and commit[repo URL] @ [commit SHA]
Compute reservation[Cluster name, accelerator type, count, dates]
Related documents[ML system design], [Datasheet], [Evaluation plan], [Model card]

Approvals

RoleNameDecisionDate
Training lead[@name][Approved / Changes requested]YYYY-MM-DD
Research lead[@name][Decision]YYYY-MM-DD
Infrastructure lead[@name][Decision]YYYY-MM-DD
Responsible AI reviewer[@name][Decision]YYYY-MM-DD
Budget holder[@name][Decision]YYYY-MM-DD

Revision history

VersionDateAuthorChange summary
0.1YYYY-MM-DD[@name]First draft
0.2YYYY-MM-DD[@name]Example: Token budget raised from 15B to 20B after ablation ABL-003
1.0YYYY-MM-DD[@name]Example: Approved for the main run

Contents

1. Objective and success criteria

1.1 Objective

[One paragraph: what the model is for, who uses it, and why a new training run is needed rather than an existing model, prompt changes or fine-tuning.]

Example: Train a 1B-parameter code completion model that runs on a single inference accelerator inside the latency budget of the editor plug-in, and that beats the current 350M-parameter production model on completion acceptance without regressing on licence and secret leakage checks.

1.2 Success criteria

IDCriterionMetricThresholdMeasured on
SC-001 (Example)Better completions than productionpass@1 on internal held-out suite≥ production + 4.0 pointsEVAL-SET-01
SC-002 (Example)Fits the latency budgetp95 time to first token at batch 1≤ 120 msServing benchmark
SC-003 (Example)No memorised secretsSecret extraction rate on canary prompts0 canaries recoveredCanary suite
SC-004 (Example)Stays within budgetAccelerator hours used≤ 110% of section 5 totalCluster billing
[SC-nnn][Criterion][Metric][Threshold][Set]

1.3 Non-goals

  • Example: Natural-language chat quality; this model completes code only.
  • Example: Languages outside the top 12 in the data mix.
  • [Non-goal]

2. Data and splits

Document each dataset in a Datasheet (ai-data/datasheet-for-datasets) and reference it here. Record the tokenizer version, because token counts change when the tokenizer changes.

2.1 Sources and mixture

IDSourceLicence / basisTokens (after filtering)Mixture weightEpochs
DS-001 (Example)Permissively licensed public repositoriesPer-file licence filter14.0B70%1.0
DS-002 (Example)Internal repositories with consentInternal agreement3.0B15%1.0
DS-003 (Example)Technical documentationPer-site terms2.0B10%1.0
DS-004 (Example)Synthetic unit tests generated by teacher modelTeacher terms of use1.0B5%1.0
Total20.0B100%

2.2 Processing pipeline

StepRuleEvidence
Licence filterExample: Keep files whose detected licence is on the allow listFilter log with counts
Exact and near-duplicate removalExample: MinHash near-duplicate removal at file levelDedup report
Secret and personal data scrubbingExample: Pattern and entropy scanners, replaced with placeholdersScanner report
Quality filterExample: Drop auto-generated and minified filesFilter thresholds and samples
DecontaminationRemove training documents that overlap evaluation sets (section 2.3)Overlap report
Tokenization[Tokenizer name and version]Token count per source

2.3 Splits and decontamination

SplitPurposeSizeFrozen onContent hash
TrainGradient updatesExample: 20.0B tokensYYYY-MM-DD[sha256]
ValidationLoss curves, early warningExample: 50M tokens, same mixtureYYYY-MM-DD[sha256]
Test (held out)Gate evaluation onlyExample: 2,000 tasks, EVAL-SET-01YYYY-MM-DD[sha256]
CanaryMemorisation checksExample: 500 unique strings inserted on purposeYYYY-MM-DD[sha256]

Check decontamination by n-gram overlap between every evaluation set and the training corpus, and record the overlap rate per set. Split by repository, not by file, so that near-identical files from one project cannot sit on both sides of the split.

3. Baseline

IDBaselineWhy it is the right comparisonScore on EVAL-SET-01
BL-001 (Example)Production 350M modelThe model users have todaypass@1 31.5
BL-002 (Example)Public open-weight 1B code modelShows whether training our own model is worth itpass@1 34.0
BL-003 (Example)Same architecture, 5B tokensShows the value of the extra dataProduced in ablation ABL-003

If a public model of similar size already meets the success criteria, stop and record the decision instead of training.

4. Architecture and training configuration

4.1 Architecture

SettingValueReason
Model familyExample: Decoder-only transformerStandard for left-to-right completion
Parameters (N)Example: 1.0BFits one inference accelerator at the latency target
Layers / hidden size / headsExample: 20 / 2,048 / 16About 12 x 20 x 2,048² ≈ 1.0B non-embedding parameters, the N used in section 5
Context lengthExample: 2,048 tokensCovers 95% of completion contexts in usage logs
Positional encodingExample: Rotary[Reason]
VocabularyExample: 49,152 tokens, byte-level BPETrained on DS-001 sample
PrecisionExample: bfloat16 compute, float32 master weights[Reason]

4.2 Optimisation

SettingValue
OptimiserExample: AdamW, beta1 0.9, beta2 0.95, weight decay 0.1
Peak learning rate and scheduleExample: 3e-4, linear warm-up 2,000 steps, cosine decay to 10%
Global batchExample: 512 sequences x 2,048 tokens = 1,048,576 tokens per step
Total steps (S)Example: 20.0B / 1,048,576 ≈ 19,073 steps
Gradient clippingExample: Global norm 1.0
ParallelismExample: Data parallel with sharded optimiser state across 16 accelerators
Random seedsExample: Model init 1234, data order 5678, recorded in the run config

5. Compute budget

5.1 Training FLOPs

For a dense transformer, the training compute is commonly approximated as about 6 floating-point operations per parameter per training token: roughly 2 for the forward pass and 4 for the backward pass. Kaplan et al. (2020) use this approximation and it ignores the attention term that depends on context length, so add a margin for long contexts.

Ctrain≈6⋅N⋅DC_{\text{train}} \approx 6 \cdot N \cdot D

where NN is the number of non-embedding parameters and DD is the number of training tokens.

Example: N=1.0×109N = 1.0 \times 10^{9} and D=2.0×1010D = 2.0 \times 10^{10}, so Ctrain≈6×109×2×1010=1.2×1020C_{\text{train}} \approx 6 \times 10^{9} \times 2 \times 10^{10} = 1.2 \times 10^{20} FLOPs.

5.2 Model and data size check

Hoffmann et al. (2022) found that, for a fixed compute budget, the number of parameters and the number of training tokens should grow in roughly equal proportion, and that many large models of the time were trained on too few tokens. A common reading of their results is a ratio of the order of 20 training tokens per parameter for compute-optimal pre-training. Treat this as a sanity check, not a rule: a model meant for cheap inference is often trained on far more tokens than the compute-optimal amount, because inference cost depends on NN, not on DD.

Example: D/N=2.0×1010/1.0×109=20D / N = 2.0 \times 10^{10} / 1.0 \times 10^{9} = 20 tokens per parameter, close to the compute-optimal ratio. The team may choose to extend DD later if ablation ABL-003 shows that loss is still falling steeply.

5.3 Accelerator hours and cost

Tacc=CtrainPpeak⋅MFUT_{\text{acc}} = \frac{C_{\text{train}}}{P_{\text{peak}} \cdot \text{MFU}}

where PpeakP_{\text{peak}} is the accelerator's peak throughput at the training precision and MFU is the model FLOPs utilisation measured in the pilot run.

ItemValueBasis
Peak throughput per acceleratorExample: 400 TFLOP/s (4.0 x 10^14)Vendor datasheet for the chosen precision
Assumed MFUExample: 40%Pilot run PIL-001; replace with the measured figure
Effective throughputExample: 1.6 x 10^14 FLOP/s4.0 x 10^14 x 0.40
Main run accelerator-secondsExample: 750,000 s1.2 x 10^20 / 1.6 x 10^14
Main run accelerator-hoursExample: 208.3 h750,000 / 3,600
Wall-clock time on 16 acceleratorsExample: 13.0 h208.3 / 16, before restarts
Budget lineAccelerator-hoursCost at USD 2.50 per hour (Example)
Main run (Example)208.33USD 520.83
Ablations, 10 runs at 1% of main compute each (Example)20.83USD 52.08
Evaluation of checkpoints (Example)10.00USD 25.00
Subtotal239.17USD 597.92
Contingency 30% for restarts and reruns (Example)71.75USD 179.38
Total requested310.92USD 777.29

Each ablation uses a 100M-parameter model on 2B tokens, so it costs 6×108×2×109=1.2×10186 \times 10^{8} \times 2 \times 10^{9} = 1.2 \times 10^{18} FLOPs, which is 1% of the main run.

5.4 Memory check

A common rule of thumb for mixed-precision training with Adam is about 16 bytes of accelerator memory per parameter for weights, gradients and optimiser state, before activations. Example: 1.0B parameters x 16 bytes = 16 GB of state, sharded across 16 accelerators, which leaves room for activations at the chosen micro-batch. Confirm the peak memory in the pilot run.

6. Ablations and pilot runs

IDQuestionDesignDecision ruleResult
PIL-001 (Example)Does the full stack run, and what is the real MFU?Main config, 200 stepsProceed if MFU ≥ 30% and no errors[Result]
ABL-001 (Example)Does adding DS-004 synthetic tests help?100M model, 2B tokens, with and without DS-004, 2 seeds eachKeep if pass@1 gain exceeds seed spread[Result]
ABL-002 (Example)Which peak learning rate is stable?100M model, 3 learning ratesPick the highest rate with no loss spikes[Result]
ABL-003 (Example)Is loss still falling at 20B tokens?Fit loss versus tokens on 3 small sizesExtend D if the fitted slope predicts ≥ 1% further gain[Result]
ABL-004 (Example)Does the tokenizer handle all 12 languages?Tokens per character per languageRetrain tokenizer if any language is 30% worse than the median[Result]
[ABL-nnn][Question][Design][Decision rule][Result]

Run ablations with at least two seeds when the expected effect is small, and report the spread. Record each run in an Experiment Tracking Report (ai-data/experiment-tracking-report).

7. Checkpoints and monitoring

7.1 Checkpoint policy

SettingValue
Save intervalExample: Every 1,000 steps (about 41 minutes of wall-clock time)
ContentsModel weights, optimiser state, scheduler state, data loader position, random number generator states
SizeExample: About 16 GB per full checkpoint
RetentionExample: Last 3 rolling, plus every 5,000 steps permanently, plus final
Storage and integrity[Bucket path]; checksum written with each checkpoint
Restore testRestore from a checkpoint in the pilot run and confirm the loss curve continues

7.2 Signals and alerts

SignalNormal rangeAlert whenWho is paged
Training lossExample: Smooth decreaseSpike above 1.5x the 200-step moving average[@on-call]
Validation lossExample: Tracks training lossRises for 3 consecutive evaluations[@training-lead]
Gradient normExample: Below clip value most stepsClipped on more than 50% of steps for 500 steps[@on-call]
ThroughputExample: Within 10% of pilotDrops more than 20%[@infra]
Hardware healthNo errorsAny device error or node loss[@infra]

8. Stop and restart criteria

IDConditionActionDecided by
STOP-001 (Example)Loss spike that does not recover within 500 stepsRoll back to last good checkpoint, skip the data batch, lower learning rate 10%Training lead
STOP-002 (Example)Second unrecovered spikePause run, open an incident, review configTraining and research leads
STOP-003 (Example)Validation loss at 50% of tokens worse than ABL-003 forecast by more than 5%Pause and review before spending the second halfResearch lead
STOP-004 (Example)Spend reaches 100% of the total requested budgetStop and request approval to continueBudget holder
STOP-005 (Example)Canary memorisation detected in an intermediate checkpointStop, investigate data pipelineResponsible AI reviewer
DONE-001 (Example)All D tokens consumed and final checkpoint savedRun final evaluation gates (section 9)Training lead

9. Evaluation gates

GateWhenChecksPass ruleOwner
G-0 (Example)Before main runPilot and ablations done, budget approvedAll decisions recordedTraining lead
G-1 (Example)25% of tokensValidation loss vs forecast; quick pass@1 on 200 tasksWithin 5% of forecastResearch lead
G-2 (Example)50% of tokensFull EVAL-SET-01, canary suiteNo canary leak; pass@1 trending to SC-001Research lead
G-3 (Example)FinalSC-001 to SC-004, safety and licence checks, serving benchmarkEvery criterion metApprovers in Document control

Use the method in the LLM Evaluation Plan (ai-data/llm-evaluation-plan) or the Benchmark Report (ai-data/benchmark-report) for confidence intervals and significance. Release documentation goes into a Model Card (ai-data/model-card).

10. Risks and mitigations

IDRiskLikelihoodImpactMitigationTrigger
RSK-001 (Example)Training instability at the chosen learning rateMediumHighABL-002; rollback rule STOP-001Loss spike alert
RSK-002 (Example)Evaluation contamination inflates scoresMediumHighDecontamination in section 2.3; fresh tasks added after the data cut-offOverlap rate above 0.5%
RSK-003 (Example)Licence-incompatible code in training dataLowHighLicence filter; legal review of allow listAudit sample finds a violation
RSK-004 (Example)Lower MFU than planned extends the runMediumMediumMeasure in PIL-001; contingency lineThroughput alert
RSK-005 (Example)Hardware failure loses progressMediumLowCheckpoint every 1,000 steps; restore testedNode loss
RSK-006 (Example)Model reproduces secrets or personal dataLowHighScrubbing, canaries, G-2 checkCanary recovered
[RSK-nnn][Risk][L][I][Mitigation][Trigger]

11. Schedule and ownership

PhaseStartEndOwnerExit condition
Data preparation and freezeYYYY-MM-DDYYYY-MM-DD[@data-lead]Split hashes recorded
Pilot and ablationsYYYY-MM-DDYYYY-MM-DD[@training-lead]G-0 passed
Main runYYYY-MM-DDYYYY-MM-DD[@training-lead]DONE-001
Final evaluationYYYY-MM-DDYYYY-MM-DD[@eval-owner]G-3 decision recorded
Documentation and handoverYYYY-MM-DDYYYY-MM-DD[@training-lead]Model card and run report published

After the run, record the actual compute, MFU, number of restarts and final scores next to the planned values. The difference between plan and actual is the most useful input to the next plan.

ItemPlannedActualVariance
Accelerator-hoursExample: 310.92[Actual][%]
MFUExample: 40%[Actual][points]
RestartsExample: 0 to 2[Actual][Count]
pass@1 on EVAL-SET-01Example: ≥ 35.5[Actual][points]

12. Glossary and references

TermMeaning
FLOPsFloating-point operations; total work, not a rate
MFUModel FLOPs utilisation: useful model FLOPs per second divided by the hardware peak
AblationA smaller controlled run that changes one factor to measure its effect
DecontaminationRemoving training data that overlaps evaluation data
CanaryA unique string inserted on purpose to detect memorisation
Compute-optimalThe split between model size and tokens that gives the lowest loss for a fixed compute budget
pass@1Share of tasks solved by the first generated sample

References:

  • Kaplan, J. et al. (2020). Scaling Laws for Neural Language Models. arXiv:2001.08361. Source of the approximation of about 6N training FLOPs per token (section 2.1 of the paper).
  • Hoffmann, J. et al. (2022). Training Compute-Optimal Large Language Models. arXiv:2203.15556. Source of the finding that parameters and tokens should scale in roughly equal proportion.
  • Pineau, J. et al. (2021). Improving Reproducibility in Machine Learning Research. arXiv:2003.12206 (JMLR).
  • Related templates: ML System Design (ai-data/ml-system-design), Fine-Tuning Report (ai-data/fine-tuning-report), Datasheet for Datasets (ai-data/datasheet-for-datasets), Experiment Tracking Report (ai-data/experiment-tracking-report), Model Card (ai-data/model-card), Reproducibility Checklist (science/reproducibility-checklist).

Standard and licence note

Based on common machine learning engineering practice and on the published research cited above (Kaplan et al. 2020; Hoffmann et al. 2022), used here as planning concepts only. No formal standard governs training plans. The text, tables and worked example are original; example figures are illustrative and must be replaced with measured values. Check the current versions of the cited papers and your hardware vendor's datasheet before relying on any number. Template licence: MIT.

Related templates