Picotron实战案例：在8个H100 GPU上训练SmolLM-1.7B模型的完整指南

Picotron实战案例在8个H100 GPU上训练SmolLM-1.7B模型的完整指南【免费下载链接】picotronMinimalistic 4D-parallelism distributed training framework for education purpose项目地址: https://gitcode.com/gh_mirrors/pi/picotronPicotron是一个极简主义的4D并行分布式训练框架专为教育目的设计。本文将详细介绍如何使用Picotron框架在8个H100 GPU上高效训练SmolLM-1.7B模型实现近50%的MFU模型浮点运算利用率。这个开源项目让分布式深度学习训练变得更加简单易学 什么是Picotron框架Picotron是一个极简且高度可定制的分布式训练框架灵感来自NanoGPT专门用于预训练类似Llama的模型。它支持完整的4D并行训练策略数据并行- 将训练数据分配到不同GPU张量并行- 将模型权重拆分到多个GPU流水线并行- 将模型层分配到不同GPU上下文并行- 处理长序列的注意力计算与复杂的工业级框架不同Picotron的代码简洁明了train.py、model.py和各个并行模块都控制在300行代码以内非常适合学习和实验 环境准备与安装系统要求8个H100 GPU或类似性能的GPUPython 3.8PyTorch 2.0CUDA 11.8快速安装# 克隆仓库 git clone https://gitcode.com/gh_mirrors/pi/picotron cd picotron # 安装依赖 pip install -e . 配置8个H100 GPU训练创建训练配置文件使用create_config.py脚本快速生成训练配置python create_config.py \ --out_dir tmp \ --exp_name smollm-1.7b-8h100 \ --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关键参数说明--dp 8使用8路数据并行--model_name HuggingFaceTB/SmolLM-1.7B指定SmolLM-1.7B模型--grad_acc_steps 32梯度累积步数--mbs 4每个GPU的微批次大小--seq_len 1024序列长度配置文件详解生成的配置文件位于tmp/smollm-1.7b-8h100/config.json包含以下核心设置{ distributed: { tp_size: 1, // 张量并行 cp_size: 1, // 上下文并行 pp_size: 1, // 流水线并行 dp_size: 8, // 数据并行 pp_engine: 1f1b }, model: { name: HuggingFaceTB/SmolLM-1.7B, num_hidden_layers: 15 }, training: { micro_batch_size: 4, gradient_accumulation_steps: 32, seq_length: 1024 } }‍♂️ 启动分布式训练本地训练启动torchrun --nproc_per_node 8 train.py --config tmp/smollm-1.7b-8h100/config.jsonSlurm集群启动python submit_slurm_jobs.py --inp_dir tmp/smollm-1.7b-8h100 --qos high --hf_token HF_TOKEN 性能监控与优化实时训练指标训练过程中Picotron会显示详细的性能指标Step: 100 | Loss: 2.3456 | Global batch size: 1.0M | Tokens/s: 125.4K | Tokens/s/GPU: 15.7K | MFU: 48.7% | Memory usage: 45.2GB关键指标解释MFU (Model FLOPs Utilization)模型浮点运算利用率达到48.7%表示接近理论峰值性能的50%Tokens/s/GPU每个GPU每秒处理的token数Memory usageGPU内存使用情况使用WandB监控在配置文件中启用WandB日志记录{ logging: { use_wandb: true, project_name: picotron, run_name: smollm-1.7b-8h100 } } 4D并行策略组合高级配置示例如果你想尝试更复杂的并行策略组合Picotron支持灵活的4D并行配置# 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并行策略选择指南并行类型适用场景优势配置建议数据并行模型能放入单个GPU实现简单dpGPU数量张量并行单个GPU内存不足减少通信开销tp2或4流水线并行模型层数很多处理超大模型pp模型层数/每GPU层数上下文并行长序列训练高效处理长文本cp序列长度/分块大小 实战技巧与最佳实践1. 内存优化策略使用梯度检查点减少内存占用启用Flash Attention加速注意力计算合理设置micro_batch_size避免OOM2. 性能调优建议从纯数据并行开始逐步增加其他并行维度监控MFU指标确保硬件利用率使用--use_fused_adam启用融合Adam优化器3. 故障排查检查GPU内存使用nvidia-smi验证数据加载速度调整num_workers参数确认并行配置正确确保tp_size * pp_size * dp_size * cp_size GPU总数 训练结果与验证预期训练性能在8个H100 GPU上训练SmolLM-1.7B模型你可以期待MFU: 接近50%的模型浮点运算利用率吞吐量: 每个GPU约15-20K tokens/秒总吞吐量: 120-160K tokens/秒内存使用: 每个GPU约40-50GB模型验证训练完成后使用检查点恢复训练或进行推理from picotron.checkpoint import CheckpointManager checkpoint_manager CheckpointManager() step, trained_tokens checkpoint_manager.load_checkpoint( model, optimizer, checkpoints/latest ) 核心模块解析Picotron的简洁设计使其易于理解和修改训练主循环 train.py完整的4D并行训练流程支持多种并行策略组合内置性能监控和日志记录模型定义 picotron/model.pyLlama架构的PyTorch实现支持Flash Attention优化简洁的Decoder层设计并行模块张量并行picotron/tensor_parallel/tensor_parallel.py流水线并行picotron/pipeline_parallel/pipeline_parallel.py上下文并行picotron/context_parallel/context_parallel.py数据并行picotron/data_parallel/data_parallel.py 常见问题解答Q: 如何调整训练批次大小A: 通过修改--mbs微批次大小和--grad_acc_steps梯度累积步数参数全局批次大小 dp_size × mbs × grad_acc_steps。Q: 支持哪些模型架构A: 目前主要支持Llama架构但代码结构清晰易于扩展支持其他架构。Q: 如何在CPU上测试A: 添加--use_cpu参数即可在CPU上运行适合学习和调试。Q: 如何继续中断的训练A: 在配置文件中设置load_path: checkpoints/latest即可从检查点恢复。 总结Picotron作为教育导向的分布式训练框架成功在8个H100 GPU上实现了SmolLM-1.7B模型的高效训练。通过简单的配置和清晰的代码结构即使是深度学习新手也能快速上手大规模模型训练。关键收获Picotron的4D并行策略让大规模模型训练变得可行简洁的代码设计降低了学习门槛近50%的MFU证明了框架的高效性灵活的配置支持多种训练场景现在就开始你的分布式训练之旅吧使用Picotron你可以在几行命令内启动8个H100 GPU上的SmolLM-1.7B训练体验工业级的大模型训练流程。提示更多详细信息和最新更新请参考项目文档和示例配置。【免费下载链接】picotronMinimalistic 4D-parallelism distributed training framework for education purpose项目地址: https://gitcode.com/gh_mirrors/pi/picotron创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考