如何在Amazon SageMaker上结合Hugging Face使用PyTorch DDP？

针对无法使用SageMaker Distributed Data Parallel（SMDistributed）的场景（如p3.8xlarge、g系列等实例），可以通过以下几种方式启动基于PyTorch DDP的Hugging Face Trainer脚本：

1. 手动使用torch.distributed.launch启动 #

在定义SageMaker训练作业时，直接指定自定义启动命令，用torch.distributed.launch初始化DDP环境：

• 命令示例（以4卡的p3.8xlarge为例）：

  python -m torch.distributed.launch --nproc_per_node=4 --use_env train.py
  

• Hugging Face Trainer会自动识别DDP环境变量，无需修改脚本即可启用DDP，替代默认的DP策略，解决内存占用不均的问题。

2. 通过SageMaker PyTorch Estimator配置DDP参数 #

使用SageMaker Python SDK时，给PyTorch Estimator添加distributed配置，让平台自动处理DDP的启动逻辑：

• 代码示例：

  from sagemaker.pytorch import PyTorch
  
  estimator = PyTorch(
      entry_point="train.py",
      instance_type="p3.8xlarge",
      instance_count=1,
      framework_version="2.0",
      py_version="py310",
      distributed={
          "distributed_training": True,
          "ddp": True
      }
  )
  estimator.fit()
  

• 该方式支持单实例多GPU和多实例多GPU场景，平台会自动配置节点通信、GPU进程数等参数。

3. 自定义训练容器配置DDP启动逻辑 #

如果需要更灵活的控制，可以自定义Docker容器，在容器的启动脚本中集成DDP启动命令：

• 推荐使用PyTorch 1.10+支持的torchrun（替代旧的launch工具），利用SageMaker提供的环境变量自动获取GPU数量：

  torchrun --nproc_per_node=$SM_NUM_GPUS train.py
  

• 容器中需确保安装对应版本的PyTorch和Hugging Face Transformers库，保证环境兼容性。

额外注意事项 #

• 避免在Trainer中手动指定device_map或冲突的分布式配置，Trainer会自动适配DDP环境。
• DDP会将模型参数均匀分配到各GPU，相比DP能大幅改善单GPU内存过载的问题，这也是解决跨实例扩展时内存不均的核心原因。

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