DeepSpeedExamples AutoTP 实战：在 HuggingFace Trainer 中启用张量并行微调（hf

示例工程

【免费下载链接】DeepSpeedExamples

Example 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 Parallel，TP）自动化能力：它能够识别典型模型中的参数模式（例如注意力投影、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.json	DeepSpeed 配置模板（含`${zero_stage}`、`${autotp_size}`占位符）
configs/ds_config.json	模板渲染后的实际生效配置
train.py	基于 HF Trainer 的 SFT 训练入口
train_bench_length.py	为基准测试而设计的变体：将序列统一 padding 到`max_length`
alpaca_data.json	Alpaca 指令微调训练数据
utils.py	数据加载 / OpenAI 解码等辅助工具
requirements.txt	运行依赖清单

三、一键启动：run.sh 的六种并行模式

README 给出的启动方式极其简洁：bash run.sh或bash run.sh MODE。MODE 参数（默认zero2tp）决定了 ZeRO 阶段与 AutoTP 大小的组合。梳理 run.sh 源码，六种模式如下：

MODE	ZERO_STAGE	AUTOTP_SIZE	每设备 batch（`mbs=2`时）	说明
`zero1tp`	1	4	`mbs * 4 = 8`	ZeRO-1 + AutoTP=4
`zero2tp`（默认）	2	4	`mbs * 4 = 8`	ZeRO-2 + AutoTP=4
`zero1`	1	0	`mbs = 2`	纯 ZeRO-1，不启用 TP
`zero2`	2	0	`mbs = 2`	纯 ZeRO-2，不启用 TP
`zero3`	3	0	`mbs = 2`	ZeRO-3 全参数切分，不启用 TP
`tp`	0	8	`mbs * 8 = 16`	纯张量并行，ZeRO 关闭

脚本开头的关键变量逻辑如下（run.sh第 1-47 行）：

# Default to a public HF model for out-of-the-box runs. weight_path=facebook/opt-125m export WANDB_MODE=disabled num_gpus=${NUM_GPUS:-8} epoch=3 mbs=2 MODE=${1:-zero2tp}

几点值得注意的细节：

默认模型为公开的facebook/opt-125m，开箱即可运行，无需额外准备权重。
WANDB_MODE=disabled：关闭 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_FILE="configs/ds_config_temp.json" OUTPUT_FILE="configs/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 侧自动决定是否启用 BF16；
"optimizer"/"scheduler"：显式声明使用AdamW与WarmupDecayLR，但所有超参均为"auto"，即完全交由 HF Trainer 从命令行超参接管，避免配置二义性；
"zero_optimization": {"stage": ${zero_stage}}：ZeRO 阶段由run.sh的 MODE 决定（0/1/2/3）；
"gather_16bit_weights_on_model_save": true：保存模型时将 16 位权重聚合成完整权重。train.py中trainer.save_model(output_dir=training_args.output_dir)的注释明确提示必须开启此项才能正常导出 HF 权重（见 train.py 第 268-271 行的注释）；
"tensor_parallel": {"autotp_size": ${autotp_size}}：AutoTP 的唯一开关。设为 0 表示关闭张量并行，设为 N 表示在 N 个 rank 间自动切分参数。这正是本示例区别于常规 HF+DS 集成示例的关键配置；
"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 时使用padding="max_length"（固定 pad 到max_length）而非padding="longest"，保证每个样本序列等长、算力可对比；
将 pad token 对应的 label 也置为IGNORE_INDEX，避免填充位参与损失计算；
缓存文件名带长度后缀（dataset_dict{max_length}.pkl），不同长度互不冲突。

5.3 内存观测回调

train.py第 226-253 行内置了一个MemoryCallback（TrainerCallback子类），在每个 step 结束时调用see_memory_usage打印：

GPU 当前显存（MA）、峰值显存（Max_MA）、缓存显存（CA）、峰值缓存（Max_CA），单位 GB；
CPU 虚拟内存用量与百分比（基于psutil）；
每次打印后调用reset_peak_memory_stats()重置峰值计数器。

该回调只在 rank 0 输出（dist.is_initialized() and not dist.get_rank() == 0时直接 return），因此无需额外改动即可在分布式训练中观察显存曲线——这正是 README 所说的"只修改了日志"。

5.4 训练启动与模型保存

trainer = Trainer( model=model, processing_class=tokenizer, args=training_args, callbacks=[MemoryCallback], **data_module, ) trainer.train() trainer.save_model(output_dir=training_args.output_dir)

保存环节依赖于配置中的gather_16bit_weights_on_model_save=true，注释中还预留了分布式 checkpoint 的加载恢复示例（trainer.save_state()与trainer.train(resume_from_checkpoint=...)，见第 264-266 行注释）。

