https://github.com/huggingface/picotron
Minimalistic 4D-parallelism distributed training framework for education purpose
https://github.com/huggingface/picotron
Last synced: 8 months ago
JSON representation
Minimalistic 4D-parallelism distributed training framework for education purpose
- Host: GitHub
- URL: https://github.com/huggingface/picotron
- Owner: huggingface
- License: apache-2.0
- Created: 2024-09-18T12:01:22.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-08-26T13:47:24.000Z (9 months ago)
- Last Synced: 2025-09-30T18:02:30.507Z (8 months ago)
- Language: Python
- Homepage:
- Size: 3.08 MB
- Stars: 1,834
- Watchers: 14
- Forks: 136
- Open Issues: 11
-
Metadata Files:
- Readme: README.md
- License: LICENSE
- Citation: CITATION.cff
Awesome Lists containing this project
README
# picotron
In the spirit of [NanoGPT](https://github.com/karpathy/nanoGPT), we created Picotron: The minimalist & most-hackable repository for pre-training Llama-like models with [4D Parallelism](https://arxiv.org/abs/2407.21783) (Data, Tensor, Pipeline, Context parallel). It is designed with simplicity and **educational** purposes in mind, making it an excellent tool for learning and experimentation.
- The code itself is simple and readable: `train.py`, `model.py` and `[data|tensor|pipeline|context]_parallel.py` are all under **300** lines of code.
- Performance is not the best but still under active development. We observed 38% MFU on a LLaMA-2-7B model using 64 H100 GPUs and nearly 50% MFU on the SmolLM-1.7B model with 8 H100 GPUs. Benchmarks will come soon
- Compared to [Nanotron](https://github.com/huggingface/nanotron/tree/main), Picotron is primarily for educational purposes, helping people quickly get familiar with all the techniques in distributed training
# Tutorial videos
- A step by step tutorial on how to build Picotron distributed training framework form scratch:
- [Picotron tutorial (playlist)](https://www.youtube.com/playlist?list=PL-_armZiJvAnhcRr6yTJ0__f3Oi-LLi9S) 🎬
- [Picotron tutorial (codebase)](https://github.com/huggingface/picotron_tutorial) 👷🏻♂️
# Install
```
pip install -e .
```
# Quick start
- Get a HF token [here](https://huggingface.co/settings/tokens) to download models from HuggingFace
- GPU
```sh
# To create a config file in json format under tmp by default
python create_config.py --out_dir tmp --exp_name llama-1B --dp 8 --model_name HuggingFaceTB/SmolLM-1.7B --num_hidden_layers 15 --grad_acc_steps 32 --mbs 4 --seq_len 1024 --hf_token
# Locally
torchrun --nproc_per_node 8 train.py --config tmp/llama-1B/config.json
# 3D Parallelism
python create_config.py --out_dir tmp --dp 4 --tp 2 --pp 2 --pp_engine 1f1b --exp_name llama-7B --model_name meta-llama/Llama-2-7b-hf --grad_acc_steps 32 --mbs 4 --seq_len 1024 --hf_token
# Slurm
python submit_slurm_jobs.py --inp_dir tmp/llama-7B --qos high --hf_token
```
- CPU (expect it to be slow)
```sh
# 3D Parallelism on CPU
python create_config.py --out_dir tmp --exp_name llama-1B-cpu --dp 2 --tp 2 --pp 2 --pp_engine 1f1b --model_name HuggingFaceTB/SmolLM-1.7B --num_hidden_layers 5 --grad_acc_steps 2 --mbs 4 --seq_len 128 --hf_token --use_cpu
# Locally
torchrun --nproc_per_node 8 train.py --config tmp/llama-1B-cpu/config.json
```
# Citation
If you use Picotron, please cite it as:
```bibtex
@misc{zhao2025picotron,
author = {Haojun Zhao and Ferdinand Mom},
title = {Picotron: Distributed training framework for education and research experimentation},
year = {2025},
publisher = {GitHub},
journal = {GitHub repository},
howpublished = {\url{https://github.com/huggingface/picotron}}
}
```
# Acknowledgements
- [Megatron-LM](https://github.com/NVIDIA/Megatron-LM)
- [FairScale](https://github.com/facebookresearch/fairscale)
- [LitGPT](https://github.com/Lightning-AI/lit-gpt)