Block interface: extract mixer and mlp config

examples/mistral.yaml

@@ -28,24 +28,27 @@ optimizer:
 model:
   base_model:
     transformer:
+      mixer:
+        type: attention
+        rotary:
+          type: default
+          theta: 10000
+        num_attention_heads: 32
+        head_groups: 8
+        kv_channels: 128
+        window_size: 4096
+        attention_dropout: 0.0
+      mlp:
+        ffn_hidden_size: 14336
+        gated: true
+        activation_type: silu
       normalization:
         type: rms_norm
         epsilon: 1.0e-05
-      rotary:
-        type: default
-        theta: 10000
       num_layers: 32
       hidden_size: 4096
-      ffn_hidden_size: 14336
-      num_attention_heads: 32
-      head_groups: 8
       add_linear_biases: false
-      gated: true
-      activation_type: silu
-      kv_channels: 128
-      window_size: 4096
       init_method_std: 0.009021
-      attention_dropout: 0.0
       hidden_dropout: 0.0
     vocab_size: 32000
     tie_word_embeddings: false

fast_llm/engine/evaluation/lm_eval/fast_llm_wrapper.py

@@ -104,7 +104,7 @@ def max_length(self):
         # check if it is absolute positional encoding and return max_position_embeddings
         if hasattr(self._config.fast_llm_config.base_model, "transformer"):
             # NOTE: will need to extend if more relative encoding types will be added
-            if isinstance(self._config.fast_llm_config.base_model.transformer.rotary, NoRotaryConfig):
+            if isinstance(self._config.fast_llm_config.base_model.transformer.mixer.rotary, NoRotaryConfig):
                 return self._config.fast_llm_config.base_model.max_position_embeddings
 
         # check if tokenizer holds model sequence leigh info

fast_llm/layers/attention/attention.py

