Transformers 文档
GPU 选择
并获取增强的文档体验
开始使用
GPU 选择
在分布式训练期间，您可以指定要使用的 GPU 数量以及顺序。当您拥有计算能力不同的 GPU，并且想要首先使用更快的 GPU 时，这会很有用。或者您也可以仅使用可用 GPU 的子集。选择过程适用于 DistributedDataParallel 和 DataParallel。您不需要 Accelerate 或 DeepSpeed 集成。
本指南将向您展示如何选择要使用的 GPU 数量以及使用顺序。
GPU 数量
例如，如果有 4 个 GPU，而您只想使用前 2 个，请运行以下命令。
使用 --nproc_per_node 选择要使用多少 GPU。
torchrun --nproc_per_node=2 trainer-program.py ...
GPU 顺序
要选择要使用的特定 GPU 及其顺序，请配置 CUDA_VISIBLE_DEVICES 环境变量。最简单的方法是在 ~/bashrc 或其他启动配置文件中设置环境变量。CUDA_VISIBLE_DEVICES 用于映射要使用的 GPU。例如，如果有 4 个 GPU（0、1、2、3），而您只想运行 GPU 0 和 2
CUDA_VISIBLE_DEVICES=0,2 torchrun trainer-program.py ...
只有 2 个物理 GPU（0 和 2）对 PyTorch 是“可见的”，并且它们分别映射到 cuda:0 和 cuda:1。您也可以颠倒 GPU 的顺序，首先使用 2。映射变为 GPU 0 的 cuda:1 和 GPU 2 的 cuda:0。
CUDA_VISIBLE_DEVICES=2,0 torchrun trainer-program.py ...
您还可以将 CUDA_VISIBLE_DEVICES 环境变量设置为空值，以创建没有 GPU 的环境。
CUDA_VISIBLE_DEVICES= python trainer-program.py ...
与任何环境变量一样，它们可以被导出，而不是添加到命令行中。但是，不建议这样做，因为如果您忘记环境变量是如何设置的，最终使用了错误的 GPU，这可能会造成混淆。相反，常见的做法是在同一命令行上为特定的训练运行设置环境变量。
CUDA_DEVICE_ORDER 是您可以使用的另一个环境变量，用于控制 GPU 的排序方式。您可以按照以下方式排序。
- PCIe 总线 ID，与 NVIDIA 的
nvidia-smi和 AMD GPU 的rocm-smi的顺序匹配。
export CUDA_DEVICE_ORDER=PCI_BUS_ID- GPU 计算能力。
export CUDA_DEVICE_ORDER=FASTEST_FIRST如果您的训练设置包含较旧和较新的 GPU，其中较旧的 GPU 首先出现，但您无法物理交换卡以使较新的 GPU 首先出现，则 CUDA_DEVICE_ORDER 尤其有用。在这种情况下，设置 CUDA_DEVICE_ORDER=FASTEST_FIRST 始终首先使用更新更快的 GPU（nvidia-smi 或 rocm-smi 仍然按照其 PCIe 顺序报告 GPU）。或者您也可以设置 export CUDA_VISIBLE_DEVICES=1,0。