[kernel] move all symlinks of kernel to colossalai._C (#1971)

2025-09-12 20:54:35 +00:00 · 2022-11-17 13:42:33 +08:00
parent 7e24b9b9ee
commit f8a7148dec
27 changed files with 463 additions and 322 deletions
--- a/colossalai/utils/multi_tensor_apply/multi_tensor_apply.py
+++ b/colossalai/utils/multi_tensor_apply/multi_tensor_apply.py
@@ -14,7 +14,7 @@ class MultiTensorApply(object):

    def __init__(self, chunk_size):
        try:
-            import colossal_C
+            import colossalai._C.fused_optim
            MultiTensorApply.available = True
            self.chunk_size = chunk_size
        except ImportError as err: