拓冰网站定制 
示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载导读本文以 training/tensor_parallel/hf_integration 目录下的官方示例为蓝本讲解如何在几乎不改动 Hugging FaceTrainer训练代码的前提下仅通过 DeepSpeed 配置文件中的tensor_parallel.autotp_size一项为 GPT/OPT/Llama 类因果语言模型开启 AutoTP自动张量并行训练。读完本文你将掌握 AutoTP 与 ZeRO 各阶段Zero-1/2/3的搭配方式、一键脚本run.sh的六种运行模式、DS 配置模板中每个关键字段的语义以及基于斯坦福 Alpaca 风格监督微调SFT数据管线的完整实战路径。一、示例背景从 Alpaca 微调脚本到 AutoTP 集成本示例由 斯坦福 Alpaca 项目改编而来。其核心思路是训练代码train.py几乎保持原样仅修改 DeepSpeed 配置与日志记录以此演示 AutoTP 与 Hugging FaceTrainer的无缝集成README 原文明确指出 We only modified the DeepSpeed config and logging, as an example use case。所谓 AutoTP即 DeepSpeed 的张量并行Tensor ParallelTP自动化能力它能够识别典型模型中的参数模式例如注意力投影、MLP 上投影/下投影自动为这些参数施加合适的切分规则因此对于受支持的模型架构如 Llama 系列无需手工编写任何切分规则。这一点在仓库的 training/tensor_parallel/basic_example/README.md 中有明确说明AutoTP recognizes supported model architectures (for example, Llama) and automatically partitions parameters, so you do not need to specify any manual partitioning rules for those models.因此这个hf_integration示例回答了一个非常实际的工程问题我已经在用 Hugging Face Trainer 做 SFT如何最省事地吃上张量并行答案是改一行 DS 配置、套一层启动脚本其余全部交给框架。二、目录结构与文件职责文件职责README.md示例说明与启动方式run.sh一键启动脚本内置六种并行模式configs/ds_config_temp.jsonDeepSpeed 配置模板含${zero_stage}、${autotp_size}占位符configs/ds_config.json模板渲染后的实际生效配置train.py基于 HF Trainer 的 SFT 训练入口train_bench_length.py为基准测试而设计的变体将序列统一 padding 到max_lengthalpaca_data.jsonAlpaca 指令微调训练数据utils.py数据加载 / OpenAI 解码等辅助工具requirements.txt运行依赖清单三、一键启动run.sh 的六种并行模式README 给出的启动方式极其简洁bash run.sh或bash run.sh MODE。MODE 参数默认zero2tp决定了 ZeRO 阶段与 AutoTP 大小的组合。梳理 run.sh 源码六种模式如下MODEZERO_STAGEAUTOTP_SIZE每设备 batchmbs2时说明zero1tp14mbs * 4 8ZeRO-1 AutoTP4zero2tp默认24mbs * 4 8ZeRO-2 AutoTP4zero110mbs 2纯 ZeRO-1不启用 TPzero220mbs 2纯 ZeRO-2不启用 TPzero330mbs 2ZeRO-3 全参数切分不启用 TPtp08mbs * 8 16纯张量并行ZeRO 关闭脚本开头的关键变量逻辑如下run.sh第 1-47 行# Default to a public HF model for out-of-the-box runs. weight_pathfacebook/opt-125m export WANDB_MODEdisabled num_gpus${NUM_GPUS:-8} epoch3 mbs2 MODE${1:-zero2tp}几点值得注意的细节默认模型为公开的facebook/opt-125m开箱即可运行无需额外准备权重。WANDB_MODEdisabled关闭 Weights Biases 在线同步避免日志阻塞。num_gpus默认 8可通过环境变量NUM_GPUS覆盖。在 TP 模式下per_device_train_batch_size mbs * AUTOTP_SIZE即每个 TP 组的整体 batch 由多个 TP rank 分摊这是张量并行的标准语义全局 batch 不变被切到各 TP rank 上。设备网格一致性保护重要坑点run.sh第 37-42 行包含一段重要的兼容性处理# HF Trainer Accelerate currently builds a 1D device mesh of size AUTOTP_SIZE. # If num_gpus AUTOTP_SIZE, ranks outside the mesh fail during init_device_mesh. if [ $AUTOTP_SIZE -gt 1 ] [ $num_gpus -ne $AUTOTP_SIZE ]; then echo Adjusting num_gpus to AUTOTP_SIZE$AUTOTP_SIZE to avoid device_mesh init failure. num_gpus$AUTOTP_SIZE fi即当启用 AutoTP 且AUTOTP_SIZE 1时Hugging FaceTrainer Accelerate 会构建一个大小为AUTOTP_SIZE的一维设备网格若实际启动的 GPU 数多于AUTOTP_SIZE网格外的 rank 会在init_device_mesh阶段初始化失败。脚本会自动把num_gpus收敛到AUTOTP_SIZE避免该问题。这是纯 HF Trainer 集成场景下与底层 AutoTP API见basic_example中自由组建 TP/DP 二维网格的显著差异务必注意。配置模板渲染脚本通过sed将模板中的占位符替换为实际数值生成生效配置run.sh第 43-47 行TEMPLATE_FILEconfigs/ds_config_temp.json OUTPUT_FILEconfigs/ds_config.json sed -e s/\${zero_stage}/${ZERO_STAGE}/g \ -e s/\${autotp_size}/${AUTOTP_SIZE}/g \ $TEMPLATE_FILE $OUTPUT_FILE默认模式zero2tp渲染出的 configs/ds_config.json 中即为zero_optimization: {stage: 2, ...}与tensor_parallel: {autotp_size: 4}。启动命令与训练超参run.sh第 50-71 行最终以deepspeed启动器拉起训练deepspeed --num_gpus $num_gpus \ --master_port 51336 train.py \ --model_name_or_path $weight_path \ --data_path ./alpaca_data.json \ --bf16 True \ --output_dir out_load_test/$MODE \ --num_train_epochs $epoch \ --gradient_checkpointing false \ --per_device_train_batch_size $per_device_train_batch_size \ --per_device_eval_batch_size 1 \ --eval_strategy no \ --save_strategy steps \ --save_steps 10000 \ --gradient_accumulation_steps 4 \ --learning_rate 0 \ --learning_rate 2e-5 \ --weight_decay 0. \ --warmup_steps 0 \ --lr_scheduler_type cosine \ --logging_steps 1 \ --tf32 True \ --deepspeed ./configs/ds_config.json关键超参一览--bf16 True半精度训练与配置模板中bf16: {enabled: auto}呼应--gradient_accumulation_steps 4梯度累积扩大有效 batch--gradient_checkpointing false本示例刻意关闭梯度检查点便于观察内存基线可通过修改启动参数开启--lr_scheduler_type cosine余弦学习率调度--tf32 True在 Ampere GPU 上启用 TF32 加速--learning_rate 2e-5注意脚本中先传了一次--learning_rate 0再传2e-5后者覆盖前者实际生效值为2e-5--save_strategy steps --save_steps 10000以步数为单位保存 checkpoint--eval_strategy no本示例不跑评估。四、DeepSpeed 配置模板逐项解析模板 configs/ds_config_temp.json 是整套示例的灵魂全文如下{ bf16: { enabled: auto }, optimizer: { type: AdamW, params: { lr: auto, betas: auto, eps: auto, weight_decay: auto } }, scheduler: { type: WarmupDecayLR, params: { total_num_steps: auto, warmup_min_lr: auto, warmup_max_lr: auto, warmup_num_steps: auto } }, zero_optimization: { stage: ${zero_stage}, gather_16bit_weights_on_model_save: true }, tensor_parallel:{ autotp_size: ${autotp_size} }, gradient_accumulation_steps: auto, gradient_clipping: auto, steps_per_print: 1, train_batch_size: auto, train_micro_batch_size_per_gpu: auto, wall_clock_breakdown: false }各字段语义bf16: {enabled: auto}与命令行--bf16 True联动由 HF Trainer 侧自动决定是否启用 BF16optimizer/scheduler显式声明使用AdamW与WarmupDecayLR但所有超参均为auto即完全交由 HF Trainer 从命令行超参接管避免配置二义性zero_optimization: {stage: ${zero_stage}}ZeRO 阶段由run.sh的 MODE 决定0/1/2/3gather_16bit_weights_on_model_save: true保存模型时将 16 位权重聚合成完整权重。train.py中trainer.save_model(output_dirtraining_args.output_dir)的注释明确提示必须开启此项才能正常导出 HF 权重见 train.py 第 268-271 行的注释tensor_parallel: {autotp_size: ${autotp_size}}AutoTP 的唯一开关。设为 0 表示关闭张量并行设为 N 表示在 N 个 rank 间自动切分参数。这正是本示例区别于常规 HFDS 集成示例的关键配置gradient_accumulation_steps、train_batch_size、train_micro_batch_size_per_gpu、gradient_clipping均为auto全部由 HF Trainer 从命令行参数推导避免配置冲突steps_per_print: 1每个 step 打印一次日志配合--logging_steps 1做密集观测wall_clock_breakdown: false关闭耗时细粒度分解减少日志开销。五、训练入口train.py 的数据管线与内存观测train.py 完整保留了 Alpaca 风格的 SFT 数据管线主要环节如下。5.1 指令模板与监督信号构造代码第 31-42 行定义了 Alpaca 指令模板区分带输入与不带输入两种形态PROMPT_DICT { prompt_input: ( Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n ### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response: ), prompt_no_input: ( Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n ### Instruction:\n{instruction}\n\n### Response: ), }在SupervisedDataset.__init__中训练样本 prompt target eos_token标签中 prompt 部分被置为IGNORE_INDEX -100第 113-125 行preprocess函数即只对回答部分计算损失这是指令微调的标准做法。5.2 数据缓存与基准变体SupervisedDataset会把 tokenize 结果以 pickle 形式缓存到本地dataset_dict.pkl二次运行直接加载节省重复 tokenize 的时间第 144-156 行。配套的 train_bench_length.py 则专门为基准测试设计tokenize 时使用paddingmax_length固定 pad 到max_length而非paddinglongest保证每个样本序列等长、算力可对比将 pad token 对应的 label 也置为IGNORE_INDEX避免填充位参与损失计算缓存文件名带长度后缀dataset_dict{max_length}.pkl不同长度互不冲突。5.3 内存观测回调train.py第 226-253 行内置了一个MemoryCallbackTrainerCallback子类在每个 step 结束时调用see_memory_usage打印GPU 当前显存MA、峰值显存Max_MA、缓存显存CA、峰值缓存Max_CA单位 GBCPU 虚拟内存用量与百分比基于psutil每次打印后调用reset_peak_memory_stats()重置峰值计数器。该回调只在 rank 0 输出dist.is_initialized() and not dist.get_rank() 0时直接 return因此无需额外改动即可在分布式训练中观察显存曲线——这正是 README 所说的只修改了日志。5.4 训练启动与模型保存trainer Trainer( modelmodel, processing_classtokenizer, argstraining_args, callbacks[MemoryCallback], **data_module, ) trainer.train() trainer.save_model(output_dirtraining_args.output_dir)保存环节依赖于配置中的gather_16bit_weights_on_model_savetrue注释中还预留了分布式 checkpoint 的加载恢复示例trainer.save_state()与trainer.train(resume_from_checkpoint...)见第 264-266 行注释。六、运行前提与依赖requirements.txt 列出核心依赖transformers4.50.1 deepspeed0.16.4 accelerate1.6.0 numpy rouge_score fire openai0.28.0 torch sentencepiece tokenizers0.13.3注意事项AutoTP 配置项tensor_parallel.autotp_size依赖deepspeed0.16.4请确保环境满足版本下限openai0.28.0来自 Alpaca 脚本的生成/评测辅助逻辑utils.py 中的openai_completion等与训练主路径解耦训练数据alpaca_data.json已随仓库提供默认模型为facebook/opt-125m可离线直接验证。七、与其他 AutoTP 示例的定位差异仓库的training/tensor_parallel/下还有两个关联示例可用于对照理解本示例的边界basic_example使用deepspeed.initialize(..., mpumpu)显式传入 TP/DP 进程组与ModelParallelUnit并以合成 token 批量数据验证 AutoTP 设置。它面向底层 API 用法验证展示了tp_size * dp_size world_size的二维网格构建方式build_tp_dp_groups。custom_patterns当模型不在 AutoTP 默认支持名单时如 GPT-NeoX/Pythia 的融合query_key_value投影通过tensor_parallel.partition_config编写layer_specs正则匹配参数名 partition_type: column/row 可选shape手动指定切分规则。而本hf_integration示例的独特价值在于用户完全不接触mpu、partition_config等底层细节仅在 HF Trainer 生态内通过run.sh DS 配置模板即可跑通 AutoTP。如果你的模型属于受支持的架构如 Llama这就是最快的上手路径若模型不受支持则需退回到 custom_patterns 方式补充分区规则。八、常见问题速查现象原因与对策启用 TP 后init_device_mesh失败Accelerate 构建的一维设备网格大小为AUTOTP_SIZE启动 GPU 数须等于AUTOTP_SIZErun.sh已自动修正手动起训时请保持--num_gpus AUTOTP_SIZE保存的模型不是完整权重需在 DS 配置中保持gather_16bit_weights_on_model_save: true再调用trainer.save_model显存观测只看到 rank 0 输出see_memory_usage有意只在 rank 0 打印属预期行为数据集重复 tokenize 耗时首次运行后dataset_dict.pkl缓存已生成后续直接加载序列不等长导致基准不稳定使用 train_bench_length.py 按max_length统一 padding结语training/tensor_parallel/hf_integration用最小的代码改动演示了 DeepSpeed AutoTP 与 Hugging Face Trainer 的集成范式配置模板中一行tensor_parallel.autotp_size脚本中一套 MODE 切换即可在 ZeRO-0/1/2/3 与 AutoTP4/8 之间自由组合。结合仓库内 basic_example 与 custom_patterns 两个示例你可以从开箱即用一路深入到自定义切分规则完整覆盖 AutoTP 训练的入门与进阶场景。赞分享示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载相关推荐蓝鲸PaaS路线图前瞻AI开发、云原生与可观测性的下一步演进蓝鲸PaaS路线图前瞻AI开发、云原生与可观测性的下一步演进 蓝鲸智云 PaaS 平台blueking paas是一个开放式的开发平台帮助开发者快速创建后端云原生微服务前端企业应用开发者门户DeepSpeed 训练场景自动张量并行AutoTP实战Tensor Parallel ZeRO 的混合并行配置与源码解析DeepSpeed 训练场景自动张量并行AutoTP实战Tensor Parallel ZeRO 的混合并行配置与源码解析 本教程面向希望在训练阶段把人工智能大模型深度学习分布式训练预训练强化学习模型优化HuggingFace Transformers教程使用Trainer API微调模型HuggingFace Transformers教程使用Trainer API微调模型 前言 在自然语言处理领域预训练模型的微调 fine tuning 已文档教程人工智能NLP深度学习上一篇react-map-gl 中 Layer 组件深度解析用 React 声明式管理 Mapbox 图层下一篇如何快速上手DCS5分钟完成性能数据收集部署创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 DeepSpeedExamples AutoTP 实战：在 HuggingFace Trainer 中启用张量并行微调（hf_integration 全解析）
示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载导读本文以 training/tensor_parallel/hf_integration 目录下的官方示例为蓝本讲解如何在几乎不改动 Hugging FaceTrainer训练代码的前提下仅通过 DeepSpeed 配置文件中的tensor_parallel.autotp_size一项为 GPT/OPT/Llama 类因果语言模型开启 AutoTP自动张量并行训练。读完本文你将掌握 AutoTP 与 ZeRO 各阶段Zero-1/2/3的搭配方式、一键脚本run.sh的六种运行模式、DS 配置模板中每个关键字段的语义以及基于斯坦福 Alpaca 风格监督微调SFT数据管线的完整实战路径。一、示例背景从 Alpaca 微调脚本到 AutoTP 集成本示例由 斯坦福 Alpaca 项目改编而来。其核心思路是训练代码train.py几乎保持原样仅修改 DeepSpeed 配置与日志记录以此演示 AutoTP 与 Hugging FaceTrainer的无缝集成README 原文明确指出 We only modified the DeepSpeed config and logging, as an example use case。所谓 AutoTP即 DeepSpeed 的张量并行Tensor ParallelTP自动化能力它能够识别典型模型中的参数模式例如注意力投影、MLP 上投影/下投影自动为这些参数施加合适的切分规则因此对于受支持的模型架构如 Llama 系列无需手工编写任何切分规则。这一点在仓库的 training/tensor_parallel/basic_example/README.md 中有明确说明AutoTP recognizes supported model architectures (for example, Llama) and automatically partitions parameters, so you do not need to specify any manual partitioning rules for those models.因此这个hf_integration示例回答了一个非常实际的工程问题我已经在用 Hugging Face Trainer 做 SFT如何最省事地吃上张量并行答案是改一行 DS 配置、套一层启动脚本其余全部交给框架。二、目录结构与文件职责文件职责README.md示例说明与启动方式run.sh一键启动脚本内置六种并行模式configs/ds_config_temp.jsonDeepSpeed 配置模板含${zero_stage}、${autotp_size}占位符configs/ds_config.json模板渲染后的实际生效配置train.py基于 HF Trainer 的 SFT 训练入口train_bench_length.py为基准测试而设计的变体将序列统一 padding 到max_lengthalpaca_data.jsonAlpaca 指令微调训练数据utils.py数据加载 / OpenAI 解码等辅助工具requirements.txt运行依赖清单三、一键启动run.sh 的六种并行模式README 给出的启动方式极其简洁bash run.sh或bash run.sh MODE。MODE 参数默认zero2tp决定了 ZeRO 阶段与 AutoTP 大小的组合。梳理 run.sh 源码六种模式如下MODEZERO_STAGEAUTOTP_SIZE每设备 batchmbs2时说明zero1tp14mbs * 4 8ZeRO-1 AutoTP4zero2tp默认24mbs * 4 8ZeRO-2 AutoTP4zero110mbs 2纯 ZeRO-1不启用 TPzero220mbs 2纯 ZeRO-2不启用 TPzero330mbs 2ZeRO-3 全参数切分不启用 TPtp08mbs * 8 16纯张量并行ZeRO 关闭脚本开头的关键变量逻辑如下run.sh第 1-47 行# Default to a public HF model for out-of-the-box runs. weight_pathfacebook/opt-125m export WANDB_MODEdisabled num_gpus${NUM_GPUS:-8} epoch3 mbs2 MODE${1:-zero2tp}几点值得注意的细节默认模型为公开的facebook/opt-125m开箱即可运行无需额外准备权重。WANDB_MODEdisabled关闭 Weights Biases 在线同步避免日志阻塞。num_gpus默认 8可通过环境变量NUM_GPUS覆盖。在 TP 模式下per_device_train_batch_size mbs * AUTOTP_SIZE即每个 TP 组的整体 batch 由多个 TP rank 分摊这是张量并行的标准语义全局 batch 不变被切到各 TP rank 上。设备网格一致性保护重要坑点run.sh第 37-42 行包含一段重要的兼容性处理# HF Trainer Accelerate currently builds a 1D device mesh of size AUTOTP_SIZE. # If num_gpus AUTOTP_SIZE, ranks outside the mesh fail during init_device_mesh. if [ $AUTOTP_SIZE -gt 1 ] [ $num_gpus -ne $AUTOTP_SIZE ]; then echo Adjusting num_gpus to AUTOTP_SIZE$AUTOTP_SIZE to avoid device_mesh init failure. num_gpus$AUTOTP_SIZE fi即当启用 AutoTP 且AUTOTP_SIZE 1时Hugging FaceTrainer Accelerate 会构建一个大小为AUTOTP_SIZE的一维设备网格若实际启动的 GPU 数多于AUTOTP_SIZE网格外的 rank 会在init_device_mesh阶段初始化失败。脚本会自动把num_gpus收敛到AUTOTP_SIZE避免该问题。这是纯 HF Trainer 集成场景下与底层 AutoTP API见basic_example中自由组建 TP/DP 二维网格的显著差异务必注意。配置模板渲染脚本通过sed将模板中的占位符替换为实际数值生成生效配置run.sh第 43-47 行TEMPLATE_FILEconfigs/ds_config_temp.json OUTPUT_FILEconfigs/ds_config.json sed -e s/\${zero_stage}/${ZERO_STAGE}/g \ -e s/\${autotp_size}/${AUTOTP_SIZE}/g \ $TEMPLATE_FILE $OUTPUT_FILE默认模式zero2tp渲染出的 configs/ds_config.json 中即为zero_optimization: {stage: 2, ...}与tensor_parallel: {autotp_size: 4}。启动命令与训练超参run.sh第 50-71 行最终以deepspeed启动器拉起训练deepspeed --num_gpus $num_gpus \ --master_port 51336 train.py \ --model_name_or_path $weight_path \ --data_path ./alpaca_data.json \ --bf16 True \ --output_dir out_load_test/$MODE \ --num_train_epochs $epoch \ --gradient_checkpointing false \ --per_device_train_batch_size $per_device_train_batch_size \ --per_device_eval_batch_size 1 \ --eval_strategy no \ --save_strategy steps \ --save_steps 10000 \ --gradient_accumulation_steps 4 \ --learning_rate 0 \ --learning_rate 2e-5 \ --weight_decay 0. \ --warmup_steps 0 \ --lr_scheduler_type cosine \ --logging_steps 1 \ --tf32 True \ --deepspeed ./configs/ds_config.json关键超参一览--bf16 True半精度训练与配置模板中bf16: {enabled: auto}呼应--gradient_accumulation_steps 4梯度累积扩大有效 batch--gradient_checkpointing false本示例刻意关闭梯度检查点便于观察内存基线可通过修改启动参数开启--lr_scheduler_type cosine余弦学习率调度--tf32 True在 Ampere GPU 上启用 TF32 加速--learning_rate 2e-5注意脚本中先传了一次--learning_rate 0再传2e-5后者覆盖前者实际生效值为2e-5--save_strategy steps --save_steps 10000以步数为单位保存 checkpoint--eval_strategy no本示例不跑评估。四、DeepSpeed 配置模板逐项解析模板 configs/ds_config_temp.json 是整套示例的灵魂全文如下{ bf16: { enabled: auto }, optimizer: { type: AdamW, params: { lr: auto, betas: auto, eps: auto, weight_decay: auto } }, scheduler: { type: WarmupDecayLR, params: { total_num_steps: auto, warmup_min_lr: auto, warmup_max_lr: auto, warmup_num_steps: auto } }, zero_optimization: { stage: ${zero_stage}, gather_16bit_weights_on_model_save: true }, tensor_parallel:{ autotp_size: ${autotp_size} }, gradient_accumulation_steps: auto, gradient_clipping: auto, steps_per_print: 1, train_batch_size: auto, train_micro_batch_size_per_gpu: auto, wall_clock_breakdown: false }各字段语义bf16: {enabled: auto}与命令行--bf16 True联动由 HF Trainer 侧自动决定是否启用 BF16optimizer/scheduler显式声明使用AdamW与WarmupDecayLR但所有超参均为auto即完全交由 HF Trainer 从命令行超参接管避免配置二义性zero_optimization: {stage: ${zero_stage}}ZeRO 阶段由run.sh的 MODE 决定0/1/2/3gather_16bit_weights_on_model_save: true保存模型时将 16 位权重聚合成完整权重。train.py中trainer.save_model(output_dirtraining_args.output_dir)的注释明确提示必须开启此项才能正常导出 HF 权重见 train.py 第 268-271 行的注释tensor_parallel: {autotp_size: ${autotp_size}}AutoTP 的唯一开关。设为 0 表示关闭张量并行设为 N 表示在 N 个 rank 间自动切分参数。这正是本示例区别于常规 HFDS 集成示例的关键配置gradient_accumulation_steps、train_batch_size、train_micro_batch_size_per_gpu、gradient_clipping均为auto全部由 HF Trainer 从命令行参数推导避免配置冲突steps_per_print: 1每个 step 打印一次日志配合--logging_steps 1做密集观测wall_clock_breakdown: false关闭耗时细粒度分解减少日志开销。五、训练入口train.py 的数据管线与内存观测train.py 完整保留了 Alpaca 风格的 SFT 数据管线主要环节如下。5.1 指令模板与监督信号构造代码第 31-42 行定义了 Alpaca 指令模板区分带输入与不带输入两种形态PROMPT_DICT { prompt_input: ( Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n ### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response: ), prompt_no_input: ( Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n ### Instruction:\n{instruction}\n\n### Response: ), }在SupervisedDataset.__init__中训练样本 prompt target eos_token标签中 prompt 部分被置为IGNORE_INDEX -100第 113-125 行preprocess函数即只对回答部分计算损失这是指令微调的标准做法。5.2 数据缓存与基准变体SupervisedDataset会把 tokenize 结果以 pickle 形式缓存到本地dataset_dict.pkl二次运行直接加载节省重复 tokenize 的时间第 144-156 行。配套的 train_bench_length.py 则专门为基准测试设计tokenize 时使用paddingmax_length固定 pad 到max_length而非paddinglongest保证每个样本序列等长、算力可对比将 pad token 对应的 label 也置为IGNORE_INDEX避免填充位参与损失计算缓存文件名带长度后缀dataset_dict{max_length}.pkl不同长度互不冲突。5.3 内存观测回调train.py第 226-253 行内置了一个MemoryCallbackTrainerCallback子类在每个 step 结束时调用see_memory_usage打印GPU 当前显存MA、峰值显存Max_MA、缓存显存CA、峰值缓存Max_CA单位 GBCPU 虚拟内存用量与百分比基于psutil每次打印后调用reset_peak_memory_stats()重置峰值计数器。该回调只在 rank 0 输出dist.is_initialized() and not dist.get_rank() 0时直接 return因此无需额外改动即可在分布式训练中观察显存曲线——这正是 README 所说的只修改了日志。5.4 训练启动与模型保存trainer Trainer( modelmodel, processing_classtokenizer, argstraining_args, callbacks[MemoryCallback], **data_module, ) trainer.train() trainer.save_model(output_dirtraining_args.output_dir)保存环节依赖于配置中的gather_16bit_weights_on_model_savetrue注释中还预留了分布式 checkpoint 的加载恢复示例trainer.save_state()与trainer.train(resume_from_checkpoint...)见第 264-266 行注释。六、运行前提与依赖requirements.txt 列出核心依赖transformers4.50.1 deepspeed0.16.4 accelerate1.6.0 numpy rouge_score fire openai0.28.0 torch sentencepiece tokenizers0.13.3注意事项AutoTP 配置项tensor_parallel.autotp_size依赖deepspeed0.16.4请确保环境满足版本下限openai0.28.0来自 Alpaca 脚本的生成/评测辅助逻辑utils.py 中的openai_completion等与训练主路径解耦训练数据alpaca_data.json已随仓库提供默认模型为facebook/opt-125m可离线直接验证。七、与其他 AutoTP 示例的定位差异仓库的training/tensor_parallel/下还有两个关联示例可用于对照理解本示例的边界basic_example使用deepspeed.initialize(..., mpumpu)显式传入 TP/DP 进程组与ModelParallelUnit并以合成 token 批量数据验证 AutoTP 设置。它面向底层 API 用法验证展示了tp_size * dp_size world_size的二维网格构建方式build_tp_dp_groups。custom_patterns当模型不在 AutoTP 默认支持名单时如 GPT-NeoX/Pythia 的融合query_key_value投影通过tensor_parallel.partition_config编写layer_specs正则匹配参数名 partition_type: column/row 可选shape手动指定切分规则。而本hf_integration示例的独特价值在于用户完全不接触mpu、partition_config等底层细节仅在 HF Trainer 生态内通过run.sh DS 配置模板即可跑通 AutoTP。如果你的模型属于受支持的架构如 Llama这就是最快的上手路径若模型不受支持则需退回到 custom_patterns 方式补充分区规则。八、常见问题速查现象原因与对策启用 TP 后init_device_mesh失败Accelerate 构建的一维设备网格大小为AUTOTP_SIZE启动 GPU 数须等于AUTOTP_SIZErun.sh已自动修正手动起训时请保持--num_gpus AUTOTP_SIZE保存的模型不是完整权重需在 DS 配置中保持gather_16bit_weights_on_model_save: true再调用trainer.save_model显存观测只看到 rank 0 输出see_memory_usage有意只在 rank 0 打印属预期行为数据集重复 tokenize 耗时首次运行后dataset_dict.pkl缓存已生成后续直接加载序列不等长导致基准不稳定使用 train_bench_length.py 按max_length统一 padding结语training/tensor_parallel/hf_integration用最小的代码改动演示了 DeepSpeed AutoTP 与 Hugging Face Trainer 的集成范式配置模板中一行tensor_parallel.autotp_size脚本中一套 MODE 切换即可在 ZeRO-0/1/2/3 与 AutoTP4/8 之间自由组合。结合仓库内 basic_example 与 custom_patterns 两个示例你可以从开箱即用一路深入到自定义切分规则完整覆盖 AutoTP 训练的入门与进阶场景。赞分享示例工程【免费下载链接】DeepSpeedExamplesExample models using DeepSpeed项目地址https://gitcode.com/gh_mirrors/de/DeepSpeedExamples点击查看免费下载相关推荐蓝鲸PaaS路线图前瞻AI开发、云原生与可观测性的下一步演进蓝鲸PaaS路线图前瞻AI开发、云原生与可观测性的下一步演进 蓝鲸智云 PaaS 平台blueking paas是一个开放式的开发平台帮助开发者快速创建后端云原生微服务前端企业应用开发者门户DeepSpeed 训练场景自动张量并行AutoTP实战Tensor Parallel ZeRO 的混合并行配置与源码解析DeepSpeed 训练场景自动张量并行AutoTP实战Tensor Parallel ZeRO 的混合并行配置与源码解析 本教程面向希望在训练阶段把人工智能大模型深度学习分布式训练预训练强化学习模型优化HuggingFace Transformers教程使用Trainer API微调模型HuggingFace Transformers教程使用Trainer API微调模型 前言 在自然语言处理领域预训练模型的微调 fine tuning 已文档教程人工智能NLP深度学习上一篇react-map-gl 中 Layer 组件深度解析用 React 声明式管理 Mapbox 图层下一篇如何快速上手DCS5分钟完成性能数据收集部署创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考