horovod/horovod

镜像概述

本镜像提供一个分布式训练框架，专为TensorFlow、Keras、PyTorch和MXNet等主流深度学习框架设计。它旨在解决大规模模型训练中计算资源不足的问题，通过分布式架构实现多节点、多GPU协同训练，提升训练效率和模型扩展性。

核心功能和特性

框架兼容性

• 全面支持TensorFlow、Keras、PyTorch和MXNet四大深度学习框架
• 提供统一的分布式训练接口，降低跨框架使用门槛

分布式训练模式

• 支持数据并行（Data Parallelism）：将训练数据拆分到多个节点/GPU，并行计算梯度
• 支持模型并行（Model Parallelism）：将大型模型拆分到不同节点/GPU，协同完成前向/反向传播
• 支持混合并行模式：结合数据并行与模型并行，适应超大规模模型训练需求

通信与优化

• 集成NCCL、gRPC等高效通信后端，确保节点间数据传输低延迟
• 内置梯度同步优化策略，支持异步/同步训练模式切换
• 自动资源调度，动态适配不同硬件环境（CPU/GPU/TPU）

易用性与可扩展性

• 提供简洁的配置接口，无需深入分布式系统细节即可快速部署
• 支持动态增减训练节点，适应弹性计算场景
• 兼容主流容器编排平台（Kubernetes、Docker Swarm）

使用场景和适用范围

适用场景

• 大规模深度学习模型训练（如Transformer、ResNet等）
• 多GPU服务器或多节点集群环境下的分布式训练
• 需要跨框架统一分布式训练流程的场景
• 资源受限但需加速模型收敛的训练任务

适用用户

• 深度学习研究者与工程师
• 需部署分布式训练系统的企业团队
• 从事大规模模型开发的科研机构

使用方法和配置说明

基础使用（docker run）

bash
# 单节点启动（示例：使用PyTorch框架，本地模式）
docker run -it --name dist-train \
  -e FRAMEWORK=pytorch \
  -e TRAIN_MODE=local \
  -v /path/to/training/code:/app/code \
  -v /path/to/dataset:/app/data \
  distributed-training-framework

# 多节点启动（示例：2节点TensorFlow分布式训练）
# 节点1（主节点）
docker run -it --name dist-train-node1 \
  -e FRAMEWORK=tensorflow \
  -e TRAIN_MODE=distributed \
  -e ROLE=master \
  -e NUM_NODES=2 \
  -e MASTER_ADDR=192.168.1.100 \
  -e MASTER_PORT=29500 \
  -p 29500:29500 \
  -v /path/to/training/code:/app/code \
  -v /path/to/dataset:/app/data \
  distributed-training-framework

# 节点2（工作节点）
docker run -it --name dist-train-node2 \
  -e FRAMEWORK=tensorflow \
  -e TRAIN_MODE=distributed \
  -e ROLE=worker \
  -e NUM_NODES=2 \
  -e MASTER_ADDR=192.168.1.100 \
  -e MASTER_PORT=29500 \
  -v /path/to/training/code:/app/code \
  -v /path/to/dataset:/app/data \
  distributed-training-framework

环境变量配置

Docker Compose配置示例（2节点分布式训练）

yaml
version: '3'
services:
  master:
    image: distributed-training-framework
    container_name: dist-train-master
    environment:
      - FRAMEWORK=pytorch
      - TRAIN_MODE=distributed
      - ROLE=master
      - NUM_NODES=2
      - MASTER_ADDR=master
      - MASTER_PORT=29500
      - PARALLEL_MODE=data_parallel
    volumes:
      - ./training_code:/app/code
      - ./dataset:/app/data
    ports:
      - "29500:29500"

  worker:
    image: distributed-training-framework
    container_name: dist-train-worker
    environment:
      - FRAMEWORK=pytorch
      - TRAIN_MODE=distributed
      - ROLE=worker
      - NUM_NODES=2
      - MASTER_ADDR=master
      - MASTER_PORT=29500
      - PARALLEL_MODE=data_parallel
    volumes:
      - ./training_code:/app/code
      - ./dataset:/app/data
    depends_on:
      - master

训练代码集成说明

1. 将训练代码挂载至容器内/app/code目录
2. 代码中需通过框架原生接口或镜像提供的统一API初始化分布式环境，例如：
   • PyTorch: torch.distributed.init_process_group(backend="nccl")
   • TensorFlow: tf.distribute.MultiWorkerMirroredStrategy()
3. 启动容器后，框架将自动执行/app/code/train.py作为训练入口脚本（可通过ENTRYPOINT环境变量自定义入口）