[BUG]: [extension] Compiling or loading the JIT-built cpu_adam kernel during runtime now

[HaixHan]

🐛 Describe the bug

when i use the "colossalai_zero" strategy to train the RM model, it will spend a lot of time to load optimizer . I am very puzzled! i ran the RM model successfully yesterday, and using the same code with less time. Who can tell me how to solve this? thanks!

Environment

No response

[HaixHan]

运行到这里，一直卡着不往下进行下去了。我怀疑是optimizer的问题

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

Running here, it has been stuck and can't go on. I suspect it is the problem of optimizer

[JThh]

Would this issue #3496 be useful?

[ZXMMD]

How to solve it? Thanks