1、查看cuda相关
1)cuda是否可用
torch.cuda.is_available()
2)gpu数量
torch.cuda.device_count()
3)gpu名字,默认从0开始
torch.cuda.get_device_name(0)
4)当前设备索引
torch.cuda.current_device()
2、多卡训练
2)
3、设置默认device
1)torch.cuda.set_device(device_ids)
2)os.environ["CUDA_VISIBLE_DEVICES"] = device_ids
【注意】
1)当已经import torch并进行了初始化之后,2可能无效,导致模型在其他卡上训练
2)torch.load(checkpoint),如果checkpoint中的tensor是gpu tensor,则会加载到0卡上;
对于1方法来说,是绝对0卡gpu: 0;
对于2方法来说,是相对0卡gpu: visible_start;
最好是存储checkpoint时,存为cpu tensor;或者load时,设置map_location=cpu
3)不推荐使用model.to(device)来指定模型gpu,因为不设置全局的默认设备的话,torch.cuda.empty_cache()操作可能会出现问题
4、GPU显存限制
torch.cuda.set_per_process_memory_fraction — PyTorch 1.12 documentation
混合精度
模型流水线(deep speed zero 3、accelerator)
grad checkpointing