Win11 WSL环境下PyTorch T5模型双GPU使用问题

先确认WSL环境的GPU可用性 #

先确保WSL能正确识别两块GPU，执行nvidia-smi确认两块GPU都被列出且GPU 1处于空闲状态。同时在Python中验证PyTorch的GPU访问能力：

import torch
print(torch.cuda.is_available())  # 应输出True
print(torch.cuda.device_count())  # 应输出2

方法1：手动指定使用空闲的GPU 1 #

如果只是想先用空闲的GPU 1缓解内存问题，直接在代码中指定设备为cuda:1：

# 初始化模型后，将模型移至GPU 1
device = torch.device("cuda:1")
model = model.to(device)

# 训练时将输入数据也移至对应设备
for batch in dataloader:
    inputs = {k: v.to(device) for k, v in batch.items()}
    outputs = model(**inputs)

方法2：使用DataParallel实现双GPU并行 #

这是单机器多GPU场景下最简单的并行方式，无需复杂配置：

import torch.nn as nn

# 用DataParallel包装模型，指定要使用的GPU编号
model = nn.DataParallel(model, device_ids=[0, 1])
model = model.to("cuda")  # 默认主设备为cuda:0

# 训练流程无需额外修改，数据会自动拆分到两块GPU计算
for batch in dataloader:
    inputs = {k: v.to("cuda") for k, v in batch.items()}
    outputs = model(**inputs)

注：DataParallel会将模型主副本放在device_ids的第一个GPU上，再将数据拆分到所有指定GPU，计算后合并结果。

方法3：使用DistributedDataParallel（更高效的并行方案） #

针对大模型训练，DDP比DataParallel效率更高，适合长期训练场景：

1. 修改训练脚本：

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

# 初始化分布式进程组
dist.init_process_group(backend="nccl")
local_rank = dist.get_rank()
torch.cuda.set_device(local_rank)

# 数据加载使用DistributedSampler，确保每个进程拿到不同的数据分片
sampler = DistributedSampler(dataset)
dataloader = torch.utils.data.DataLoader(dataset, sampler=sampler, batch_size=your_batch_size)

# 模型移至对应GPU并包装DDP
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])

# 训练循环
for epoch in range(epochs):
    sampler.set_epoch(epoch)  # 保证不同epoch数据分片不同
    for batch in dataloader:
        inputs = {k: v.to(local_rank) for k, v in batch.items()}
        outputs = model(**inputs)
        # 损失计算、反向传播等流程与单GPU一致

# 训练结束后销毁进程组
dist.destroy_process_group()

2. 用torchrun启动脚本（自动分配多进程管理GPU）：

torchrun --nproc_per_node=2 your_train_script.py

额外优化建议 #

• 结合梯度累积：通过多次前向/反向传播再更新参数，降低单步内存占用。
• 模型量化：使用torch.quantization将模型权重转为低精度，进一步节省显存。
• 确认Windows与WSL的NVIDIA驱动版本一致，避免兼容性问题。

内容的提问来源于stack exchange，提问作者Ram