事件起因:

我需要微调一个多模态大模型，为了压缩图片输入的token数量，防止上下文数量过长导致训练时间和效果变差，在多模态大模型的Vision Encoder后加入一个模块用于把每张图片的token数压缩至32个.

奇妙的bug

我加入的模块大概长这样:

python
def apply_compress(self, features):
    """
    features: [image_size,original_token,hidden_dim]
    """
    new_image_feature = []
    for i in range(image_size):
        new_image_feature.append(compress(features[i]))
    
    return new_image_feature

省去了很多细节，在训练的时候，奇怪的事情发生了: 在单卡上训练速度很正常，但一旦我切换到双卡, 训练在第一个batch都没完成的情况下就卡死了(batch_size=1)

经过排查，我发现rank0和rank1都执行到了这个函数, rank0的image_size为47, rank1的image_size为24. 在rank1处理完所有的图片后, rank0和rank1一起卡死.

问题解决

我一开始以为是compress函数有问题(但其实它是一个库函数), 结果发现我对torch的多卡运算有误,下面直接贴出AI的解释:

您好！从您提供的详尽代码和问题描述来看，这是一个非常典型的分布式训练（DDP, Distributed Data Parallelism）中遇到的 死锁（Deadlock）或挂起（Hang） 问题。

根本原因在于：在分布式训练中，每个进程（GPU/rank）的计算图和执行路径必须完全一致。而您当前的代码和数据加载方式破坏了这个核心原则。

问题根源：执行路径不一致导致同步失败

在 PyTorch DDP 中，forward 传播过程会记录模型中每个参数参与的运算。在 backward 传播时，DDP 会为每个参数注册一个钩子（hook），当该参数的梯度计算完成后，DDP 会立即启动一个异步的 AllReduce 操作，将所有进程上的这个梯度进行求和平均，从而保证所有进程的梯度一致。

这个机制能正常工作的前提是：所有进程必须以完全相同的顺序、对完全相同的参数集计算梯度。如果任何一个进程的计算路径与其他进程不同，就会导致：

某个进程期望收到一个它永远不会收到的梯度（因为它没有执行相应的计算）。
或者，进程之间 AllReduce 的顺序错乱。

这两种情况都会导致某些进程无限期地等待其他进程，从而造成整个训练过程的挂起。

1. 罪魁祸首：不同数量的图像 (47 vs 24)

您描述的场景是问题的核心：

Rank 0 加载了包含 47 个图像的样本。
Rank 1 加载了包含 24 个图像的样本。
执行次数不同：Rank 0 的 for 循环会执行 47 次，而 Rank 1 只会执行 24 次。
计算图不同：这意味着 self.compress 这个模块（它是一个拥有独立参数的神经网络）在 Rank 0 上被调用了 47 次，而在 Rank 1 上只被调用了 24 次。
DDP 崩溃：当反向传播开始时，Rank 0 会计算 47 次 compress 的梯度，并准备进行 47 次相应的 AllReduce 同步。而 Rank 1 只会计算 24 次。当 Rank 1 完成它自己的 24 次计算和同步后，它就会进入等待状态。与此同时，Rank 0 在完成了 24 次同步后，会继续尝试发起第 25 次 AllReduce，但此时它永远等不到 Rank 1 的响应，因为 Rank 1 的计算图里已经没有这个操作了。于是，Rank 0 就被永久地卡住了。

您遇到的问题是 DDP 的经典陷阱。其核心是由于数据不均匀（不同数量的图像）导致了模型中带参数模块（PerceiverResampler）的调用次数不一致，破坏了 DDP 要求的“计算图一致性”原则，从而在梯度同步时发生死锁。

最直接的解决方案是：修改您的数据加载和预处理流程，确保送入模型的每个 batch 在所有 GPU 上的结构（尤其是可变长度的元素，如图像数量）都是统一的，通常通过填充（Padding）来实现。

后续的拓展