多卡训练torchrun

<think>嗯，用户想将PyTorch的单卡训练代码改成多卡训练的。首先，我需要回忆一下PyTorch多卡训练的方法。通常，PyTorch支持两种主要的多卡训练方式：DataParallel（DP）和DistributedDataParallel（DDP）。DataParallel比较简单，但效率可能不如DDP，特别是在多机多卡的情况下。而DDP虽然复杂一些，但性能更好，适合大规模训练。 用户可能对这两种方法不太了解，需要解释清楚区别。可能他们需要的是快速上手，所以得先给出两种方法的步骤，然后比较优缺点。另外，用户可能在使用过程中遇到问题，比如GPU负载不均衡，或者多进程启动的问题，所以注意事项部分也很重要。 首先，对于DataParallel，修改起来确实简单，只需要用DataParallel包装模型，然后调整batch size。但要注意，DP有一个主卡，其他卡的计算结果会汇总到主卡，可能导致主卡内存不足。这时候需要提醒用户如果遇到内存问题，可能需要换用DDP。 然后是DDP的部分，DDP需要更多的设置，比如初始化进程组，配置本地rank，使用DistributedSampler，还有启动脚本。用户可能不熟悉多进程启动，这里需要给出具体的代码示例和启动命令，比如用torch.distributed.launch或者torchrun。 还要提到DDP中每个进程处理自己的数据分片，所以Sampler需要调整，避免数据重复。另外，保存模型的时候要注意只在主进程保存，否则可能会多个进程同时写文件，导致问题。 混合精度训练也是一个常见的需求，所以可以补充一下如何结合DDP和AMP，比如使用autocast和GradScaler。 用户可能关心代码修改的步骤是否清晰，所以需要分点说明，每个步骤的关键代码部分，让用户能跟着做。同时，注意事项部分需要提醒常见的坑，比如batch size的设置、学习率调整、GPU选择、版本兼容性等。 最后，总结一下两种方法的适用场景，帮助用户根据自身情况选择。比如，快速实验用DP，大规模训练用DDP。可能用户的需求是希望提升训练速度，但资源有限，所以需要给出实际的建议，比如如果卡不多，DP可能够用，但卡多的话DDP更好。 检查一下是否有遗漏的部分，比如是否需要修改优化器或者学习率调度器？在DDP中，优化器会自动处理梯度同步，所以用户不需要额外修改，这部分可能不需要提。但要注意模型参数只在主进程初始化，或者需要同步，可能需要用到torch.save的时候加上module属性。 另外，启动命令中的参数是否正确，比如--nproc_per_node指定每个节点的进程数，通常等于GPU数量。还有是否需要设置MASTER_ADDR和MASTER_PORT，特别是在多机情况下，但单机多卡可能不需要，所以示例中可以省略，但提一下如果有需要的情况。 确保代码示例正确，比如在DDP中，model = DDP(model, device_ids=[local_rank])，这个是否正确？或者是否应该将模型移到device之后再用DDP包装？可能需要调整顺序，先.to(device)，再包装。 总之，步骤要详细但不过于复杂，让用户能一步步操作，同时指出可能的错误和解决方法。这样用户即使遇到问题，也能根据注意事项进行排查。</think>将PyTorch单卡训练代码改为多卡训练需要以下核心步骤，分为`DataParallel`（简单但效率有限）和`DistributedDataParallel`（推荐的高效方法）两种方案： --- ### 方案一：DataParallel（适合快速实验） **优点**：代码改动小，适合快速验证 **缺点**：效率较低（主卡瓶颈），不支持多机 ```python import torch.nn as nn # 原单卡代码 model = MyModel().cuda() # 改为多卡（自动分配数据到所有可见GPU） model = nn.DataParallel(model) # 注意：batch_size需调整为单卡的n倍（n为GPU数量） ``` --- ### 方案二：DistributedDataParallel（DDP，生产环境推荐） **优点**：高效无瓶颈，支持多机多卡 **步骤**： 1. **初始化进程组** ```python import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backend='nccl', init_method='env://') local_rank = int(os.environ['LOCAL_RANK']) torch.cuda.set_device(local_rank) ``` 2. **模型包装** ```python model = MyModel().to(local_rank) model = DDP(model, device_ids=[local_rank]) ``` 3. **数据加载器改造** ```python from torch.utils.data.distributed import DistributedSampler train_sampler = DistributedSampler(dataset, shuffle=True) dataloader = DataLoader(dataset, batch_size=args.batch_size, sampler=train_sampler) ``` 4. **训练循环调整** ```python for epoch in range(epochs): dataloader.sampler.set_epoch(epoch) # 保证shuffle有效性 for x, y in dataloader: x, y = x.to(local_rank), y.to(local_rank) # 后续与单卡训练相同 ``` 5. **启动命令** ```bash # 单机4卡示例 python -m torch.distributed.launch --nproc_per_node=4 train.py # 或使用新版本torchrun（推荐） torchrun --nproc_per_node=4 train.py ``` --- ### 关键注意事项 1. **Batch Size设置**：总batch_size = 单卡batch_size × GPU数量 2. **学习率调整**：通常需要线性缩放（lr = base_lr × GPU数量） 3. **模型保存**：通过`model.module`获取原始模型对象 4. **日志处理**：只在`local_rank==0`的进程执行打印/保存操作 5. **混合精度训练**：需配合`torch.cuda.amp`使用 --- ### 常见问题解决 1. **GPU显存不均**：检查是否漏写`torch.cuda.set_device` 2. **进程死锁**：确保`DistributedSampler`正确设置 3. **NCCL错误**：添加`NCCL_DEBUG=INFO`环境变量调试 推荐优先使用DDP方案，虽然代码改动较多，但能获得更好的扩展性。实际修改时建议使用`torchrun`作为启动器（PyTorch 1.9+），它自动处理环境变量设置。

