You can deploy this model and start earning money today!

tokyotech-llm/
Llama-3.1-8B-math-ablation-exp1-LR2.5e-5-WD0.1-iter0002500

Name: Llama-3.1-8B-math-ablation-exp1-LR2.5e-5-WD0.1-iter0002500
Availability: InStock

safetensors

llama

tokyotech-llm/swallow-math

2505.02881

meta-llama/Llama-3.1-8B

finetune

llama3.3

Views

No views yet

Model Card

Model Summary

This model is a continual pre-training of Llama-3.1-8B on a mix of mathematical datasets from finemath-4+ and multilingual text datasets. The model was trained to evaluate the performance of mathematical reasoning and problem-solving as part of the SwallowMath ablation experiments (experiment 1).

It was trained on 50 billion tokens using a mix of 4.8% Finemath-4+, 13.1% Code, and 82% multilingual text, following the setup described in the SwallowMath paper. Training was performed using Megatron-LM.

Use

Generation

python

# pip install -q transformers
from transformers import AutoModelForCausalLM, AutoTokenizer
model = "tokyotech-llm/<model-name>"
device = "cuda"  # for GPU usage or "cpu" for CPU usage
tokenizer = AutoTokenizer.from_pretrained(model)
model = AutoModelForCausalLM.from_pretrained(model).to(device)
inputs = tokenizer.encode("Solve the equation 2x + 3 = 7:", return_tensors="pt").to(device)
outputs = model.generate(inputs, max_length=100)
print(tokenizer.decode(outputs[0]))

Training

Model

Architecture: Llama-3.1
Pretraining tokens: 50B
Precision: bfloat16
Sequence length: 8,192
Tokenizer: Llama-3 tokenizer

Data

The training mix consists of:

Mathematical Data (~4.84%):
- Target Math Data: 2.4B tokens
Code Data (~13.12%):
- SwallowCode (Syntax, Pylint Filtered): 6.5B tokens
Multilingual Text (~82.04%):
- Japanese Wikipedia: 0.84B tokens
- Japanese Swallow Corpus v2: 33.0B tokens
- Laboro-ParaCorpus: 0.22B tokens
- English Wikipedia: 1.1B tokens
- English Cosmopedia: 3.3B tokens
- English DCLM: 2.2B tokens

Details are in the paper’s Appendix.

Hardware

GPUs: 64 NVIDIA H100 (94GB)
Interconnect: InfiniBand NDR200
Supercomputer: TSUBAME, Institute of Science Tokyo

Software

Megatron-LM (version core_r0.9.0) for training
lm-evaluation-harness for evaluation
BigCodeBench for code evaluation

Evaluation

The model was evaluated using the setup described in the SwallowMath paper, with the lm-evaluation-harness and BigCodeBench. Benchmarks include mathematical reasoning (GSM8K, MATH), code generation (HumanEval), and general tasks (OpenBookQA, TriviaQA, HellaSwag, SQuAD 2.0, XWINO, MMLU, BBH). Results are reported for checkpoints at 10B, 20B, 30B, 40B, and 50B tokens.

Evaluation Results (Finemath-4+ experiment 1)

Tokens (B)	OpenBookQA	TriviaQA	HellaSwag	SQuAD2.0	XWINO	MMLU	HumanEval	GSM8K	BBH	MATH
10	0.3700	0.6626	0.5990	0.3350	0.8985	0.6243	0.3439	0.4685	0.6057	0.1760
20	0.3720	0.6536	0.5963	0.3510	0.9032	0.6261	0.3622	0.5011	0.5896	0.2080
30	0.3700	0.6574	0.5999	0.3506	0.8998	0.6253	0.3561	0.5019	0.5971	0.2260
40	0.3720	0.6577	0.6024	0.3499	0.9049	0.6312	0.3701	0.5231	0.6054	0.2260
50	0.3740	0.6608	0.6001	0.3550	0.9058	0.6329	0.3561	0.5292	0.6166	0.2400

Citation

bibtex

@misc{fujii2025rewritingpretrainingdataboosts,
      title={Rewriting Pre-Training Data Boosts LLM Performance in Math and Code}, 
      author={Kazuki Fujii and Yukito Tajima and Sakae Mizuki and Hinari Shimada and Taihei Shiotani and Koshiro Saito and Masanari Ohi and Masaki Kawamura and Taishi Nakamura and Takumi Okamoto and Shigeki Ishida and Kakeru Hattori and Youmi Ma and Hiroya Takamura and Rio Yokota and Naoaki Okazaki},
      year={2025},
      eprint={2505.02881},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2505.02881}, 
}

tokyotech-llm/Llama-3.1-8B-math-ablation-exp1-LR2.5e-5-WD0.1-iter0002500