提交分布式训练任务时，系统会构建容器计算环境和设置相应的环境变量，该章节用于介绍常见的环境变量，用户也可根据实际训练任务自定义环境变量。

基本环境变量

  • 适用于所有分布式训练环境

    变量名 值（示例） 说明 用法

    TENSORBOARD_LOG_PATH

    /tensorboard_logs

    TensorBoard 日志存放路径

    指定 TensorBoard 的日志文件存储路径，用于训练过程中的可视化。

  • 适用于使用 NVIDIA GPU 的分布式训练环境

    变量名 值（示例） 说明 用法

    NVIDIA_VISIBLE_DEVICES

    GPU-b9807d54-6689-a9b5-3fd3-79e733fa28c5

    可见的 GPU 设备

    指定 PyTorch 训练过程中可使用的 GPU 设备。

    NVIDIA_REQUIRE_CUDA

    cuda>=12.3 brand=…​

    CUDA 版本要求和兼容的 GPU 品牌及驱动版本

    用于指定 PyTorch 训练所需的 CUDA 版本及兼容的 GPU 类型和驱动版本。

    NVIDIA_DRIVER_CAPABILITIES

    compute，utility

    GPU 驱动的功能

    指定 GPU 驱动需要启用的功能，以确保计算任务正常进行。

    LD_LIBRARY_PATH

    /usr/local/nvidia/lib:/usr/local/nvidia/lib64

    动态库路径

    设置动态链接库的搜索路径，确保 PyTorch 能找到所需的 CUDA 库。

    CUDA_VERSION

    12.3.1

    CUDA 版本

    指定当前使用的 CUDA 版本。

    NV_CUDA_CUDART_VERSION

    12.3.101-1

    CUDA Runtime 版本

    指定当前使用的 CUDA Runtime 版本。

    NV_CUDA_COMPAT_PACKAGE

    cuda-compat-12-3

    CUDA 兼容包

    指定当前使用的 CUDA 兼容包版本。

    NVARCH

    x86_64

    架构类型

    指定当前使用的硬件架构。

MPI 分布式训练环境变量

变量名 值（示例） 说明 用法

OMPI_MCA_orte_default_hostfile

/etc/mpi/hostfile

指定默认的主机文件路径，用于 Open MPI 作业。

配置 Open MPI 以使用指定的主机文件定义计算节点。

OMPI_MCA_plm_rsh_agent

/etc/mpi/kubexec.sh

指定用于远程启动 MPI 进程的代理脚本。

配置 Open MPI 以使用特定脚本远程启动 MPI 进程。

PyTorch 分布式训练环境变量

针对 PyTorch 分布式训练存在以下变量可使用，禁止变更。

变量名 值（示例） 说明 用法

PET_NNODES

1

PyTorch 作业中使用的节点数

用于多节点分布式训练节点。

PET_NODE_RANK

0

当前节点的排名

用于多节点分布式训练中的节点识别。

MASTER_PORT

23456

主节点上的端口号

设置分布式训练中的主节点端口，用于节点间通信。

PET_MASTER_PORT

23456

PyTorch 主节点的端口号

和 MASTER_PORT 一样，用于设置分布式训练中的主节点端口。

WORLD_SIZE

1

总进程数（节点数与每节点进程数的乘积）

用于分布式训练中定义全局进程数，决定总的并行度。

PET_NPROC_PER_NODE

auto

每节点的进程数

自动配置每个节点的进程数，用于分布式训练。

PET_MASTER_ADDR

tn-dp4iad3q4vsw-master-0

分布式训练中主节点的地址

设置主节点的 IP 地址或主机名，其他节点将连接到该地址以进行通信。

RANK

0

当前进程在分布式训练中的全局排名

用于识别分布式训练中的每个进程。

TENSORFLOW 分布式训练环境变量

变量名 值（示例） 说明 用法

TF_CONFIG

{"cluster":{"worker":["tn-dp8mcg9lv11c-worker-0.usr-rticrzez.svc:2222","tn-dp8mcg9lv11c-worker-1.usr-rticrzez.svc:2222"]},"task":{"type":"worker","index":0},"environment":"cloud"}

TensorFlow 集群配置。包含集群角色信息和任务信息。

在分布式 TensorFlow 作业中，TF_CONFIG 用于配置集群的作业角色、任务索引以及环境类型。