哈喽，请问是否尝试过使用ZeRO3进行视频图像的混合训练？ #126
Description
使用DeepSpeed ZeRO3训练模型时，如果在batch（per_device_bs>1)内同时存在image和video的话，我发现程序会挂起，弹出一条非logger日志，类似于“Invalidate trace cache @ step 3: expected module 2, but got module 1”，当然也可能没有任何消息弹出，取决于DeepSpeed版本。
查了一下相关的DeepSpeed Issue，发现大概是指ZeRO3会做一个关于前向传播和反向传播的trace cache，但是同时存在video和image的情况下，self.visual在不同的GPU上，有可能有的是是调用一次self.visual.forward，有的是两次，取决于micro_batch内部的构成，这样可能会破坏ZeRO3所谓的“trace cache”（我也并不是特别了解，只是一个大概结论），然后可能GPU之间就不同步了，程序挂起。
ZeRO2/per_device_bs=1 with ZeRO3都不会出现这个问题，都是可以正常训练的。
猜测两个分支合二为一的话或许没有这个问题了？还是我ZeRO3的打开方式不正确？