This model is a continual pre-training of Llama-3.1-8B on a mix of the Python subset of The-Stack-v2-train-smol-ids (from SwallowCode, Experiment 1) and multilingual text datasets.
The model was trained to evaluate the baseline performance of unfiltered Python code from The-Stack-v2 in the SwallowCode ablation experiments.
The model shows baseline performance on code generation tasks (HumanEval and HumanEval+) and maintains general capabilities on knowledge, reasoning, and common sense benchmarks.
It serves as a reference point for comparing subsequent ablation experiments in the SwallowCode pipeline.
It was trained on 50 billion tokens using a mix of 16% SwallowCode (Experiment 1, Python subset) and 84% multilingual text, following the setup described in the SwallowCode paper.
Training was performed using Megatron-LM.
Use
Generation
python
1# pip install -q transformers2from transformers import AutoModelForCausalLM, AutoTokenizer
34model ="tokyotech-llm/<model-name>"5device ="cuda"# for GPU usage or "cpu" for CPU usage67tokenizer = AutoTokenizer.from_pretrained(model)8model = AutoModelForCausalLM.from_pretrained(model).to(device)910inputs = tokenizer.encode("def fibonacci(n):", return_tensors="pt").to(device)11outputs = model.generate(inputs, max_length=100)12print(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:
16% Code: Python subset of The-Stack-v2-train-smol-ids (8B tokens), from SwallowCode, Experiment 1.
84% Multilingual Text:
Japanese Wikipedia (0.84B tokens)
Japanese Swallow Corpus v2 (26.1B tokens)
Laboro-ParaCorpus (0.22B tokens)
English Wikipedia (1.1B tokens)
English Cosmopedia (3.7B tokens)
English DCLM (10.0B 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 SwallowCode paper, with the lm-evaluation-harness and BigCodeBench. Benchmarks include code generation (HumanEval, HumanEval+) and general tasks (OpenBookQA, TriviaQA, HellaSwag, SQuAD 2.0, XWINO, MMLU, GSM8K, BBH). Results are reported for checkpoints at 10B, 20B, 30B, 40B, and 50B tokens.
Citation
bibtex
1@misc{fujii2025rewritingpretrainingdataboosts,
2 title={Rewriting Pre-Training Data Boosts LLM Performance in Math and Code},
3 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},
4 year={2025},
5 eprint={2505.02881},
6 archivePrefix={arXiv},
7 primaryClass={cs.LG},
8 url={https://arxiv.org/abs/2505.02881},
9}
