从零预训练 TinyLlama 1.1B基于 LitGPT 的数据处理、FSDP 训练与断点续训完整指南【免费下载链接】litgpt20 high-performance LLMs with recipes to pretrain, finetune and deploy at scale.项目地址: https://gitcode.com/GitHub_Trending/li/litgpt本指南以 tutorials/pretrain_tinyllama.md 为主线完整演示如何在 LitGPT 仓库中复现 TinyLlama 1.1B 的预训练流程从 SlimPajama 与 Starcoder 数据集的下载、tokenize 与二进制分块预处理到使用litgpt pretrain启动多卡训练、断点恢复与 checkpoint 导出。读完本文你将掌握一条可直接落地运行的预训练流水线并能读懂 config_hub/pretrain/tinyllama.yaml 与 litgpt/pretrain.py 中对应的实现细节。TinyLlama 是什么TinyLlama 与 Meta AI 的 LLaMA 2 在架构上完全一致但参数量只有 1.1B并在多个 epoch 上使用 SlimPajama 与 Starcoder 混合数据集训练。它最初由 jzhang38/TinyLlama 项目提出目标是让单卡也能研究的 LLaMA 架构训练。以下事实表来源于 TinyLlama 作者的 README名称说明参数量1.1B模型规模Layers: 22Heads: 32Query Groups: 4Embedding Size: 2048Intermediate Size: 5632序列长度2048学习率4e-4学习率调度余弦退火2000 步 warmup训练数据SlimPajama893 GB、Starcoder290 GB组合数据集规模约 950B tokens训练期间总 token 数3 万亿3 个 epoch完成训练耗时64 张 A100 约 4 周MFUModel FLOPs Utilization52%在 LitGPT 仓库中TinyLlama 的架构定义位于 litgpt/config.py其hf_config直接指向TinyLlama/TinyLlama-1.1B关键超参数为block_size2048、vocab_size32000、n_layer22、n_head32、n_embd2048、n_query_groups4、intermediate_size5632使用RMSNorm归一化与LLaMAMLP前馈层、biasFalse——这些与事实表一一对应也说明config_hub中的配置与官方 TinyLlama 保持一致。下载数据集预训练需要两个数据集SlimPajama约 893 GB与 Starcoder约 290 GB两者合计约 1.2 TB 磁盘空间。使用 git-lfs 下载# 确保已安装 git-lfshttps://git-lfs.com sudo apt install git-lfs git clone https://huggingface.co/datasets/cerebras/slimpajama-627b data/slimpajama-raw git clone https://huggingface.co/datasets/bigcode/starcoderdata data/starcoderdata-raw将数据预处理为训练格式LitGPT 的预训练脚本不能直接读取原始文本/parquet 文件而是要求数据被读取、tokenize 并写入二进制分块binary chunks。这一步依赖litdata的优化流水线与 streaming 数据集。首先安装预处理所需的附加依赖pip install .[all]然后准备 tokenizer 配置。本教程假设使用与 LLaMA/TinyLlama 相同的 tokenizer因此先下载 LLaMA-2-7B 的 tokenizer任何经过训练的、词表大小为 32000 的 SentencePiece tokenizer 也可以litgpt download meta-llama/Llama-2-7b-hf \ --access_token your_hf_token \ --tokenizer_only true下载完成后分别对每个数据集的每个 split 运行预处理脚本。Starcoder需要约 1.1 TB 磁盘空间python litgpt/data/prepare_starcoder.py \ --input_dir data/starcoderdata-raw \ --output_dir data/starcoder \ --tokenizer_path checkpoints/meta-llama/Llama-2-7b-hfSlimPajama需要约 2.5 TB 磁盘空间其 validation/test/train 三个 split 需要分别处理python litgpt/data/prepare_slimpajama.py \ --input_dir data/slimpajama-raw/validation \ --output_dir data/slimpajama/val \ --tokenizer_path checkpoints/meta-llama/Llama-2-7b-hf python litgpt/data/prepare_slimpajama.py \ --input_dir data/slimpajama-raw/test \ --output_dir data/slimpajama/test \ --tokenizer_path checkpoints/meta-llama/Llama-2-7b-hf python litgpt/data/prepare_slimpajama.py \ --input_dir data/slimpajama-raw/train \ --output_dir data/slimpajama/train \ --tokenizer_path checkpoints/meta-llama/Llama-2-7b-hf如果想先在数据的小切片上试运行为上述命令追加--fast_dev_runtrue即可。从源码看这两个脚本都构建在litdata的DataProcessorDataChunkRecipe之上见 litgpt/data/prepare_starcoder.py 与 litgpt/data/prepare_slimpajama.pyStarcoderDataRecipe用Path.rglob(*.parquet)递归收集文件并通过pyarrow.parquet.ParquetFile.iter_batches(batch_size8192, columns[content])分批量读取 content 列以控制 RAM 占用然后逐条tokenizer.encode(text, bosFalse, eosTrue)SlimPajamaDataRecipe用Path.rglob(*.zst)收集 zstd 压缩的 JSONL 文件并显式跳过RedPajamaGithub来源的数据——因为该部分与 Starcoder 数据重叠避免重复见 litgpt/data/prepare_slimpajama.py默认chunk_size分别为2049 * 8192Starcoder与2049 * 16384SlimPajama2049对应2048 token 序列 1 个下一 token两个脚本都通过CLI(prepare)暴露命令行参数litgpt/utils.py 中的CLInum_workers默认取os.cpu_count()。启动预训练预处理完成后即可启动预训练。脚本的默认配置至少需要 8 张 A100 GPUlitgpt pretrain --config config_hub/pretrain/tinyllama.yaml提示可以使用litgpt pretrain --data.help TinyLlama列出额外的数据集选项。脚本会周期性地将 checkpoint 保存到out/目录。默认情况下pretrain脚本使用FSDP bfloat16 混合精度 梯度累积进行训练。需要特别说明的是pretrain并不是针对单一模型的训练脚本你可以尝试 config_hub 中的其他配置或通过传入不同的模型名字符串来更换模型类型与规模例如litgpt pretrain Gemma-2b当前支持的所有模型名可以通过不带参数直接执行litgpt pretrain列出在 litgpt/pretrain.py 中传入model_namelist会打印name_to_config中全部可用模型并退出。tinyllama.yaml 配置逐项解读config_hub/pretrain/tinyllama.yaml 是本次预训练的完整配置其核心字段如下model_name: tiny-llama-1.1b指定预训练模型来自litgpt.config中的名字与model_config互斥out_dir: out/pretrain/tiny-llamacheckpoint 与日志输出目录若在 Lightning Studio Job 中运行可在/teamspace/jobs/job-name/share找到precision: bf16-mixed可选bf16-true、bf16-mixed、32-trueresume: false支持true从out_dir最新 checkpoint 恢复找不到会报错、auto恢复但不报错或具体 checkpoint 路径data: TinyLlama数据模块默认即litgpt.data.TinyLlamatrain段对应 litgpt/args.py 的TrainArgssave_interval: 1000每 1000 个优化器步保存一次 checkpointlog_interval: 1每次迭代记录日志global_batch_size: 512跨所有数据并行 rank 的、两次优化器更新之间的样本数micro_batch_size: 4每个数据并行 rank 的微批大小梯度累积次数 batch_size(devices, num_nodes) // micro_batch_size见TrainArgs.gradient_accumulation_iterslr_warmup_steps: 2000学习率 warmup 迭代数max_tokens: 30000000000003 万亿总 token与 TinyLlama 官方训练规模一致max_seq_length: 2048限制样本长度max_norm: 1.0梯度裁剪范数min_lr: 4.0e-05余弦退火的最低学习率eval段对应EvalArgsinterval: 1000每 1000 步评估一次、max_iters: 100、initial_validation/final_validation: falseoptimizer段torch.optim.AdamWlr: 4e-4、weight_decay: 0.1、betas: (0.9, 0.95)devices: auto、num_nodes: 1、seed: 42tokenizer_dir: checkpoints/meta-llama/Llama-2-7b-hf预处理时使用的 tokenizer 目录logger_name: tensorboard默认 logger。数据模块SlimPajama 与 Starcoder 的混合默认数据模块TinyLlama定义在 litgpt/data/tinyllama.py。它要求data_path下存在slimpajama/train、slimpajama/val以及启用时还包括starcoder目录即上一步预处理脚本的输出connect()中seq_length max_seq_length 1——多出的 1 个 token 作为下一 token 的预测目标训练集使用CombinedStreamingDataset以0.693584 : 0.306416的比例混合 SlimPajama 与 Starcoder见 litgpt/data/tinyllama.py并可用--data.use_starcoderFalse关闭 Starcoder 数据每个StreamingDataset使用TokensLoader(block_sizeseq_length)切分 token 块训练时shuffleTrue, drop_lastTrue校验集仅使用 SlimPajama 的 val split。训练循环中的实现细节在 litgpt/pretrain.py 中多卡devices * num_nodes 1时使用FSDPStrategy(auto_wrap_policy{Block}, state_dict_typefull, sharding_strategyHYBRID_SHARD)即按Block自动包裹的混合分片 FSDP单卡时strategyauto模型以fabric.init_module(empty_initTrue)方式在 meta 设备上初始化随后加载权重torch.compile(model)默认开启学习率调度在get_lr()中实现见 litgpt/pretrain.py前warmup_iters步线性上升之后按余弦曲线衰减到min_lr与 TinyLlama 官方Cosine with 2000 warmup steps一致训练期间通过ThroughputMonitor(fabric, window_size5)监控吞吐并用measure_flops打印实测 TFLOPs训练结束会打印总 token 数、训练耗时、tok/s 与显存占用并保存out_dir/final/lit_model.pth最终 checkpoint。调整超参数与避免 OOM脚本通过命令行暴露了大量可调超参数。例如--train.micro_batch_size应根据可用 GPU 显存调整如果显存不足可以减小micro_batch_size并相应增大梯度累积步数来保持全局批大小不变。更详细的显存优化建议请参阅 tutorials/oom.md。实验跟踪预训练脚本默认只做最小化日志但长周期实验建议切换到正式的实验追踪器。LitGPT 支持多种 logger通过--logger_name指定TensorBoard默认本地可视化CSV Logger本地 CSV 日志WandB云端实验追踪--logger_namewandb仓库示例配置中已演示如何接入任意实验追踪框架MLflowMLflow 实验追踪LitLoggerLightning AI 原生实验追踪--logger_namelitlogger。多机扩展单机训练需要数周时间加速需要集群。进入集群后可参照 Lightning Fabric 的多机启动说明Lightning AI 云、SLURM、裸机集群、MPI 等方案跨机器启动脚本。仓库中的litgpt pretrain本身支持--num_nodes参数FSDP 的HYBRID_SHARD分片策略也正是为多机场景设计的。恢复训练预训练期间保存的 checkpoint 包含恢复所需的全部信息。只需在重新运行脚本时加上--resume参数litgpt pretrain tiny-llama\ --config config_hub/pretrain/tinyllama.yaml \ --resume out/pretrain/tiny-llama/step-00060500重要每个 checkpoint 是一个目录。请指向该目录本身而不是指向其中的lit_model.pth文件。恢复逻辑在 litgpt/pretrain.pystate中包含 model、optimizer、train_dataloader、iter_num与step_countfind_resume_path()解析 resume 参数后由fabric.load(resume, state)一次性恢复全部状态。resume还支持true/auto两种快捷方式true从out_dir的最新 checkpoint 恢复无 checkpoint 时报错auto恢复最新 checkpoint 但不报错详见 litgpt/pretrain.py 的参数说明。导出 checkpoint训练完成后需要把 checkpoint 转换为可被评估、推理、微调等流程加载的格式litgpt convert_pretrained_checkpoint out/pretrain/tiny-llama/step-00060500 \ --output_dir checkpoints/tiny-llama/final转换由 litgpt/scripts/convert_pretrained_checkpoint.py 中的convert_pretrained_checkpoint(checkpoint_dir, output_dir)实现。转换完成后输出目录应包含以下文件checkpoints/tiny-llama/final ├── model_config.yaml ├── lit_model.pth ├── tokenizer_config.json ├── tokenizer.json └── tokenizer.model此后你可以直接用这个 checkpoint 目录运行 评估、推理、LoRA 微调或继续 转换模型格式。项目模板可选Lightning AI 官方提供了一系列与 LitGPT 预训练相关的 Lightning Studio 模板可在具备多 GPU 与多节点支持的可复现环境中直接运行包括Prepare the TinyLlama 1T token dataset准备 TinyLlama 1T token 数据集、Pretrain LLMs - TinyLlama 1.1B预训练 TinyLlama 1.1B、Continued Pretraining with TinyLlama 1.1BTinyLlama 1.1B 持续预训练。无需本地环境时可直接从这些模板入手需要零配置快速体验预训练流程的读者也可以参考 Lightning AI 官方的 TinyLlama 预训练 Studio。【免费下载链接】litgpt20 high-performance LLMs with recipes to pretrain, finetune and deploy at scale.项目地址: https://gitcode.com/GitHub_Trending/li/litgpt创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
