Pytorch基础指令

1. pytorch 指定使用GPU卡

##方法一：在运行脚本前，使用 CUDA_VISIBLE_DEVICES 环境变量
export CUDA_VISIBLE_DEVICES=0,1  # 使用 GPU 卡号为 0 和 1
python your_script.py

#或者
CUDA_VISIBLE_DEVICES=0,1 python your_script.py （常用方法）

## 方法二：使用torch.cuda.device
import torch

# 指定使用的 GPU 卡，例如使用 GPU 卡号为 0 和 1
torch.cuda.device(0)  # 使用第一个 GPU
torch.cuda.device(1)  # 使用第二个 GPU

## 方法三：在 PyTorch 中设置 CUDA_VISIBLE_DEVICES
import os
import torch

# 指定使用的 GPU 卡，例如使用 GPU 卡号为 0 和 1 （发现有的时候不生效）
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

# 在这之后的 PyTorch 操作将在指定的 GPU 卡上执行

2. 模型加载时指定device

# 通过.to(device)的形式，可以正常加载模型。device = "cuda:3"
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        trust_remote_code=True,
    ) .to(device)

# 在加载模型内部指定device=device，则报错 “RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'。
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        trust_remote_code=True,
        device= device    可修改为：device_map=device
    ) 
#

3. device_map

首先，检查num_gpus（即GPU数量）是否小于2。其次，检查变量device_map是否为None。device_map通常被用于定义模型如何在多个设备上分配，如果它为None，则表示模型应在单个设备上运行。通常设置device_map=2