Move pipeline logic to its own train module

allenai · Jan 31, 2025 · a96f127 · a96f127
1 parent 7e4a566
commit a96f127
Show file tree

Hide file tree

Showing 4 changed files with 1,099 additions and 478 deletions.
diff --git a/src/olmo_core/train/train_module/__init__.py b/src/olmo_core/train/train_module/__init__.py
@@ -9,11 +9,15 @@
     TransformerActivationCheckpointingMode,
     TransformerDataParallelConfig,
     TransformerDataParallelWrappingStrategy,
-    TransformerPipelineParallelConfig,
     TransformerTensorParallelConfig,
     TransformerTrainModule,
     TransformerTrainModuleConfig,
 )
+from .transformer_pipeline import (
+    TransformerPipelineParallelConfig,
+    TransformerPipelineTrainModule,
+    TransformerPipelineTrainModuleConfig,
+)
 
 __all__ = [
     "TrainModule",
@@ -22,6 +26,8 @@
     "BasicTrainModule",
     "TransformerTrainModule",
     "TransformerTrainModuleConfig",
+    "TransformerPipelineTrainModule",
+    "TransformerPipelineTrainModuleConfig",
     "TransformerActivationCheckpointingConfig",
     "TransformerActivationCheckpointingMode",
     "TransformerDataParallelConfig",