GPU 使用¶

申请 GPU¶

使用 --gres=gpu:<数量> 申请 GPU：

srun --input=none -N1 -n1 --gres=gpu:1 nvidia-smi

批处理脚本中：

#SBATCH --gres=gpu:1

Slurm 会为作业设置 CUDA_VISIBLE_DEVICES。程序应只看到分配到的 GPU。

srun --input=none -N1 -n1 --gres=gpu:1 bash -lc 'echo $CUDA_VISIBLE_DEVICES'

srun --input=none -N1 -n1 --gres=gpu:1 python - <<'PY'
import torch
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))
PY

多节点训练建议使用 torchrun、MPI、Ray 或 Dask，并通过 Slurm 分配节点。具体启动方式取决于训练框架。

基本原则：