用途
Picotron 是一个极简且高度可修改的分布式训练框架，专为教育目的而设计。它受 NanoGPT 启发，旨在帮助用户快速理解并实践分布式训练中的各项技术，特别是用于预训练 Llama 类模型。
主要能力
Picotron 支持 4D 并行，即数据并行（Data Parallel）、张量并行（Tensor Parallel）、流水线并行（Pipeline Parallel）和上下文并行（Context Parallel）。其代码库非常简洁可读，核心文件 train.py、model.py 以及 data_parallel.py、tensor_parallel.py、pipeline_parallel.py、context_parallel.py 均少于 300 行代码。
性能方面，项目仍处于积极开发中，并非最优。根据官方信息，在 64 张 H100 GPU 上训练 LLaMA-2-7B 模型时观察到 38% 的 MFU（模型浮点运算利用率），在 8 张 H100 GPU 上训练 SmolLM-1.7B 模型时接近 50% MFU。更多基准测试即将发布。
与 Nanotron 相比，Picotron 主要面向教育目的，帮助人们快速熟悉分布式训练中的所有技术。
使用方式
用户需要先从 Hugging Face 获取 token（HF_TOKEN）以下载模型。
GPU 环境
- 使用 create_config.py 生成 JSON 格式的配置文件，默认输出到 tmp 目录。例如：
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 < HF_TOKEN >
- 本地运行：
torchrun --nproc_per_node 8 train.py --config tmp/llama-1B/config.json
- 3D 并行示例：
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 < HF_TOKEN >
- Slurm 提交：
python submit_slurm_jobs.py --inp_dir tmp/llama-7B --qos high --hf_token < HF_TOKEN >
CPU 环境（预期较慢）
3D 并行示例：
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 < HF_TOKEN > --use_cpu
本地运行：
torchrun --nproc_per_node 8 train.py --config tmp/llama-1B-cpu/config.json
适用场景
Picotron 适合以下场景：分布式训练的学习与教学、快速实验不同的并行策略、研究或验证 4D 并行技术。它不适用于追求最高训练性能的生产环境。
限制
- 性能并非最优，仍处于积极开发中。
- 主要面向教育目的，生产使用需谨慎评估。
- 具体安装步骤、完整依赖列表和详细基准测试暂未提供。
引用
如果使用 Picotron，请引用项目。教程视频和代码库可通过官方提供的链接获取。
