多GPU模型并行拆分求助：遇「All tensors must be on same GPU」报错

问题根源 #

你同时使用了nn.DataParallel和手动模型并行逻辑，两者冲突；且forward中未手动将张量迁移到对应Block的GPU设备，导致张量与模型层设备不匹配，触发「All tensors must be on same GPU」错误。

解决方案步骤 #

1. 移除nn.DataParallel包装，改用nn.ModuleList存储Blocks，方便单独分配GPU
2. 手动将每个Block分配到不同GPU设备
3. 在forward中逐块处理，每次将张量迁移到当前Block的GPU后再计算

修改后的完整代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F

# 假设这些是你的全局参数，根据实际情况调整
vocab_size = 50257
n_embd = 768
block_size = 1024
n_layer = 6
n_head = 12
devices = [torch.device("cuda:0"), torch.device("cuda:1")]  # 至少2个GPU

class Block(nn.Module):
    # 补充Block基础实现，你可替换为自己的版本
    def __init__(self, n_embd, n_head):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = nn.MultiheadAttention(n_embd, n_head, batch_first=True)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd)
        )
    
    def forward(self, x):
        x = x + self.attn(self.ln1(x), self.ln1(x), self.ln1(x))[0]
        x = x + self.mlp(self.ln2(x))
        return x

class LanguageModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd).to(devices[0])
        self.position_embedding_table = nn.Embedding(block_size, n_embd).to(devices[0])
        # 用ModuleList存储Blocks，移除DataParallel
        self.blocks = nn.ModuleList([Block(n_embd, n_head) for _ in range(n_layer)])
        # 手动将每个Block分配到不同GPU
        for i, block in enumerate(self.blocks):
            block.to(devices[i % len(devices)])
        self.ln_f = nn.LayerNorm(n_embd).to(devices[0])
        self.lm_head = nn.Linear(n_embd, vocab_size).to(devices[0])

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        # 将输入idx移到embedding层所在的GPU（devices[0]）
        idx = idx.to(devices[0])
        tok_emb = self.token_embedding_table(idx)  # (B,T,C)
        pos_emb = self.position_embedding_table(torch.arange(T, device=devices[0]))  # (T,C)
        x = tok_emb + pos_emb  # (B,T,C)

        # 逐块处理，手动迁移张量到对应Block的GPU
        for block in self.blocks:
            x = x.to(block.device)  # 把输入移到当前Block的GPU
            x = block(x)
        
        # 将Blocks的输出移回主GPU（devices[0]）做后续处理
        x = x.to(devices[0])
        x = self.ln_f(x)  # (B,T,C)
        logits = self.lm_head(x)  # (B,T,vocab_size)

        loss = None
        if targets is not None:
            # 确保targets和logits在同一GPU
            targets = targets.to(devices[0])
            B, T, C = logits.shape
            logits = logits.view(B*T, C)
            targets = targets.view(B*T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss

    def generate(self, idx, max_new_tokens):
        idx = idx.to(devices[0])
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -block_size:]
            logits, loss = self(idx_cond)
            logits = logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

关键修改说明 #

• 移除DataParallel：nn.DataParallel是数据并行逻辑，会将整个模型复制到每个GPU，与手动拆分模型到不同GPU的逻辑完全冲突，必须移除。
• Block设备分配：通过ModuleList单独管理每个Block，循环将Block分配到不同GPU，确保权重分散在多卡。
• 张量手动迁移：在forward中，每次处理Block前将输入张量移到该Block所在的GPU，处理完成后再传递给下一个Block；最后将结果移回主GPU处理后续层（LayerNorm、分类头）。
• 输入/目标张量对齐：确保输入idx和targets都迁移到对应GPU，避免设备不匹配。

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