refactor!: rename stage block types

README.md

@@ -46,12 +46,13 @@ MBLM can be used with the default Transformer Decoder or Mamba block. The below
 
 ```py
 import torch
+
 from mblm import (
     MBLM,
-    MambaBlockConfig,
+    MambaBlock,
     MBLMModelConfig,
     MBLMReturnType,
-    TransformerBlockConfig,
+    TransformerBlock,
 )
 
 mblm = MBLM(
@@ -63,14 +64,14 @@ mblm = MBLM(
         pad_token_id=256,
         train_checkpoint_chunks=None,
         block=[
-            MambaBlockConfig(
+            MambaBlock(
                 d_state=128,
                 d_conv=4,
                 expand=2,
                 headdim=64,
                 pos_emb_type=None,
             ),
-            TransformerBlockConfig(
+            TransformerBlock(
                 attn_head_dims=64,
                 attn_num_heads=16,
                 attn_use_rot_embs=True,
@@ -97,6 +98,7 @@ Alternatively, you can read configuration from a YAML string (or file):
 ```py
 import torch
 import yaml
+
 from mblm import MBLM, MBLMModelConfig, MBLMReturnType
 
 yml_model_config = """
@@ -131,17 +133,13 @@ You can define custom stage blocks for MBLM as follows. A stageblock must provid
 
 ```py
 import torch
-from mblm import (
-    MBLM,
-    MBLMModelConfig,
-    MBLMReturnType,
-    TransformerBlockConfig,
-)
+from pydantic import Field
+
+from mblm import MBLM, MBLMModelConfig, MBLMReturnType, TransformerBlock
 from mblm.model.block import StageBlock
-from pydantic import BaseModel, Field
 
 # Define any custom model
-class MyLSTM(torch.nn.Module):
+class LSTM(torch.nn.Module):
     def __init__(self, lstm: torch.nn.LSTM):
         super().__init__()
         self.lstm = lstm
@@ -151,15 +149,15 @@ class MyLSTM(torch.nn.Module):
         out, _ = self.lstm(input_ids)
         return out
 
-# Add a block config and inherit from StageBlock and BaseModel
-class LSTMBlockConfig(StageBlock, BaseModel):
+# Add a block config and inherit from StageBlock
+class LSTMBlock(StageBlock):
     block_type: str = Field(init=False, default="lstm")
 
     # Add whatever is needed
     dropout: float
 
     def to_model(self, model_dim: int, num_layers: int) -> torch.nn.Module:
-        return MyLSTM(
+        return LSTM(
             torch.nn.LSTM(
                 input_size=model_dim,
                 hidden_size=model_dim,
@@ -178,11 +176,11 @@ mblm = MBLM(
         pad_token_id=256,
         train_checkpoint_chunks=None,
         block=[
-            LSTMBlockConfig(
+            LSTMBlock(
                 dropout=0.1,
                 pos_emb_type=None,
             ),
-            TransformerBlockConfig(
+            TransformerBlock(
                 attn_head_dims=64,
                 attn_num_heads=16,
                 attn_use_rot_embs=True,
@@ -202,10 +200,11 @@ If you want to parse a YAML config to a custom block, **register the block** bef
 ```py
 import torch
 import yaml
+from pydantic import Field
+
 from mblm import MBLM, MBLMModelConfig, MBLMReturnType
 from mblm.model.block import StageBlock
 from mblm.model.config import block_registry  # Add this!
-from pydantic import BaseModel, Field
 
 # Define any custom model
 class MyLSTM(torch.nn.Module):
@@ -218,8 +217,8 @@ class MyLSTM(torch.nn.Module):
         out, _ = self.lstm(input_ids)
         return out
 
-# Add a block config and inherit from StageBlock and BaseModel
-class LSTMBlockConfig(StageBlock, BaseModel):
+# Add a block config and inherit from StageBlock
+class LSTMBlockConfig(StageBlock):
     block_type: str = Field(init=False, default="lstm")
 
     # Add whatever is needed
@@ -269,10 +268,11 @@ If you want to use the MBLM trainer with [torchrun](https://pytorch.org/docs/sta
 # Filename: train_my_mblm.py
 
 import torch
+from typing_extensions import Unpack
+
+from mblm import MambaBlock, TransformerBlock
 from mblm.data.datasets import DistributedDataset, DistributedDatasetConfig
 from mblm.data.types import BatchWithLossMask, ModelMode
-from mblm.model.mamba import MambaBlockConfig
-from mblm.model.transformer import TransformerBlockConfig
 from mblm.train.core.config import CoreTrainConfig
 from mblm.train.mblm import (
     TrainEntryConfig,
@@ -281,7 +281,6 @@ from mblm.train.mblm import (
     dataset_registry,
     train_mblm,
 )
-from typing_extensions import Unpack
 
 
 class MyDataset(DistributedDataset[BatchWithLossMask]):
@@ -372,14 +371,14 @@ config = TrainEntryConfig(
         pad_token_id=256,
         train_checkpoint_chunks=None,
         block=[
-            MambaBlockConfig(
+            MambaBlock(
                 d_state=128,
                 d_conv=4,
                 expand=2,
                 headdim=64,
                 pos_emb_type=None,
             ),
-            TransformerBlockConfig(
+            TransformerBlock(
                 attn_head_dims=64,
                 attn_num_heads=16,
                 attn_use_rot_embs=True,
@@ -392,6 +391,7 @@ config = TrainEntryConfig(
 
 if __name__ == "__main__":
     train_mblm(config)
+
 ```
 
 Then, run the above file with:

src/mblm/__init__.py

@@ -11,14 +11,14 @@
 
 
 from mblm.model.config import MBLMModelConfig, MBLMReturnType
-from mblm.model.mamba import MambaBlockConfig
+from mblm.model.mamba import MambaBlock
 from mblm.model.mblm import MBLM
-from mblm.model.transformer import TransformerBlockConfig
+from mblm.model.transformer import TransformerBlock
 
 __all__ = [
     "MBLM",
     "MBLMModelConfig",
     "MBLMReturnType",
-    "TransformerBlockConfig",
-    "MambaBlockConfig",
+    "TransformerBlock",
+    "MambaBlock",
 ]

src/mblm/model/config.py

@@ -27,12 +27,12 @@
 from pydantic import BaseModel, computed_field, field_validator, model_validator
 
 from mblm.model.block import StageBlock, StageBlockRegistry
-from mblm.model.mamba import MambaBlockConfig
-from mblm.model.transformer import TransformerBlockConfig
+from mblm.model.mamba import MambaBlock
+from mblm.model.transformer import TransformerBlock
 
 block_registry = StageBlockRegistry()
-block_registry.register(TransformerBlockConfig)
-block_registry.register(MambaBlockConfig)
+block_registry.register(TransformerBlock)
+block_registry.register(MambaBlock)
 
 
 class MBLMReturnType(str, Enum):

src/mblm/model/mamba.py

@@ -21,13 +21,13 @@
 SOFTWARE."""
 
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from mblm.model.block import StageBlock
 from mblm.model.mamba_shim import Mamba1, Mamba1Config, Mamba2Mixer
 
 
