【Bug已解决】PyTorch multiprocessing error with Hogwild 解决方案
问题描述
在 PyTorch 中使用 Hogwild 训练模式（即多进程共享模型参数，不加锁异步更新）时，开发者经常遇到各种多进程相关的错误。Hogwild 是一种高效的并行训练方法，特别适用于稀疏梯度场景（如 NLP 中的词嵌入），但它的配置和使用有许多陷阱。
典型报错信息：
RuntimeError: Cannot re-initialize CUDA in forked subprocess.或者：
RuntimeError: DataLoader worker (pid(s) XXXX) exited unexpectedly或者：
RuntimeError: Unable to open shared memory lock或者更隐蔽的问题：多进程训练时模型参数不共享，各进程独立训练，最终结果不正确。
错误复现
复现场景一：CUDA fork 错误
import torch import torch.multiprocessing as mp import torch.nn as nn # 错误：在 fork 之前初始化了 CUDA model = nn.Linear(100, 10).cuda() # CUDA 已初始化 def train_worker(rank): # 子进程中尝试使用 CUDA model = model.cuda() # 报错！ # RuntimeError: Cannot re-initialize CUDA in forked subprocess mp.spawn(train_worker, nprocs=4)复现场景二：共享内存错误
import torch import torch.multiprocessing as mp import torch.nn as nn def train_worker(rank, model): # 子进程中尝试访问共享模型 optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # RuntimeError: Unable to open shared memory lock # 如果模型没有正确共享到共享内存 model = nn.Linear(100, 10) # 错误：没有调用 share_memory() mp.spawn(train_worker, args=(model,), nprocs=4)复现场景三：参数不共享
import torch import torch.multiprocessing as mp import torch.nn as nn def train_worker(rank, model): # 每个进程独立训练，参数不共享 optimizer = torch.optim.SGD(model.parameters(), lr=0.01) for i in range(100): x = torch.randn(32, 100) y = torch.randint(0, 10, (32,)) loss = nn.CrossEntropyLoss()(model(x), y) optimizer.zero_grad() loss.backward() optimizer.step() # 各进程的模型参数不同步！ model = nn.Linear(100, 10) # 没有共享内存，各进程有独立的模型副本 mp.spawn(train_worker, args=(model,), nprocs=4)根因分析
1. Hogwild 的原理
Hogwild 是一种无锁的并行随机梯度下降方法。核心思想：
- 多个进程/线程共享同一份模型参数
- 每个进程独立计算梯度并直接更新共享参数
- 不使用锁来同步更新
- 在稀疏梯度场景下（每次只更新少量参数），冲突概率很低
进程1 ──→ 计算梯度 ──→ 更新共享参数 ←── 计算梯度 ←── 进程2 进程3 ──→ 计算梯度 ──→ 更新共享参数 ←── 计算梯度 ←── 进程4 （无锁，异步更新）2. PyTorch 的共享内存机制
PyTorch 使用共享内存（shared memory）在进程间共享张量：
# 将张量移到共享内存 tensor.share_memory_() # nn.Module 的共享 model.share_memory()share_memory_()将模型的所有参数张量移到共享内存区域，使多个进程可以访问同一份参数。
3. CUDA fork 限制
CUDA 不支持在 fork 的子进程中重新初始化。当主进程已经初始化了 CUDA（例如调用了.cuda()或torch.cuda.is_available()），fork 出的子进程无法使用 CUDA。
解决方案：
- 使用
spawn启动方式（而不是fork） - 在子进程内部初始化 CUDA
- 使用
if __name__ == "__main__":保护主进程代码
4. 进程启动方式
Python 的multiprocessing支持三种启动方式：
fork: Unix 默认，复制父进程的内存空间 spawn: Windows 默认，重新启动 Python 解释器 forkserver: 新进程通过 forkserver 创建PyTorch 推荐使用spawn，因为它不继承父进程的 CUDA 状态。
解决方案
方案一：CPU 上的 Hogwild 训练
import torch import torch.multiprocessing as mp import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset def train_worker(rank, model, dataset, num_epochs, batch_size, learning_rate): """每个进程的训练函数""" # 设置随机种子（不同进程不同种子） torch.manual_seed(42 + rank) # 创建 DataLoader dataloader = DataLoader( dataset, batch_size=batch_size, shuffle=True, num_workers=0 # 子进程中不再创建子进程 ) # 创建优化器（共享模型参数） optimizer = optim.SGD(model.parameters(), lr=learning_rate) criterion = nn.CrossEntropyLoss() for epoch in range(num_epochs): total_loss = 0 num_batches = 0 for batch_x, batch_y in dataloader: # 前向传播 output = model(batch_x) loss = criterion(output, batch_y) # 反向传播（直接更新共享参数） optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() num_batches += 1 if rank == 0: # 只有主进程打印 print(f"Process {rank}, Epoch {epoch}, " f"Loss: {total_loss/num_batches:.4f}") def run_hogwild_training(): """运行 Hogwild 训练""" # 必须在 __main__ 中调用 # 创建模型并共享内存 model = nn.Sequential( nn.Linear(100, 256), nn.ReLU(), nn.Linear(256, 10) ) # 关键：将模型参数移到共享内存 model.share_memory() # 创建数据集 X = torch.randn(10000, 100) y = torch.randint(0, 10, (10000,)) dataset = TensorDataset(X, y) # 启动多进程训练 num_processes = 4 mp.spawn( train_worker, args=(model, dataset, 10, 32, 0.01), nprocs=num_processes, join=True ) # 训练完成后，model 的参数已被所有进程更新 print("Hogwild 训练完成！") return model if __name__ == "__main__": model = run_hogwild_training()方案二：使用 DistributedDataParallel 替代 Hogwild
import torch import torch.distributed as dist import torch.nn as nn import torch.optim as optim from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, TensorDataset from torch.utils.data.distributed import DistributedSampler import os def setup(rank, world_size): """初始化进程组""" os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '12355' dist.init_process_group("gloo", rank=rank, world_size=world_size) def cleanup(): """清理进程组""" dist.destroy_process_group() def train_worker(rank, world_size, dataset, num_epochs, batch_size): """DDP 训练函数""" setup(rank, world_size) # 创建模型并移到当前进程 model = nn.Sequential( nn.Linear(100, 256), nn.ReLU(), nn.Linear(256, 10) ) # 包装为 DDP model = DDP(model) # 分布式采样器 sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler) optimizer = optim.SGD(model.parameters(), lr=0.01) criterion = nn.CrossEntropyLoss() for epoch in range(num_epochs): sampler.set_epoch(epoch) # 确保每个 epoch 的数据打乱不同 total_loss = 0 num_batches = 0 for batch_x, batch_y in dataloader: output = model(batch_x) loss = criterion(output, batch_y) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() num_batches += 1 if rank == 0: print(f"Epoch {epoch}, Loss: {total_loss/num_batches:.4f}") cleanup() def run_ddp_training(): """运行 DDP 训练""" X = torch.randn(10000, 100) y = torch.randint(0, 10, (10000,)) dataset = TensorDataset(X, y) world_size = 4 mp.spawn( train_worker, args=(world_size, dataset, 10, 32), nprocs=world_size, join=True ) if __name__ == "__main__": run_ddp_training()方案三：使用线程替代进程
import torch import torch.nn as nn import torch.optim as optim from threading import Thread, Lock import queue from torch.utils.data import DataLoader, TensorDataset class HogwildThreadTrainer: """使用线程的 Hogwild 训练器（避免进程间通信开销）""" def __init__(self, model, dataset, num_threads=4, batch_size=32, lr=0.01): self.model = model self.dataset = dataset self.num_threads = num_threads self.batch_size = batch_size self.lr = lr self.criterion = nn.CrossEntropyLoss() # 创建数据队列 self.data_queue = queue.Queue() self.total_batches = 0 def _prepare_data(self): """准备数据到队列""" dataloader = DataLoader( self.dataset, batch_size=self.batch_size, shuffle=True, num_workers=0 ) for batch in dataloader: self.data_queue.put(batch) self.total_batches += 1 # 放入结束标记 for _ in range(self.num_threads): self.data_queue.put(None) def _train_thread(self, thread_id): """单个训练线程""" # 每个线程有自己的优化器，但共享模型参数 optimizer = optim.SGD(self.model.parameters(), lr=self.lr) while True: batch = self.data_queue.get() if batch is None: break batch_x, batch_y = batch # 前向传播 output = self.model(batch_x) loss = self.criterion(output, batch_y) # 反向传播（无锁更新） optimizer.zero_grad() loss.backward() optimizer.step() if thread_id == 0 and self.data_queue.qsize() % 100 == 0: print(f"Thread {thread_id}, Loss: {loss.item():.4f}, " f"Queue size: {self.data_queue.qsize()}") def train(self, num_epochs=10): """训练""" for epoch in range(num_epochs): self._prepare_data() threads = [] for i in range(self.num_threads): t = Thread(target=self._train_thread, args=(i,)) threads.append(t) t.start() for t in threads: t.join() print(f"Epoch {epoch} 完成") # 使用示例 if __name__ == "__main__": model = nn.Sequential( nn.Linear(100, 256), nn.ReLU(), nn.Linear(256, 10) ) X = torch.randn(10000, 100) y = torch.randint(0, 10, (10000,)) dataset = TensorDataset(X, y) trainer = HogwildThreadTrainer(model, dataset, num_threads=4, batch_size=32, lr=0.01) trainer.train(num_epochs=5)完整修复代码
""" PyTorch Hogwild 多进程训练完整方案 涵盖：CPU Hogwild、DDP、线程训练、错误处理 """ import torch import torch.nn as nn import torch.optim as optim import torch.multiprocessing as mp import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, TensorDataset, Dataset from torch.utils.data.distributed import DistributedSampler from threading import Thread import queue import os import time from typing import Optional, Callable # ==================== 模型定义 ==================== class SimpleModel(nn.Module): def __init__(self, input_dim=100, hidden_dim=256, num_classes=10): super().__init__() self.fc1 = nn.Linear(input_dim, hidden_dim) self.fc2 = nn.Linear(hidden_dim, num_classes) def forward(self, x): return self.fc2(torch.relu(self.fc1(x))) # ==================== 方案一：CPU Hogwild ==================== class HogwildTrainer: """CPU 上的 Hogwild 训练器""" def __init__(self, model, dataset, num_workers=4, batch_size=32, lr=0.01): self.model = model self.dataset = dataset self.num_workers = num_workers self.batch_size = batch_size self.lr = lr def _worker(self, rank, num_epochs, result_queue): """单个训练进程""" torch.manual_seed(42 + rank) # 每个进程创建独立的 DataLoader dataloader = DataLoader( self.dataset, batch_size=self.batch_size, shuffle=True, num_workers=0 ) optimizer = optim.SGD(self.model.parameters(), lr=self.lr) criterion = nn.CrossEntropyLoss() epoch_losses = [] for epoch in range(num_epochs): total_loss = 0 num_batches = 0 for batch_x, batch_y in dataloader: output = self.model(batch_x) loss = criterion(output, batch_y) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() num_batches += 1 avg_loss = total_loss / max(num_batches, 1) epoch_losses.append(avg_loss) if rank == 0: print(f" Epoch {epoch}, Avg Loss: {avg_loss:.4f}") result_queue.put((rank, epoch_losses)) def train(self, num_epochs=10): """启动 Hogwild 训练""" print(f"启动 Hogwild 训练: {self.num_workers} 个进程") # 关键：共享模型参数 self.model.share_memory() # 结果队列 ctx = mp.get_context("spawn") result_queue = ctx.Queue() processes = [] for rank in range(self.num_workers): p = ctx.Process( target=self._worker, args=(rank, num_epochs, result_queue) ) p.start() processes.append(p) # 等待所有进程完成 for p in processes: p.join() # 收集结果 results = [] while not result_queue.empty(): results.append(result_queue.get()) print("Hogwild 训练完成！") return results # ==================== 方案二：DDP 训练 ==================== class DDPTrainer: """DistributedDataParallel 训练器""" @staticmethod def _setup(rank, world_size): os.environ['MASTER_ADDR'] = 'localhost' os.environ['MASTER_PORT'] = '12355' dist.init_process_group("gloo", rank=rank, world_size=world_size) @staticmethod def _cleanup(): if dist.is_initialized(): dist.destroy_process_group() @staticmethod def _worker(rank, world_size, dataset, num_epochs, batch_size, lr): DDPTrainer._setup(rank, world_size) model = SimpleModel() model = DDP(model) sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler) optimizer = optim.SGD(model.parameters(), lr=lr) criterion = nn.CrossEntropyLoss() for epoch in range(num_epochs): sampler.set_epoch(epoch) total_loss = 0 num_batches = 0 for batch_x, batch_y in dataloader: output = model(batch_x) loss = criterion(output, batch_y) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() num_batches += 1 if rank == 0: print(f" DDP Epoch {epoch}, Loss: {total_loss/num_batches:.4f}") DDPTrainer._cleanup() @staticmethod def train(dataset, num_epochs=10, world_size=4, batch_size=32, lr=0.01): print(f"启动 DDP 训练: {world_size} 个进程") mp.spawn( DDPTrainer._worker, args=(world_size, dataset, num_epochs, batch_size, lr), nprocs=world_size, join=True ) print("DDP 训练完成！") # ==================== 方案三：线程训练 ==================== class ThreadHogwildTrainer: """使用线程的 Hogwild 训练器""" def __init__(self, model, dataset, num_threads=4, batch_size=32, lr=0.01): self.model = model self.dataset = dataset self.num_threads = num_threads self.batch_size = batch_size self.lr = lr self.criterion = nn.CrossEntropyLoss() def _train_thread(self, thread_id, data_queue, stats_dict): optimizer = optim.SGD(self.model.parameters(), lr=self.lr) batches_processed = 0 total_loss = 0 while True: batch = data_queue.get() if batch is None: break batch_x, batch_y = batch output = self.model(batch_x) loss = self.criterion(output, batch_y) optimizer.zero_grad() loss.backward() optimizer.step() batches_processed += 1 total_loss += loss.item() stats_dict[thread_id] = { "batches": batches_processed, "avg_loss": total_loss / max(batches_processed, 1) } def train(self, num_epochs=10): print(f"启动线程 Hogwild 训练: {self.num_threads} 个线程") for epoch in range(num_epochs): # 准备数据 dataloader = DataLoader( self.dataset, batch_size=self.batch_size, shuffle=True ) data_queue = queue.Queue() for batch in dataloader: data_queue.put(batch) for _ in range(self.num_threads): data_queue.put(None) # 启动线程 stats_dict = {} threads = [] for i in range(self.num_threads): t = Thread(target=self._train_thread, args=(i, data_queue, stats_dict)) t.start() threads.append(t) for t in threads: t.join() # 统计 total_batches = sum(s["batches"] for s in stats_dict.values()) avg_loss = sum(s["avg_loss"] * s["batches"] for s in stats_dict.values()) / max(total_batches, 1) print(f" Epoch {epoch}, Batches: {total_batches}, Avg Loss: {avg_loss:.4f}") print("线程训练完成！") # ==================== 演示 ==================== def demo_hogwild(): """演示 Hogwild 训练""" print("=" * 60) print("PyTorch Hogwild 多进程训练演示") print("=" * 60) # 创建数据 X = torch.randn(5000, 100) y = torch.randint(0, 10, (5000,)) dataset = TensorDataset(X, y) # 1. CPU Hogwild print("\n--- CPU Hogwild ---") model = SimpleModel() trainer = HogwildTrainer(model, dataset, num_workers=4, batch_size=32, lr=0.01) trainer.train(num_epochs=3) # 2. 线程训练 print("\n--- 线程 Hogwild ---") model2 = SimpleModel() trainer2 = ThreadHogwildTrainer(model2, dataset, num_threads=4, batch_size=32, lr=0.01) trainer2.train(num_epochs=3) if __name__ == "__main__": demo_hogwild()常见陷阱与注意事项
1. 必须在main中启动多进程
# 错误：不在 __main__ 中启动 def some_function(): mp.spawn(train, nprocs=4) # 可能导致递归创建进程 # 正确 if __name__ == "__main__": mp.spawn(train, nprocs=4)2. share_memory 必须在 spawn 之前调用
model = SimpleModel() model.share_memory() # 在 spawn 之前 mp.spawn(worker, args=(model,), nprocs=4)3. 子进程中不要创建子进程
def worker(rank): # 错误：在子进程中创建 DataLoader 的子进程 dataloader = DataLoader(dataset, num_workers=4) # 嵌套多进程！ # 正确：子进程中 num_workers=0 dataloader = DataLoader(dataset, num_workers=0)4. CUDA 和多进程
# CUDA 场景下不要使用 Hogwild，使用 DDP # Hogwild 主要用于 CPU 训练 # 如果必须用 CUDA + 多进程： # 1. 使用 spawn 启动方式 # 2. 在子进程内部初始化 CUDA # 3. 每个进程使用独立的 GPU5. 梯度累积与 Hogwild
# Hogwild 中梯度累积可能导致更多冲突 # 建议每个 step 直接更新参数 # 不要在 Hogwild 中使用梯度累积总结
PyTorch 的 Hogwild 多进程训练是一种高效的 CPU 并行训练方法，特别适用于稀疏梯度场景。其核心原理是多个进程共享同一份模型参数（通过共享内存），各自独立计算梯度并直接更新参数，不使用锁来同步。
使用 Hogwild 的关键要点：
- 共享内存：在启动多进程之前调用
model.share_memory()，将参数移到共享内存。 - spawn 启动：使用
mp.spawn或mp.get_context("spawn")启动进程，避免 CUDA fork 问题。 - 子进程限制：子进程中 DataLoader 的
num_workers必须为 0，避免嵌套多进程。 - 适用场景：Hogwild 适用于 CPU 训练和稀疏梯度（如 NLP 词嵌入）。对于 GPU 训练或稠密梯度，推荐使用 DDP。
对于现代 PyTorch 训练，DDP（DistributedDataParallel）通常是更好的选择——它支持 GPU、有梯度同步保证、且性能优秀。Hogwild 更适合特定的 CPU 场景或作为轻量级并行方案。
空间解法)
