[builder] unified cpu_optim fused_optim inferface (#2190)

2025-09-08 12:30:42 +00:00 · 2022-12-23 20:57:41 +08:00
parent 9587b080ba
commit 355ffb386e
9 changed files with 28 additions and 50 deletions
--- a/colossalai/nn/optimizer/fused_sgd.py
+++ b/colossalai/nn/optimizer/fused_sgd.py
@@ -80,11 +80,8 @@ class FusedSGD(Optimizer):
        self.wd_after_momentum = wd_after_momentum

        if multi_tensor_applier.available:
-            try:
-                from colossalai._C import fused_optim
-            except:
-                from colossalai.kernel.op_builder import FusedOptimBuilder
-                fused_optim = FusedOptimBuilder().load()
+            from colossalai.kernel import fused_optim
+
            # Skip buffer
            self._dummy_overflow_buf = torch.tensor([0],
                                                    dtype=torch.int,