-class MambaBlockConfig(StageBlock, BaseModel):
+class MambaBlock(StageBlock):
     """
     General config for creating a Mamba block inside MBLM.
     Uses roughly 3 * expand * d_model^2 parameters.

src/mblm/model/transformer.py

@@ -25,12 +25,12 @@
 
 import torch
 from MEGABYTE_pytorch.megabyte import Attention, FeedForward, RMSNorm, RotaryEmbedding, token_shift
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from mblm.model.block import StageBlock
 
 
-class TransformerBlockConfig(StageBlock, BaseModel):
+class TransformerBlock(StageBlock):
     """
     General config for creating a Transformer Decocer block inside MBLM.
     """

tests/integration/config/test_sample_config_to_model.py

@@ -3,10 +3,8 @@
 
 import pytest
 
-from mblm import MBLM, MBLMModelConfig
+from mblm import MBLM, MambaBlock, MBLMModelConfig, TransformerBlock
 from mblm.data.dataset.clevr import ClevrOptionalArgs
-from mblm.model.mamba import MambaBlockConfig
-from mblm.model.transformer import TransformerBlockConfig
 from mblm.train.mblm import TrainEntryConfig
 from mblm.utils.io import load_yml
 
@@ -31,8 +29,8 @@ def ensure_dataset_args_are_valid(self, config: TrainEntryConfig) -> None:
 
     def ensure_model_is_created(self, config: TrainEntryConfig) -> None:
         for b in config.params.stage_blocks:
-            assert isinstance(b, (TransformerBlockConfig, MambaBlockConfig))
-            if isinstance(b, TransformerBlockConfig):
+            assert isinstance(b, (TransformerBlock, MambaBlock))
+            if isinstance(b, TransformerBlock):
                 assert b.block_type == "transformer"
             else:
                 # mamba1, can be mamba2 (only if tested on Linux with mamba_ssm installed)

tests/integration/install/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mblm-test"
-version = "0.0.1"
+version = "0.0.0"
 description = "Multiscale Byte Language Model - Test project"
 authors = [
     { name = "Eric Egli", email = "[email protected]" },

tests/integration/install/test_custom_block.py

@@ -5,17 +5,13 @@
 
 def test_from_config():
     import torch
-    from mblm import (
-        MBLM,
-        MBLMModelConfig,
-        MBLMReturnType,
-        TransformerBlockConfig,
-    )
+    from pydantic import Field
+
+    from mblm import MBLM, MBLMModelConfig, MBLMReturnType, TransformerBlock
     from mblm.model.block import StageBlock
-    from pydantic import BaseModel, Field
 
     # Define any custom model
-    class MyLSTM(torch.nn.Module):
+    class LSTM(torch.nn.Module):
         def __init__(self, lstm: torch.nn.LSTM):
             super().__init__()
             self.lstm = lstm
@@ -25,15 +21,15 @@ def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
             out, _ = self.lstm(input_ids)
             return out
 
-    # Add a block config and inherit from StageBlock and BaseModel
-    class LSTMBlockConfig(StageBlock, BaseModel):
+    # Add a block config and inherit from StageBlock
+    class LSTMBlock(StageBlock):
         block_type: str = Field(init=False, default="lstm")
 
         # Add whatever is needed
         dropout: float
 
         def to_model(self, model_dim: int, num_layers: int) -> torch.nn.Module:
-            return MyLSTM(
+            return LSTM(
                 torch.nn.LSTM(
                     input_size=model_dim,
                     hidden_size=model_dim,
@@ -52,11 +48,11 @@ def to_model(self, model_dim: int, num_layers: int) -> torch.nn.Module:
             pad_token_id=256,
             train_checkpoint_chunks=None,
             block=[
-                LSTMBlockConfig(
+                LSTMBlock(
                     dropout=0.1,
                     pos_emb_type=None,
                 ),
-                TransformerBlockConfig(
+                TransformerBlock(
                     attn_head_dims=64,
                     attn_num_heads=16,
                     attn_use_rot_embs=True,
@@ -74,10 +70,11 @@ def to_model(self, model_dim: int, num_layers: int) -> torch.nn.Module:
 def test_from_yaml():
     import torch
     import yaml
+    from pydantic import Field
+
     from mblm import MBLM, MBLMModelConfig, MBLMReturnType
     from mblm.model.block import StageBlock
     from mblm.model.config import block_registry  # Add this!
-    from pydantic import BaseModel, Field
 
     # Define any custom model
     class MyLSTM(torch.nn.Module):
@@ -90,8 +87,8 @@ def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
             out, _ = self.lstm(input_ids)
             return out
 
-    # Add a block config and inherit from StageBlock and BaseModel
-    class LSTMBlockConfig(StageBlock, BaseModel):
+    # Add a block config and inherit from StageBlock
+    class LSTMBlockConfig(StageBlock):
         block_type: str = Field(init=False, default="lstm")
 
         # Add whatever is needed

tests/integration/install/test_custom_dataset.py

@@ -1,10 +1,11 @@
 # Filename: train_my_mblm.py
 
 import torch
+from typing_extensions import Unpack
+
+from mblm import MambaBlock, TransformerBlock
 from mblm.data.datasets import DistributedDataset, DistributedDatasetConfig
 from mblm.data.types import BatchWithLossMask, ModelMode
-from mblm.model.mamba import MambaBlockConfig
-from mblm.model.transformer import TransformerBlockConfig
 from mblm.train.core.config import CoreTrainConfig
 from mblm.train.mblm import (
     TrainEntryConfig,
@@ -13,7 +14,6 @@
     dataset_registry,
     train_mblm,
 )
-from typing_extensions import Unpack
 
 
 class MyDataset(DistributedDataset[BatchWithLossMask]):
@@ -104,14 +104,14 @@ def supports_test_mode() -> bool:
         pad_token_id=256,
         train_checkpoint_chunks=None,
         block=[
-            MambaBlockConfig(
+            MambaBlock(
                 d_state=128,
                 d_conv=4,
                 expand=2,
                 headdim=64,
                 pos_emb_type=None,
             ),
-            TransformerBlockConfig(
+            TransformerBlock(
                 attn_head_dims=64,
                 attn_num_heads=16,
                 attn_use_rot_embs=True,