Vertex AI多节点分布式训练是否支持HuggingFace Trainer+Deepspeed？

1. 是否支持DeepSpeed？ #

是的，Vertex AI Training完全支持结合Hugging Face Trainer与DeepSpeed进行分布式训练。尽管官方示例多基于PyTorch原生分布式策略，但Vertex AI的多节点训练环境提供了DeepSpeed运行所需的网络互连、进程管理基础，与Hugging Face Trainer的DeepSpeed集成逻辑完全兼容。

2. 多节点训练中集成DeepSpeed的步骤 #

步骤1：准备训练脚本与DeepSpeed配置 #

在训练脚本中，通过TrainingArguments指定deepspeed参数关联配置文件，无需手动配置PyTorch原生分布式策略：

from transformers import TrainingArguments, Trainer
from your_model_module import YourModel
from your_data_module import get_train_dataset

# 初始化模型与数据集
model = YourModel.from_pretrained("your-model-name")
train_dataset = get_train_dataset()

# 设置训练参数，指定DeepSpeed配置文件
training_args = TrainingArguments(
    output_dir="./results",
    per_device_train_batch_size=8,
    num_train_epochs=3,
    deepspeed="./ds_config.json",  # 指向DeepSpeed配置文件
    logging_dir="./logs",
    logging_steps=10,
    # 其他训练相关参数
)

# 初始化Trainer并启动训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

trainer.train()

DeepSpeed配置文件示例（ds_config.json），可根据训练需求调整Zero优化阶段、批量大小等：

{
    "train_batch_size": 64,
    "train_micro_batch_size_per_gpu": 8,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 3e-5,
            "betas": [0.9, 0.999],
            "eps": 1e-8
        }
    },
    "gradient_clipping": 1.0,
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "logging": {
        "steps_per_print": 10
    }
}

步骤2：配置Vertex AI训练环境与资源 #

• 机器选型：选择带GPU的机器类型（如n1-standard-8搭配Tesla T4，或a2-highgpu-1g搭配A100），多节点训练时通过machine-count指定节点数量（如machine-count=2）。
• 依赖配置：确保环境包含transformers、datasets、deepspeed、torch等依赖。可使用Vertex AI官方PyTorch GPU镜像，并通过requirements.txt声明依赖：
  requirements.txt示例：

  transformers>=4.30.0
  datasets>=2.13.0
  deepspeed>=0.9.2
  torch>=2.0.0
  

步骤3：提交Vertex AI训练作业 #

使用gcloud命令提交作业，无需额外编写分布式启动命令，Hugging Face Trainer会结合DeepSpeed自动完成多节点进程初始化：

gcloud ai jobs submit training my-deepspeed-training-job \
    --region=us-central1 \
    --master-image-uri=gcr.io/deeplearning-platform-release/pytorch-gpu.2.0 \
    --machine-type=n1-standard-8 \
    --accelerator-type=NVIDIA_TESLA_T4 \
    --accelerator-count=1 \
    --machine-count=2 \
    --package-path=./training_package \
    --module-name=training_package.train \
    --job-dir=gs://my-bucket/training-jobs/deepspeed-demo \
    --python-version=3.10 \
    --requirements-file=./requirements.txt

若使用自定义镜像，替换--master-image-uri为你的镜像地址即可。

步骤4：验证训练状态 #

在Vertex AI控制台查看作业日志，若出现DeepSpeed初始化日志（如Zero stage 2 enabled、Number of processes: X），则说明多节点分布式训练已正常启动。

内容的提问来源于stack exchange，提问作者esdy