### 如何在Linux上使用PyTorch进行多GPU训练 为了实现多GPU训练，在启动脚本时可以利用`torch.distributed.launch`模块来简化分布式训练环境的设置[^1]。下面是一个典型的命令行调用方式： ```bash python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script) ``` 在这个命令中，`--nproc_per_node`参数指定了每台机器要使用的GPU数量；而`YOUR_TRAINING_SCRIPT.py`则是实际执行训练过程的Python文件名。 对于更复杂的场景或自定义配置，还可以采用另一种方法——通过编写专门的入口函数并手动初始化进程组来进行控制。这通常涉及到导入必要的库以及指定通信后端和初始方法等操作[^2]： ```python import torch import torch.nn as nn import torch.optim as optim import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP def setup(rank, world_size): # 初始化进程组 torch.distributed.init_process_group( backend='nccl', # 或者 'gloo' 如果不支持NCCL的话 init_method='env://', world_size=world_size, rank=rank ) def cleanup(): torch.distributed.destroy_process_group() class ToyModel(nn.Module): def __init__(self): super(ToyModel, self).__init__() ... def train_model(rank, size): setup(rank, size) model = ToyModel().to(rank) ddp_model = DDP(model, device_ids=[rank]) loss_fn = nn.MSELoss() optimizer = optim.SGD(ddp_model.parameters(), lr=0.001) ... cleanup() ``` 上述代码片段展示了如何创建一个简单的模型并通过DDP封装它以便于跨多个设备运行。需要注意的是，这里假设读者已经熟悉基本的PyTorch编程概念，并且环境中已正确安装了CUDA工具包和支持多GPU硬件设施[^3]。