extensions/cuda_extension.py · ffffc32dc71c767aac625279593a7ab339855645 · OpenDAS / ColossalAI

[checkpointio] fix gemini and hybrid parallel optim checkpoint (#5347) · ffffc32d

Hongxin Liu authored Feb 01, 2024

* [checkpointio] fix hybrid parallel optim checkpoint

* [extension] fix cuda extension

* [checkpointio] fix gemini optimizer checkpoint

* polish code

ffffc32d

cuda_extension.py 3.8 KB

Replace cuda_extension.py