@@ -94,17 +94,6 @@ def __init__(
 
         self._softmax_scale = self._config.kv_channels ** (-self._config.attention_softmax_scale_power)
 
-        init_method_qkv = init_normal_(
-            std=self._config.init_method_std_qkv,
-            min_val=self._config.init_method_min_qkv,
-            max_val=self._config.init_method_max_qkv,
-        )
-        init_method_std_attn_proj = init_normal_(
-            std=self._config.init_method_std_attn_proj,
-            min_val=self._config.init_method_min_attn_proj,
-            max_val=self._config.init_method_max_attn_proj,
-        )
-
         lr_scale = combine_lr_scales(
             self._lr_scale,
             self._config.attention_lr_scale,
@@ -114,7 +103,7 @@ def __init__(
         self.query = self._config.query_layer.get_layer(
             hidden_dim,
             query_dim,
-            default_weight_initializer=init_method_qkv,
+            default_weight_initializer=init_normal_(std=self._block_config.init_method_std),
             default_add_bias=self._block_config.add_linear_biases,
             sequence_parallel=self._sequence_parallel,
             lr_scale=lr_scale,
@@ -123,7 +112,7 @@ def __init__(
         self.key_value = self._config.query_layer.get_layer(
             hidden_dim,
             key_value_dim,
-            default_weight_initializer=init_method_qkv,
+            default_weight_initializer=init_normal_(std=self._block_config.init_method_std),
             default_add_bias=self._block_config.add_linear_biases,
             sequence_parallel=self._sequence_parallel,
             lr_scale=lr_scale,
@@ -137,7 +126,9 @@ def __init__(
         self.dense = self._config.dense_layer.get_layer(
             dense_dim,
             hidden_dim,
-            default_weight_initializer=init_method_std_attn_proj,
+            default_weight_initializer=init_normal_(
+                std=self._block_config.init_method_std / max(2 * self._block_config.num_layers, 1) ** 0.5,
+            ),
             default_add_bias=self._block_config.add_linear_biases,
             sequence_parallel=self._sequence_parallel,
             lr_scale=lr_scale,

fast_llm/layers/attention/block.py

@@ -1,8 +1,6 @@
-import functools
 import logging
 import typing
 
-from fast_llm.layers.attention.attention import Attention
 from fast_llm.layers.attention.config import AttentionConfig, TransformerConfig
 from fast_llm.layers.block.block import Block
 
@@ -13,10 +11,6 @@ class TransformerBlock[ConfigType: TransformerConfig](Block[ConfigType]):
     # TODO: Standardize to `mixer`
     _mixer_module_name: typing.ClassVar[str] = "self_attn"
 
-    @functools.cached_property
-    def _mixer_class(self) -> type[Attention]:
-        return Attention
-
     @property
     def _mixer_config(self) -> AttentionConfig:
-        return self._config
+        return self._config.mixer

fast_llm/layers/attention/config.py

@@ -1,16 +1,20 @@
 import functools
 import logging
+import typing
 import warnings
 
-from fast_llm.config import Config, Field, FieldHint, check_field, config_class, skip_valid_if_none
+from fast_llm.config import Field, FieldHint, FieldUpdate, check_field, config_class, skip_valid_if_none
 from fast_llm.engine.config_utils.data_type import DataType
 from fast_llm.engine.distributed.config import DistributedConfig
 from fast_llm.functional.config import TritonConfig
 from fast_llm.layers.attention.rotary.config import RotaryConfig
-from fast_llm.layers.block.config import BlockConfig, BlockKwargs
+from fast_llm.layers.block.config import BlockConfig, BlockKwargs, MixerConfig
 from fast_llm.layers.common.linear.config import AffineLinearConfig
 from fast_llm.utils import Assert, div
 
+if typing.TYPE_CHECKING:
+    from fast_llm.layers.attention.attention import Attention
+
 logger = logging.getLogger(__name__)
 
 
@@ -28,8 +32,8 @@ class AttentionKwargs(BlockKwargs):
     past_key_values = "past_key_values"
 
 
-@config_class()
-class AttentionConfig(Config):
+@config_class(dynamic_type={MixerConfig: "attention"})
+class AttentionConfig(MixerConfig):
     # TODO: Make mixer class dynamic.
     _abstract = False
 
@@ -106,72 +110,26 @@ class AttentionConfig(Config):
         " Under muP (if scaling number of heads instead of kv_channels): use 0.5.",
         valid=skip_valid_if_none(check_field(Assert.geq, 0)),
     )
-    # TODO: Review initialization
-    init_method_std_qkv: float = Field(
-        default=None,
-        desc="Scale for the query, key and value weight initialization. Default: init_method_std",
-        hint=FieldHint.optional,
-        valid=check_field(Assert.geq, 0),
-    )
-    init_method_max_qkv: float | None = Field(
-        default=None,
-        desc="Max value for clamping initialized weights for query, key and value matrices. Default: float('inf')",
-        hint=FieldHint.optional,
-    )
-    init_method_min_qkv: float | None = Field(
-        default=None,
-        desc="Min value for clamping initialized weights for query, key and value matrices. Default: -float('inf')",
-        hint=FieldHint.optional,
-    )
-    init_method_std_attn_proj: float = Field(
-        default=None,
-        desc="Scale for the attention projection weight initialization. Default: init_method_std",
-        hint=FieldHint.optional,
-        valid=check_field(Assert.geq, 0),
-    )
-    init_method_max_attn_proj: float | None = Field(
-        default=None,
-        desc="Max value for clamping initialized weights for attention projection. Default: float('inf')",
-        hint=FieldHint.optional,
-    )
-    init_method_min_attn_proj: float | None = Field(
-        default=None,
-        desc="Min value for clamping initialized weights for attention projection. Default: -float('inf')",
-        hint=FieldHint.optional,
-    )
 
-    def _validate(self) -> None:
-        with self._set_implicit_default():
-            # TODO: Make this work without inheritance.
-            if self.kv_channels is None:
-                self.kv_channels = div(self.hidden_size, self.num_attention_heads)
-            # TODO: Review initialization
-            if self.init_method_std_qkv is None:
-                self.init_method_std_qkv = self.init_method_std
-            if self.init_method_std_attn_proj is None:
-                self.init_method_std_attn_proj = self.init_method_std / max(2 * self.num_layers, 1) ** 0.5
-            if self.init_method_max_qkv is None:
-                self.init_method_max_qkv = self.init_method_max
-            if self.init_method_min_qkv is None:
-                self.init_method_min_qkv = self.init_method_min
-            if self.init_method_max_attn_proj is None:
-                self.init_method_max_attn_proj = self.init_method_max
-            if self.init_method_min_attn_proj is None:
-                self.init_method_min_attn_proj = self.init_method_min
-            if self.init_method_min_qkv is not None and self.init_method_max_qkv is not None:
-                Assert.leq(self.init_method_min, self.init_method_max)
-            if self.init_method_min_qkv is not None and self.init_method_max_qkv is not None:
-                Assert.leq(self.init_method_min_qkv, self.init_method_max_qkv)
-            if self.init_method_min_attn_proj is not None and self.init_method_max_attn_proj is not None:
-                Assert.leq(self.init_method_min_attn_proj, self.init_method_max_attn_proj)
+    def set_defaults(self, hidden_size: int):
+        if self.kv_channels is None:
+            with self._set_implicit_default():
+                self.kv_channels = div(hidden_size, self.num_attention_heads)
 
+    def _validate(self) -> None:
         super()._validate()
 
         if not TritonConfig.TRITON_ENABLED:
             warnings.warn("Triton is disabled, but triton rotary kernel will be used anyway.")
 
         Assert.multiple(self.num_attention_heads, self.head_groups)
 
+    @property
+    def layer_class(self) -> "type[Attention]":
+        from fast_llm.layers.attention.attention import Attention
+
+        return Attention
+
     @functools.cached_property
     def projection_size(self):
         assert self._validated
@@ -183,16 +141,7 @@ def do_use_flash_attention(self, distributed_config: DistributedConfig) -> bool:
 
 @config_class()
 # TODO: Use composition instead
-class TransformerConfig(AttentionConfig, BlockConfig):
+class TransformerConfig(BlockConfig):
     _abstract = False
-
-    def _validate(self) -> None:
-        with self._set_implicit_default():
-            # Kept here for initialization order.
-            # TODO: Review initialization
-            if self.init_method_std is None:
-                self.init_method_std = self.hidden_size**-0.5
-            if self.init_method_min is not None and self.init_method_max is not None:
-                Assert.leq(self.init_method_min, self.init_method_max)
-
-        super()._validate()
+    # TODO: Make this unnecessary
+    mixer: AttentionConfig = FieldUpdate()

fast_llm/layers/attention/preprocessing.py

@@ -86,9 +86,7 @@ def preprocess_meta(self, kwargs: dict[str, typing.Any]) -> None:
 
 class FlashAttnVarlenPreprocessor(Preprocessor):
     def __init__(self, config: AttentionConfig, distributed_config: DistributedConfig):
-        self._config = config
-        self._distributed_config = distributed_config
-        assert self._config.do_use_flash_attention(self._distributed_config)
+        assert config.do_use_flash_attention(distributed_config)
 
     def preprocess(self, batch: torch.Tensor, kwargs: dict[str, typing.Any]) -> None:
         """

fast_llm/layers/block/block.py

@@ -12,7 +12,7 @@
 from fast_llm.engine.config_utils.tensor_dim import TensorDim
 from fast_llm.engine.distributed.config import DistributedConfig
 from fast_llm.engine.distributed.distributed import Distributed
-from fast_llm.layers.block.config import BlockConfig, BlockKwargs
+from fast_llm.layers.block.config import BlockConfig, BlockKwargs, MixerConfig
 from fast_llm.logging import log_distributed_grad, log_distributed_tensor, log_memory_usage
 from fast_llm.tensor import TensorMeta
 
@@ -174,8 +174,7 @@ def __init__(
         setattr(
             self,
             self._mixer_module_name,
-            self._mixer_class(
-                self._mixer_config,
+            self._mixer_config.get_layer(
                 self._config,
                 self._distributed_config,
                 self._hidden_dim,
@@ -185,12 +184,7 @@ def __init__(
             ),
         )
 
-        # TODO: Use dynamic type.
-        from fast_llm.layers.block.mlp.mixture_of_experts import MixtureOfExpertMLP
-        from fast_llm.layers.block.mlp.mlp import MLP
-
-        self.mlp = (MixtureOfExpertMLP if self._config.num_experts > 1 else MLP)(
-            self._config,
+        self.mlp = self._config.mlp.get_layer(
             self._config,
             self._distributed_config,
             self._hidden_dim,
@@ -199,14 +193,9 @@ def __init__(
             self._lr_scale,
         )
 
-    @functools.cached_property
-    @abc.abstractmethod
-    def _mixer_class(self) -> type[BlockLayer]:
-        pass
-
     @property
     @abc.abstractmethod
-    def _mixer_config(self) -> Config:
+    def _mixer_config(self) -> MixerConfig:
         pass
 
     def setup(self, distributed: Distributed) -> None: