It seems that the FSDP2 integration crashed multiple things in the checkpoint conversion. The FSDP2 integration PR is already pretty stable and we can refactor the checkpoint conversion accordingly: https://github.com/Modalities/modalities/pull/312
