在AWS SageMaker多GPU节点上如何配置torchrun作为容器入口？

问题说明 #

• 在AWS SageMaker的ml.p4d.24xlarge实例上运行训练任务时，出现训练时间大幅变慢的问题，表现与transformers仓库中编号28916的已知问题一致。
• 打算用torchrun解决该问题，但不想依赖SageMaker提供的HuggingFace或PyTorch estimator，希望自定义容器配置。
• 当前Dockerfile的入口设置为：

  ENTRYPOINT ["python3", "/opt/program/entrypoint.py"]
  

• 疑问：是否可以直接修改为以下形式？

  ENTRYPOINT ["torchrun --nproc_per_node 8", "/opt/program/entrypoint.py"]
  

正确修改方式 #

你给出的写法存在问题，Docker的ENTRYPOINT数组格式（exec形式）要求每个命令和参数都作为独立的数组元素，不能把torchrun和它的参数合并成一个字符串。正确的配置应该是：

ENTRYPOINT ["torchrun", "--nproc_per_node", "8", "/opt/program/entrypoint.py"]

额外注意事项 #

• ml.p4d.24xlarge实例配备8个GPU，--nproc_per_node 8的设置刚好让每个GPU对应一个训练进程，能充分利用硬件资源。
• 确保你的容器镜像中已安装合适版本的PyTorch，torchrun是PyTorch自带的工具，版本过低可能没有该命令。

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