六、运行前提与依赖

requirements.txt 列出核心依赖：

transformers==4.50.1 deepspeed>=0.16.4 accelerate==1.6.0 numpy rouge_score fire openai==0.28.0 torch sentencepiece tokenizers>=0.13.3

注意事项：

AutoTP 配置项（tensor_parallel.autotp_size）依赖deepspeed>=0.16.4，请确保环境满足版本下限；
openai==0.28.0来自 Alpaca 脚本的生成/评测辅助逻辑（utils.py 中的openai_completion等），与训练主路径解耦；
训练数据alpaca_data.json已随仓库提供，默认模型为facebook/opt-125m，可离线直接验证。

七、与其他 AutoTP 示例的定位差异

仓库的training/tensor_parallel/下还有两个关联示例，可用于对照理解本示例的边界：

basic_example：使用deepspeed.initialize(..., mpu=mpu)显式传入 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_SIZE`；`run.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 与 AutoTP=4/8 之间自由组合。结合仓库内 basic_example 与 custom_patterns 两个示例，你可以从"开箱即用"一路深入到"自定义切分规则"，完整覆盖 AutoTP 训练的入门与进阶场景。

示例工程

【免费下载链接】DeepSpeedExamples

Example models using DeepSpeed

项目地址：https://gitcode.com/gh_mirrors/de/DeepSpeedExamples

点击查看免费下载

DeepSpeedExamples AutoTP 实战：在 HuggingFace Trainer 中启用张量并行微调（hf_integration 全解析）

【免费下载链接】DeepSpeedExamples

导读

一、示例背景：从 Alpaca 微调脚本到 AutoTP 集成

二、目录结构与文件职责

三、一键启动：run.sh 的六种并行模式

设备网格一致性保护（重要坑点）

配置模板渲染

启动命令与训练超参

四、DeepSpeed 配置模板逐项解析

五、训练入口：train.py 的数据管线与内存观测

5.1 指令模板与监督信号构造

5.2 数据缓存与基准变体

5.3 内存观测回调

5.4 训练启动与模型保存

六、运行前提与依赖

七、与其他 AutoTP 示例的定位差异

八、常见问题速查

结语

【免费下载链接】DeepSpeedExamples

相关推荐

蓝鲸PaaS路线图前瞻：AI开发、云原生与可观测性的下一步演进

DeepSpeed 训练场景自动张量并行（AutoTP）实战：Tensor Parallel + ZeRO 的混合并行配置与源码解析

HuggingFace Transformers教程：使用Trainer API微调模型

告别 2 小时上限：Wand-Enhancer 本地解锁 Wand 专业版完整指南

Windows 8 Build 7779安装教程

STM32驱动WS2812呼吸灯：PWM+DMA方案实现顺滑渐变

new Object() 到底做了啥？

python-mastery Exercise 9.3 实战：用 `all` 控制包导出符号并拆分模块结构

Cadence Capture CIS建库实战：K9F2G08U0M NAND Flash原理图符号与数据库管理

【免费下载链接】DeepSpeedExamples

导读

一、示例背景：从 Alpaca 微调脚本到 AutoTP 集成

二、目录结构与文件职责

三、一键启动：run.sh 的六种并行模式

设备网格一致性保护（重要坑点）

配置模板渲染

启动命令与训练超参

四、DeepSpeed 配置模板逐项解析

五、训练入口：train.py 的数据管线与内存观测

5.1 指令模板与监督信号构造

5.2 数据缓存与基准变体

5.3 内存观测回调

5.4 训练启动与模型保存

六、运行前提与依赖

七、与其他 AutoTP 示例的定位差异

八、常见问题速查

结语

【免费下载链接】DeepSpeedExamples

相关推荐

蓝鲸PaaS路线图前瞻：AI开发、云原生与可观测性的下一步演进

DeepSpeed 训练场景自动张量并行（AutoTP）实战：Tensor Parallel + ZeRO 的混合并行配置与源码解析

HuggingFace Transformers教程：使用Trainer API微调模型

告别 2 小时上限：Wand-Enhancer 本地解锁 Wand 专业版完整指南

Windows 8 Build 7779安装教程

STM32驱动WS2812呼吸灯：PWM+DMA方案实现顺滑渐变

new Object() 到底做了啥？

python-mastery Exercise 9.3 实战：用 `__all__` 控制包导出符号并拆分模块结构

Cadence Capture CIS建库实战：K9F2G08U0M NAND Flash原理图符号与数据库管理

python-mastery Exercise 9.3 实战：用 `all` 控制包导出符号并拆分模块结构