### 使用 DeepSpeed 进行单机多卡训练 ChatGLM 模型 #### 配置文件设置 对于使用 DeepSpeed 训练大型语言模型（LLMs），配置文件至关重要。通常情况下，DeepSpeed 的配置通过 JSON 文件指定。下面是一个适用于 ChatGLM 模型的典型 DeepSpeed 配置文件 `ds_config.json`： ```json { "fp16": { "enabled": true, "loss_scale": 0, "initial_scale_power": 32, "loss_scale_window": 1000, "hysteresis": 2, "min_loss_scale": 1 }, "optimizer": { "type": "AdamW", "params": { "lr": 5e-5, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.01 } }, "scheduler": { "type": "WarmupLR", "params": { "warmup_min_lr": 0, "warmup_max_lr": 5e-5, "warmup_num_steps": 500 } }, "zero_optimization": { "stage": 2, "allgather_partitions": true, "allgather_bucket_size": 2e8, "reduce_scatter": true, "reduce_bucket_size": 2e8, "overlap_comm": true, "load_from_fp32_weights": true, "cpu_offload": false }, "gradient_accumulation_steps": 1, "steps_per_print": 2000, "wall_clock_breakdown": false } ``` 此配置启用了 FP16 和 ZeRO stage 2 来优化内存效率和加速训练过程[^1]。 #### 启动命令 要利用 DeepSpeed 执行单机多卡训练，可以采用如下 bash 脚本启动训练任务。假设已经准备好了上述提到的 `ds_config.json` 文件以及必要的 Python 环境，则可以通过以下方式启动训练： ```bash #!/bin/bash set -ex export PYTHONPATH=$PYTHONPATH:/path/to/your/project/ export DEEPSPEED=1 torchrun \ --nnodes=1 \ --nproc_per_node=4 \ # 假设每台机器有四个 GPU 可用 /path/to/fine_tune.py \ --deepspeed \ --deepspeed_config=/path/to/ds_config.json \ --model_name_or_path=chatglm-model-path \ --output_dir=output_directory \ --train_file=train_data_path \ --validation_file=val_data_path \ --do_train \ --per_device_train_batch_size=4 \ --learning_rate=5e-5 \ --num_train_epochs=3 \ --save_strategy="epoch" ``` 这段脚本同样遵循了良好的实践习惯，在遇到错误时会自动终止，并且显示执行过的每一行命令以便于调试[^3]。 #### 常见问题解决方法 当尝试使用 DeepSpeed 对 LLM 如 ChatGLM 进行微调时可能会碰到一些挑战或障碍。以下是几个常见的解决方案： - **显存不足**：如果发现显存在训练过程中耗尽，考虑调整 batch size 或者启用更高级别的 Zero Redundancy Optimizer (ZeRO)，即增加 `"zero_optimization.stage"` 参数至更高的数值。 - **性能瓶颈**：若观察到 I/O 成为性能瓶颈，可探索异步加载数据集的方法或是提高读取速度的数据预处理技术；另外也可以适当减少日志记录频率以减轻磁盘压力。 - **通信延迟过高**：针对跨多个 GPU 协同工作的场景下可能出现较高的网络通讯开销问题，建议检查集群内部网路状况确保其稳定高效运作的同时合理规划节点间连接拓扑结构。