[feat] refactored extension module (#5298)

* [feat] refactored extension module * polish * polish * polish * polish * polish * polish * polish * polish * polish * polish
2025-09-15 06:00:07 +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/nn/optimizer/fused_sgd.py
+++ b/colossalai/nn/optimizer/fused_sgd.py
@@ -72,9 +72,9 @@ class FusedSGD(Optimizer):
        self.wd_after_momentum = wd_after_momentum

        if multi_tensor_applier.available:
-            from colossalai.kernel.op_builder import FusedOptimBuilder
+            from colossalai.kernel.kernel_loader import FusedOptimizerLoader

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

            # Skip buffer
            self._dummy_overflow_buf = torch.tensor(