Support tensor parallel (#2)

2023-03-22 04:45:42 +08:00
parent cfae35b861
commit 2f49f15585
24 changed files with 2480 additions and 174 deletions
--- a/cacheflow/parallel_utils/init.py
+++ b/cacheflow/parallel_utils/init.py
@@ -0,0 +1,12 @@
+import cacheflow.parallel_utils.parallel_state
+import cacheflow.parallel_utils.tensor_parallel
+import cacheflow.parallel_utils.utils
+
+# Alias parallel_state as mpu, its legacy name
+mpu = parallel_state
+
+__all__ = [
+    "parallel_state",
+    "tensor_parallel",
+    "utils",
+]