Language model block (#372)

Author: jlamypoirier

Dockerfile

@@ -47,3 +47,8 @@ COPY --chmod=777 ./tests tests
 COPY --chmod=777 ./tools tools
 COPY --chmod=777 ./fast_llm_external_models fast_llm_external_models
 COPY --chmod=777 --exclude=./fast_llm/csrc/ ./fast_llm/ fast_llm/
+
+# Set a dummy default user so we don't run in root by default.
+# The image is still compatible with any user id.
+RUN useradd user
+USER user

examples/mistral.yaml

@@ -28,7 +28,6 @@ optimizer:
 model:
   base_model:
     embeddings:
-      hidden_size: 4096
       vocab_size: 32000
       dropout: 0.0
     decoder:
@@ -58,6 +57,7 @@ model:
       normalization:
         type: rms_norm
         epsilon: 1.0e-05
+    hidden_size: 4096
     tied_embedding_weight: false
   multi_stage:
     zero_stage: 2

fast_llm/engine/checkpoint/huggingface.py

@@ -150,7 +150,6 @@ def _load_weights(
                 ].values()
             }
         elif (config.path / transformers.utils.WEIGHTS_NAME).is_file():
-            # TODO: Prevent unsafe by default
             paths = {config.path / transformers.utils.WEIGHTS_NAME}
         elif (config.path / transformers.utils.WEIGHTS_INDEX_NAME).is_file():
             logger.info(f"Loading index from {config.path / transformers.utils.WEIGHTS_INDEX_NAME}")
