脚本内指定GPU部署

CUDA设备配置

设置脚本可见的GPU卡号

在脚本开头添加以下代码，以指定脚本可见的GPU卡号范围：

import os

# 设置环境变量，只显示4, 5, 6, 7号GPU
os.environ['CUDA_VISIBLE_DEVICES'] = '4,5,6,7'

注意事项

• 经过上述设置后，对于脚本而言，这四张卡的编号将会是 "0, 1, 2, 3"，即根据可见卡重新排列了索引。
• 在此脚本内的任何GPU部署都不会影响到其他卡，因为此脚本看不到其他卡了。

模型加载示例

以下是使用指定设备映射加载模型的示例代码：

from transformers import AutoModelForCausalLM

# 模型路径
model_path = 'path_to_your_model'

# 加载模型，指定数据类型、内存使用优化、信任远程代码，并自动设置设备映射
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
    device_map="auto"
)

device_map 参数可以自由设置，不会影响其他人的卡。

附赠

查看GPU进程全名

nvidia-smi --query-compute-apps=process\_name --format=csv

查看GPU信息

nvidia-smi