How to load unsharded pre-trained weights into hybrid parallel model?

[ShinoharaHare]

I refer to this example to build my own hybrid parallel model (TP + ZeRO + PP) for bloom and it seems to work well so far.

But now i want to initialize the model with the pre-trained weights which is unsharded, the thing is that i have no idea what slices of the unsharded weights is corespond to the sharded weights.

I wonder if there is a simple solution for this intention, or i need to convert it by myself.

I would appreciate it if someone could help.

Edit:
I think this is similar to this issue #2770.

[kurisusnowdeng]

Hi @ShinoharaHare , could you please try colossalai.utils.checkpoint.load_checkpoint. BTW, you shall use colossalai.utils.checkpoint.save_checkpoint to save gathered weights.

[ShinoharaHare]

Hi, @kurisusnowdeng
不好意思，我想您應該會說中文，我就直接以中文繼續提問了

關於 colossalai.utils.checkpoint.load_checkpoint，雖然還沒測試過，但我覺得應該是無法使用的，因為我說的「unsharded pre-trained weights」指的是從這裡下載的原始權重，它的格式應該是不同於 CAI 的吧？