@@ -170,7 +169,7 @@ def _load_weights(
                     for key in f.keys():
                         yield key, "weights", f.get_slice(key)
             elif path.suffix == ".bin":
-                # TODO: Prevent unsafe by default
-                yield from torch.load(path)
+                # TODO: Confirm that loading works with `weights_only=True`
+                yield from torch.load(path, weights_only=True)
             else:
                 raise NotImplementedError(f"Unknown file format for {path}")

fast_llm/engine/config_utils/parameter.py

@@ -1,7 +1,8 @@
 import math
 import typing
 
-from fast_llm.config import Config, Field, FieldHint, config_class
+from fast_llm.config import Field, FieldHint, config_class
+from fast_llm.engine.base_model.config import ModuleConfig
 from fast_llm.engine.config_utils.initialization import Initialization, InitializationConfig
 from fast_llm.engine.config_utils.tensor_dim import TensorDim
 from fast_llm.layers.common.peft.config import PeftConfig
@@ -36,7 +37,7 @@ def combine_lr_scales(*lr_scales: float | None | tuple[float | None, ...]):
 
 
 @config_class()
-class ParameterConfig(Config):
+class ParameterConfig(ModuleConfig):
     initialization: InitializationConfig = Field(
         desc="If provided, override the default initialization method set by the parent layer.",
         hint=FieldHint.feature,

fast_llm/layers/block/block.py

@@ -103,6 +103,7 @@ def __init__(
         config: ConfigType,
         distributed_config: DistributedConfig,
         *,
+        # TODO: Review. Use `input_dim(s)` and `output_dim(s)` instead?
         hidden_dim: TensorDim,
         lr_scale: float | None,
         peft: PeftConfig | None,

fast_llm/layers/language_model/config.py

@@ -2,7 +2,6 @@
 import typing
 
 from fast_llm.config import Field, FieldHint, check_field, config_class, skip_valid_if_none
-from fast_llm.engine.base_model.config import ModuleConfig
 from fast_llm.engine.config_utils.parameter import OptionalParameterConfig, ParameterConfig, combine_lr_scales
 from fast_llm.engine.config_utils.tensor_dim import TensorDim
 from fast_llm.engine.distributed.config import DistributedConfig
@@ -16,6 +15,7 @@
 if typing.TYPE_CHECKING:
     from fast_llm.layers.language_model.embedding import LanguageModelEmbedding
     from fast_llm.layers.language_model.head import LanguageModelHead, LanguageModelHeadBase
+    from fast_llm.layers.language_model.language_model import LanguageModel
     from fast_llm.layers.language_model.multi_token_prediction import MultiTokenPrediction
 
 
@@ -41,12 +41,6 @@ class LanguageModelEmbeddingsConfig(BlockConfig):
         desc="Configuration for the word embedding (weight).",
         hint=FieldHint.architecture,
     )
-    hidden_size: int = Field(
-        default=1024,
-        desc="Size of the model's main hidden dimension, e.g., for its input and output layers.",
-        hint=FieldHint.architecture,
-        valid=check_field(Assert.gt, 0),
-    )
     vocab_size: int = Field(
         default=49152,
         desc="Size of the vocabulary, i.e., number of vocabulary embeddings and logits.",
@@ -295,24 +289,29 @@ def max_prediction_distance(self) -> int:
 
 
 @config_class()
-class LanguageModelConfig(ModuleConfig):
-    # TODO: block
+class LanguageModelConfig(BlockConfig):
     decoder: BlockSequenceConfig = Field(
         desc="Configuration for the language model decoder.",
         hint=FieldHint.architecture,
     )
-    embeddings: LanguageModelEmbeddingsConfig = Field()
-    head: LanguageModelHeadBaseConfig = Field()
-    # TODO: Allow overriding in sub-models?
-    peft: PeftConfig = Field(
-        desc="Configuration for parameter-efficient fine tuning.",
+    embeddings: LanguageModelEmbeddingsConfig = Field(
         hint=FieldHint.architecture,
+        desc="Configuration for the language model embeddings.",
+    )
+    head: LanguageModelHeadBaseConfig = Field(
+        hint=FieldHint.architecture, desc="Configuration for the language model head(s)."
     )
     tied_embedding_weight: bool = Field(
         default=False,
         desc="Tie the output weights (logits) with the vocabulary embedding.",
         hint=FieldHint.architecture,
     )
+    hidden_size: int = Field(
+        default=1024,
+        desc="Size of the model's main hidden dimension, e.g., for its input and output layers.",
+        hint=FieldHint.architecture,
+        valid=check_field(Assert.gt, 0),
+    )
     sequence_first: bool | None = Field(
         default=None,
         desc="Override the default dimension ordering",
@@ -321,3 +320,9 @@ class LanguageModelConfig(ModuleConfig):
         " Setting this parameter overrides the default choice. Note that setting to `False` will either do nothing or raise an error.",
         hint=FieldHint.testing,
     )
+
+    @property
+    def layer_class(self) -> "type[LanguageModel]":
+        from fast_llm.layers.language_model.language_model import LanguageModel
+
+        return LanguageModel

fast_llm/layers/language_model/language_model.py

@@ -1,52 +1,64 @@
 import logging
 import typing
 
-from fast_llm.config import Configurable
-from fast_llm.engine.base_model.base_model import Layer, LayerBase
+import torch
+
+from fast_llm.engine.base_model.base_model import Layer
 from fast_llm.engine.base_model.config import LossDef
 from fast_llm.engine.config_utils.tensor_dim import TensorDim
 from fast_llm.engine.distributed.config import DistributedConfig
+from fast_llm.layers.block.block import BlockBase
+from fast_llm.layers.common.peft.config import PeftConfig
 from fast_llm.layers.language_model.config import LanguageModelConfig
 from fast_llm.layers.language_model.embedding import LanguageModelEmbedding
 
 logger = logging.getLogger(__name__)
 
 
-class LanguageModel[ConfigType: LanguageModelConfig](Configurable[ConfigType], LayerBase):
+class LanguageModel[ConfigType: LanguageModelConfig](BlockBase[ConfigType]):
     _config: ConfigType
 
     def __init__(
         self,
         config: ConfigType,
         distributed_config: DistributedConfig,
+        *,
+        # TODO: Unused, but required by the `BlockBase` interface.
+        hidden_dim: TensorDim | None = None,
+        lr_scale: float | None,
+        peft: PeftConfig | None,
     ):
-        super().__init__(config, distributed_config)
-
-        self._hidden_dim = TensorDim("hidden", config.embeddings.hidden_size)
+        super().__init__(
+            config,
+            distributed_config,
+            hidden_dim=TensorDim("hidden", config.hidden_size),
+            lr_scale=lr_scale,
+            peft=peft,
+        )
         self.embeddings: LanguageModelEmbedding = self._config.embeddings.get_layer(
             distributed_config,
             hidden_dim=self._hidden_dim,
-            lr_scale=None,
-            peft=self._config.peft,
+            lr_scale=self._lr_scale,
+            peft=self._peft,
         )
         self.decoder = self._config.decoder.get_layer(
             distributed_config,
             self._hidden_dim,
-            lr_scale=None,
-            peft=self._config.peft,
+            lr_scale=self._lr_scale,
+            peft=self._peft,
         )
         self.head = self._config.head.get_layer(
             distributed_config,
             self._config.embeddings,
             hidden_dim=self._hidden_dim,
-            lr_scale=None,
-            peft=self._config.peft,
+            lr_scale=self._lr_scale,
+            peft=self._peft,
         )
 
-    def get_layers(self) -> list["Layer"]:
+    def get_layers(self) -> list[Layer]:
         return self.embeddings.get_layers() + self.decoder.get_layers() + self.head.get_layers()
 
-    def preprocess(self, batch: "torch.Tensor", kwargs: dict[str, typing.Any]) -> None:
+    def preprocess(self, batch: torch.Tensor, kwargs: dict[str, typing.Any]) -> None:
         # Needed because the base class uses `get_layers` which may bypass the decoder and head. TODO: Avoidable?
         self.embeddings.preprocess(batch, kwargs)
         self.decoder.preprocess(batch, kwargs)

fast_llm/models/gpt/config.py

@@ -10,6 +10,7 @@
 from fast_llm.engine.multi_stage.config import FastLLMModelConfig, PretrainedFastLLMModelConfig
 from fast_llm.engine.schedule.config import BatchConfig
 from fast_llm.engine.training.config import TrainerConfig
+from fast_llm.layers.common.peft.config import PeftConfig
 from fast_llm.layers.language_model.config import LanguageModelConfig, MultiTokenPredictionConfig
 from fast_llm.models.gpt.conversion.config import (
     AprielHybridSSMCheckpointFormat,
@@ -84,6 +85,11 @@ def micro_batch_splits(self) -> int:
 class GPTBaseModelConfig(LanguageModelConfig, BaseModelConfig):
     _abstract = False
 
+    # TODO: Allow overriding in sub-models?
+    peft: PeftConfig = Field(
+        desc="Configuration for parameter-efficient fine tuning.",
+        hint=FieldHint.architecture,
+    )
     # Debug, to get an exact match with megatron init.
     use_megatron_initialization: bool = Field(
         default=False, desc="Exactly match the initialization of a Megatron model.", hint=FieldHint.testing

fast_llm/models/gpt/conversion/llama.py

@@ -449,19 +449,13 @@ def get_converters(
 class LlamaEmbeddingsConverter:
     @classmethod
     def import_config(cls, config: dict) -> dict:
-        return {
-            "vocab_size": config["vocab_size"],
-            "hidden_size": config["hidden_size"],
-        }
+        return {"vocab_size": config["vocab_size"]}
 
     @classmethod
     def export_config(cls, config: LanguageModelEmbeddingsConfig) -> dict:
         Assert.custom(isinstance, config, LanguageModelEmbeddingsConfig)
         assert not config.position_embeddings.enabled
-        return {
-            "vocab_size": config.vocab_size,
-            "hidden_size": config.hidden_size,
-        }
+        return {"vocab_size": config.vocab_size}
 
     @classmethod
     def get_converters(
@@ -516,6 +510,7 @@ def import_config(cls, config: dict) -> dict:
             "embeddings": cls.embeddings_converter_class.import_config(config),
             "decoder": cls.decoder_converter_class.import_config(config),
             "head": cls.head_converter_class.import_config(config),
+            "hidden_size": config["hidden_size"],
             "tied_embedding_weight": config["tie_word_embeddings"],
         }
 
@@ -526,7 +521,10 @@ def export_config(cls, config: GPTBaseModelConfig) -> dict:
             cls.embeddings_converter_class.export_config(config.embeddings),
             cls.decoder_converter_class.export_config(config.decoder),
             cls.head_converter_class.export_config(config.head),
-            {"tie_word_embeddings": config.tied_embedding_weight},
+            {
+                "tie_word_embeddings": config.tied_embedding_weight,
+                "hidden_size": config.hidden_size,
+            },
         )
 
     @classmethod

fast_llm/models/gpt/model.py

@@ -30,16 +30,14 @@ class GPTBaseModel[ConfigType: GPTBaseModelConfig](LanguageModel[ConfigType], Ba
 
     def __init__(
         self,
-        config: GPTBaseModelConfig,
+        config: ConfigType,
         distributed_config: DistributedConfig,
     ):
-        super().__init__(config, distributed_config)
+        super().__init__(config, distributed_config, lr_scale=config.lr_scale, peft=config.peft)
         if self._config.use_megatron_initialization:
             for param in self.parameters():
                 Assert.custom(isinstance, param, ParameterMeta)
-                param.init_parameter = get_init_megatron(
-                    param, self._config.decoder.block, config.embeddings.hidden_size
-                )  # Noqa
+                param.init_parameter = get_init_megatron(param, self._config.decoder.block, config.hidden_size)  # Noqa
 
     def preprocess_meta(
         self, batch_meta: GPTBatchConfig | torch.Tensor, phase: PhaseType