我的机器上有4张H100卡,我现在只想用某一张卡跑程序,该如何设置。
代码里面设置
import os
# 记住要写在impot torch前
os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1"
命令行设置
export CUDA_VISIBLE_DEVICES=0,2 # Linux 环境
python test.py # 运行程序
docker 设置
docker run --gpus device=2 --shm-size 1024g
docker run --gpus all --shm-size 1024g
docker run --gpus device=1,2 --shm-size 1024g
拉取镜像
docker pull vllm/vllm-openai
备注,国内用户,可以使用代理源来加速镜像下载:
docker pull docker.1ms.run/vllm/vllm-openai
无法 import 某个package
import sys
sys.path.append('/your/home/package/')
docker设置代理