StopIteration: Caught StopIteration in replica 0 on device 0. 问题排查与解决

Yore_999

已于 2022-06-18 15:05:42 修改

阅读量2.1k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： Pytorch 文章标签： pytorch 人工智能深度学习

于 2022-06-18 15:05:01 首次发布

本文链接：https://blog.csdn.net/qq_35284646/article/details/125346731

Pytorch 专栏收录该内容

4 篇文章

订阅专栏

博主在升级PyTorch至1.11后遇到TransformerXL训练错误，通过排查发现是数据精度不一致导致。通过将数据类型强制转换为torch.float32并修改内存初始化代码，成功解决了跨GPU训练问题。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

首先是错误内容截图：（抱歉因为打码有点糊）

我在训练修改后的TransformerXL时，发现了如上的错误，此前代码已经成功地在单GPU下运行过，切换到多卡运行出现该问题。尝试进行解决。

使用的环境是： Pytorch1.11 transformers：4.18

在网上进行查阅后大部分人都说可能是pytorch版本的问题，当前所使用的pytorch版本过高，需要降级到1.4.0版本。

降级听起来比较简单，但是我不想降级到太低的版本，只能走第二条路，修改代码。

首先定位到出错的非源码的最后一行，

param = next(self.parameters())

经过上网查找，发现可能是在训练过程中部分数据的精度不同导致的问题，可能同时存在16位精度和32位精度的数据，尝试在这里进行修改，将其直接指定为torch.float32 进行训练。

原始代码为：

    def init_mems(self):
        if self.mem_len > 0:
            mems = []
            param = next(self.parameters())
            for i in range(self.n_layer+1):
                empty = torch.empty(0, dtype=param.dtype, device=param.device)
                mems.append(empty)
            return mems
        else:
            return None

更改后的代码是：

    def init_mems(self):
        if self.mem_len > 0:
            mems = []
            for i in range(self.n_layer+1):
                empty = torch.empty(0, dtype=torch.float32).cuda()
                mems.append(empty)
            return mems
        else:
            return None

成功！问题解决！