[feat] refactored extension module (#5298)

* [feat] refactored extension module * polish * polish * polish * polish * polish * polish * polish * polish * polish * polish
2025-09-08 04:24:47 +00:00 · 2024-01-25 17:01:48 +08:00
parent d7f8db8e21
commit 7cfed5f076
157 changed files with 1353 additions and 8966 deletions
--- a/colossalai/legacy/utils/common.py
+++ b/colossalai/legacy/utils/common.py
@@ -96,9 +96,9 @@ def _calc_l2_norm(grads):
    global fused_optim

    if fused_optim is None:
-        from colossalai.kernel.op_builder import FusedOptimBuilder
+        from colossalai.kernel.kernel_loader import FusedOptimizerLoader

-        fused_optim = FusedOptimBuilder().load()
+        fused_optim = FusedOptimizerLoader().load()

    norm = 0.0
    if len(grads) > 0: