使用 Huggingface Trainer 时如何指定使用哪个 GPU

Question

HuggingFace 提供如下的training_args。当我使用 HF trainer 训练模型时，我发现默认使用 cuda:0。

我浏览了 HuggingFace 文档，但仍然不知道在使用 HF 训练器时如何指定在哪个 GPU 上运行。

training_args = TrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=3,              # total # of training epochs
    per_device_train_batch_size=16,  # batch size per device during training
    per_device_eval_batch_size=64,   # batch size for evaluation
    warmup_steps=500,                # number of warmup steps for learning rate scheduler
    weight_decay=0.01,               # strength of weight decay
    logging_dir='./logs',            # directory for storing logs
)

Answer 1

控制使用哪个 GPU 的最常见和实用的方法是设置 CUDA_VISIBLE_DEVICES 环境变量。

如果你想在运行Python脚本时在命令行中使用这个选项，你可以这样做：

CUDA_VISIBLE_DEVICES=1 python train.py

或者，您可以在导入 PyTorch 或任何其他基于 CUDA 的库（如 HuggingFace Transformers）之前插入此代码：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"  # or "0,1" for multiple GPUs

这样，无论您的计算机上有多少个 GPU，Hugging Face Trainer 都只能查看和使用您指定的 GPU。