[Feature] Support pipeline parallelism model wrapper

open-mmlab · Sep 5, 2023 · ae13282 · ae13282
1 parent 45ee96d
commit ae13282
Show file tree

Hide file tree

Showing 2 changed files with 409 additions and 0 deletions.
diff --git a/mmengine/model/wrappers/__init__.py b/mmengine/model/wrappers/__init__.py
@@ -14,3 +14,7 @@
     from .fully_sharded_distributed import \
         MMFullyShardedDataParallel  # noqa:F401
     __all__.append('MMFullyShardedDataParallel')
+
+    from .pipeline_distributed import \
+        MMPipelineParallel  # noqa:F401
+    __all__.append('MMPipelineParallel')