Move transformer config components to its own submodule

allenai · Oct 21, 2024 · d90292e · d90292e
1 parent 4d3b231
commit d90292e
Show file tree

Hide file tree

Showing 3 changed files with 574 additions and 549 deletions.
diff --git a/src/olmo_core/nn/transformer/__init__.py b/src/olmo_core/nn/transformer/__init__.py
@@ -11,13 +11,15 @@
     TransformerBlockConfig,
     TransformerBlockType,
 )
+from .config import (
+    TransformerActivationCheckpointingConfig,
+    TransformerConfig,
+    TransformerDataParallelConfig,
+)
 from .init import InitMethod
 from .model import (
     Transformer,
-    TransformerActivationCheckpointingConfig,
     TransformerActivationCheckpointingMode,
-    TransformerConfig,
-    TransformerDataParallelConfig,
     TransformerDataParallelWrappingStrategy,
 )