Generalize preprocessor (#224)

Author: jlamypoirier

fast_llm/engine/base_model/base_model.py

@@ -6,7 +6,7 @@
 import torch.nn
 
 from fast_llm.config import Configurable
-from fast_llm.engine.base_model.config import BaseModelArchitectureConfig, BaseModelConfig
+from fast_llm.engine.base_model.config import BaseModelArchitectureConfig, BaseModelConfig, Preprocessor
 from fast_llm.engine.config_utils.tensor_space import TensorSpace
 from fast_llm.engine.distributed.config import DistributedConfig, PhaseType
 from fast_llm.engine.distributed.distributed import Distributed
@@ -135,3 +135,7 @@ def get_tied_weights(self) -> dict[str, tuple[ParameterMeta, tuple[int, ...]]]:
     @abc.abstractmethod
     def loss_defs(self) -> list[LossDef]:
         pass
+
+    def add_preprocessor(self, preprocessor: Preprocessor):
+        # TODO: Generalize preprocessors.
+        raise NotImplementedError()

fast_llm/engine/base_model/config.py

@@ -1,3 +1,4 @@
+import abc
 import typing
 
 from fast_llm.config import Config, config_class
@@ -40,3 +41,12 @@ class BaseModelConfig(BaseModelArchitectureConfig):
 
     def get_architecture(self) -> BaseModelArchitectureConfig:
         return self.architecture_class.from_dict(self, strict=False)
+
+
+class Preprocessor(abc.ABC):
+    def preprocess_meta(self, kwargs: dict[str, typing.Any]) -> None:
+        pass
+
+    @abc.abstractmethod
+    def preprocess(self, batch, kwargs: dict[str, typing.Any]) -> None:
+        pass

fast_llm/layers/language_model/preprocessing.py

@@ -3,6 +3,7 @@
 
 import torch
 
+from fast_llm.engine.base_model.config import Preprocessor
 from fast_llm.engine.config_utils.tensor_space import DefaultDimNames, TensorDim, TensorSpace
 from fast_llm.layers.language_model.config import LanguageModelBaseConfig, LanguageModelKwargs
 from fast_llm.layers.transformer.config import TransformerKwargs
@@ -12,7 +13,7 @@
 logger = logging.getLogger(__name__)
 
 
-class PositionEmbeddingPreprocessor:
+class PositionEmbeddingPreprocessor(Preprocessor):
     _scalar_dim: TensorDim
     _rotary_embedding_frequencies: torch.Tensor
     _position_ids: torch.Tensor
@@ -29,7 +30,7 @@ def __init__(
         self._distributed_config = self._tensor_space.distributed_config
         self._scalar_dim = self._tensor_space.get_tensor_dim(DefaultDimNames.scalar)
 
-    def create_tensors(self, sequence_length: int) -> None:
+    def _create_tensors(self, sequence_length: int) -> None:
         if sequence_length <= self._tensor_cache_max_sequence_length:
             return
         self._tensor_cache_max_sequence_length = sequence_length
@@ -39,7 +40,8 @@ def create_tensors(self, sequence_length: int) -> None:
             0, sequence_length, device=self._tensor_space.distributed.device, dtype=torch.int64
         )
 
-    def preprocess(self, kwargs: dict[str, typing.Any]) -> None:
+    def preprocess(self, batch, kwargs: dict[str, typing.Any]) -> None:
+        self._create_tensors(kwargs[TransformerKwargs.sequence_length])
         sequence_k = kwargs[TransformerKwargs.sequence_k_dim].size
         sequence_q = kwargs[TransformerKwargs.sequence_q_dim].size
         if (sequence_lengths := kwargs.get(TransformerKwargs.sequence_lengths)) is not None:

fast_llm/layers/transformer/preprocessing.py

@@ -4,6 +4,7 @@
 
 import torch
 
+from fast_llm.engine.base_model.config import Preprocessor
 from fast_llm.engine.config_utils.tensor_space import DefaultDimNames, TensorDim, TensorSpace
 from fast_llm.functional.rotary import convert_rotary_complex_to_real
 from fast_llm.layers.transformer.config import (
@@ -129,7 +130,7 @@ def get_rotary_frequencies(
     return frequencies
 
 
-class RotaryEmbeddingPreprocessor:
+class RotaryEmbeddingPreprocessor(Preprocessor):
     _scalar_dim: TensorDim
     _kv_channels_dim: TensorDim
     _rotary_embedding_frequencies: torch.Tensor
@@ -149,7 +150,7 @@ def __init__(
         self._scalar_dim = self._tensor_space.get_tensor_dim(DefaultDimNames.scalar)
         self._kv_channels_dim = self._tensor_space.get_tensor_dim(TransformerDimNames.kv_channels)
 
-    def create_tensors(self, sequence_length: int) -> None:
+    def _create_tensors(self, sequence_length: int) -> None:
         if sequence_length <= self._tensor_cache_max_sequence_length:
             return
         self._tensor_cache_max_sequence_length = sequence_length
@@ -161,7 +162,8 @@ def create_tensors(self, sequence_length: int) -> None:
             device=self._tensor_space.distributed.device,
         )
 
-    def preprocess(self, kwargs: dict[str, typing.Any]) -> None:
+    def preprocess(self, batch, kwargs: dict[str, typing.Any]) -> None:
+        self._create_tensors(kwargs[TransformerKwargs.sequence_length])
         sequence_k = kwargs[TransformerKwargs.sequence_k_dim].size
         kwargs[TransformerKwargs.rotary_freq_q] = self._rotary_embedding_frequencies[
             :, sequence_k - kwargs[TransformerKwargs.sequence_q_dim].size : sequence_k
@@ -189,7 +191,7 @@ def preprocess_meta(self, kwargs: dict[str, typing.Any]) -> None:
         )
 
 
-class BackupAttentionPreprocessor:
+class BackupAttentionPreprocessor(Preprocessor):
     _scalar_dim: TensorDim
     _kv_channels_dim: TensorDim
     _rotary_embedding_frequencies: torch.Tensor
@@ -208,7 +210,7 @@ def __init__(
         assert not self._config.do_use_flash_attention(self._distributed_config)
         self._scalar_dim = self._tensor_space.get_tensor_dim(DefaultDimNames.scalar)
 
-    def create_tensors(self, sequence_length: int) -> None:
+    def _create_tensors(self, sequence_length: int) -> None:
         if sequence_length <= self._tensor_cache_max_sequence_length:
             return
         self._tensor_cache_max_sequence_length = sequence_length
@@ -228,7 +230,8 @@ def create_tensors(self, sequence_length: int) -> None:
             device=self._tensor_space.distributed.device,
         )
 
-    def preprocess(self, kwargs: dict[str, typing.Any]) -> None:
+    def preprocess(self, batch, kwargs: dict[str, typing.Any]) -> None:
+        self._create_tensors(kwargs[TransformerKwargs.sequence_length])
         sequence_k = kwargs[TransformerKwargs.sequence_k_dim].size
         sequence_q = kwargs[TransformerKwargs.sequence_q_dim].size
         kwargs[TransformerKwargs.attention_mask] = self._mask[
@@ -264,14 +267,14 @@ def preprocess_meta(self, kwargs: dict[str, typing.Any]) -> None:
         )
 
 
-class FlashAttnVarlenPreprocessor:
+class FlashAttnVarlenPreprocessor(Preprocessor):
     def __init__(self, config: TransformerConfig, tensor_space: TensorSpace):
         self._config = config
         self._tensor_space = tensor_space
         self._distributed_config = self._tensor_space.distributed_config
         assert self._config.do_use_flash_attention(self._distributed_config)
 
-    def preprocess(self, kwargs: dict[str, typing.Any]) -> None:
+    def preprocess(self, batch, kwargs: dict[str, typing.Any]) -> None:
         """
         Prepares cu_seqlens_q and cu_seqlens_k for flash_attn_varlen_func:
         https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_interface.py#L1375
@@ -281,7 +284,9 @@ def preprocess(self, kwargs: dict[str, typing.Any]) -> None:
         also contain previous tokens from the first document in micro-sequence.
         We use individual sequence lengths of each document to (optionally) find the micro-sequences in the batch and compute the cumulative lengths.
         """
-        sequence_lengths = kwargs.get(TransformerKwargs.sequence_lengths)
+        if TransformerKwargs.sequence_lengths not in kwargs:
+            return
+        sequence_lengths = kwargs[TransformerKwargs.sequence_lengths]
         sequence_k = kwargs[TransformerKwargs.sequence_k_dim].size
         sequence_q = kwargs[TransformerKwargs.sequence_q_dim].size
         if sequence_q < kwargs[TransformerKwargs.sequence_length]:

fast_llm/models/gpt/model.py

@@ -5,6 +5,7 @@
 
 from fast_llm.data.data.gpt.data import GPTBatch
 from fast_llm.engine.base_model.base_model import BaseModel, Layer, LossDef
+from fast_llm.engine.base_model.config import Preprocessor
 from fast_llm.engine.config_utils.tensor_space import TensorDim
 from fast_llm.engine.distributed.config import DistributedConfig, DistributedDimNames, PhaseType
 from fast_llm.engine.distributed.distributed import Distributed
@@ -58,18 +59,17 @@ def __init__(
             for param in self.parameters():
                 Assert.custom(isinstance, param, ParameterMeta)
                 param.init_parameter = get_init_megatron(param, self._config.transformer)  # Noqa
+        self._preprocessors: list[Preprocessor] = []
         if self._config.use_absolute_position_embeddings:
-            self._position_embedding_preprocessor = PositionEmbeddingPreprocessor(self._config, self._tensor_space)
+            self._preprocessors.append(PositionEmbeddingPreprocessor(self._config, self._tensor_space))
         if self._config.transformer.rotary.enabled:
-            self._rotary_embedding_preprocessor = RotaryEmbeddingPreprocessor(
-                self._config.transformer.rotary, self._tensor_space
-            )
-        if not self._use_flash_attention:
-            self._backup_attention_preprocessor = BackupAttentionPreprocessor(
-                self._config.transformer, self._tensor_space
+            self._preprocessors.append(
+                RotaryEmbeddingPreprocessor(self._config.transformer.rotary, self._tensor_space)
             )
+        if self._use_flash_attention:
+            self._preprocessors.append(FlashAttnVarlenPreprocessor(self._config.transformer, self._tensor_space))
         else:
-            self._flash_varlen_preprocessor = FlashAttnVarlenPreprocessor(self._config.transformer, self._tensor_space)
+            self._preprocessors.append(BackupAttentionPreprocessor(self._config.transformer, self._tensor_space))
 
     def get_output_layers(self) -> list[Layer]:
         return [
@@ -207,12 +207,8 @@ def preprocess_meta(
                 kwargs[LanguageModelKwargs.labels] = TensorMeta.from_dims(
                     hidden_dims[:2], tensor_name="labels", dtype=torch.int64
                 )
-            if self._config.use_absolute_position_embeddings:
-                self._position_embedding_preprocessor.preprocess_meta(kwargs)
-            if self._config.transformer.rotary.enabled:
-                self._rotary_embedding_preprocessor.preprocess_meta(kwargs)
-            if not self._use_flash_attention:
-                self._backup_attention_preprocessor.preprocess_meta(kwargs)
+            for preprocessor in self._preprocessors:
+                preprocessor.preprocess_meta(kwargs)
             preprocessed_meta.append((tokens, kwargs))
 
         return preprocessed_meta
@@ -235,7 +231,6 @@ def preprocess(
         _, common_kwargs = preprocessed_meta[0]
         sequence_q = common_kwargs[TransformerKwargs.sequence_q_dim].size
         sequence_first = common_kwargs[TransformerKwargs.sequence_first]
-        sequence_length = common_kwargs[TransformerKwargs.sequence_length]
 
         batch.token_ids = batch.token_ids.to(
             device=self._tensor_space.distributed.device,
@@ -246,13 +241,6 @@ def preprocess(
             # Move the sequence dimension first to make sequence parallel ops more efficient.
             batch.token_ids = batch.token_ids.transpose(0, 1).contiguous()
 
-        if self._config.use_absolute_position_embeddings:
-            self._position_embedding_preprocessor.create_tensors(sequence_length)
-        if self._config.transformer.rotary.enabled:
-            self._rotary_embedding_preprocessor.create_tensors(sequence_length)
-        if not self._use_flash_attention:
-            self._backup_attention_preprocessor.create_tensors(sequence_length)
-
         preprocessed = []
         presents = None
         for i, (tokens_meta, kwargs_meta) in enumerate(preprocessed_meta):
@@ -264,8 +252,6 @@ def preprocess(
                 tokens = batch.token_ids[:, sequence_k - sequence_q : sequence_k].contiguous()
             if batch.sequence_lengths is not None:
                 kwargs_meta[TransformerKwargs.sequence_lengths] = batch.sequence_lengths
-                if self._use_flash_attention:
-                    self._flash_varlen_preprocessor.preprocess(kwargs_meta)
 
             # TODO: Add pasts/presents to meta input?
             # Use lists as pointers so `past_key_values` is populated during the previous micro_sequence.
@@ -300,12 +286,8 @@ def preprocess(
                                 else:
                                     labels[i, start : end + 1] = -100
                 kwargs[LanguageModelKwargs.labels] = labels
-            if self._config.use_absolute_position_embeddings:
-                self._position_embedding_preprocessor.preprocess(kwargs)
-            if self._config.transformer.rotary.enabled:
-                self._rotary_embedding_preprocessor.preprocess(kwargs)
-            if not self._use_flash_attention:
-                self._backup_attention_preprocessor.preprocess(kwargs)
+            for preprocessor in self._preprocessors:
+                preprocessor.preprocess(tokens, kwargs)
             preprocessed.append((tokens, kwargs))
 
         return preprocessed
@@ -379,6 +361,10 @@ def loss_defs(self) -> list[LossDef]:
             )
         return loss_defs
 
+    def add_preprocessor(self, preprocessor: Preprocessor):
+        assert not self._is_setup
+        self._preprocessors.append(preprocessor)
+
 
 class GPTModel[ConfigType: GPTModelConfig](FastLLMModel[ConfigType]):
     config_class: typing.ClassVar[type[GPTModelConfig]] = GPTModelConfig

tests/test_attention.py

@@ -84,6 +84,6 @@ def test_varlen_preprocessor():
             TransformerKwargs.sequence_length: sequence_length,
             TransformerKwargs.sequence_lengths: sequence_lengths,
         }
-        varlen_preprocessor.preprocess(kwargs)
+        varlen_preprocessor.preprocess(None, kwargs)
         Assert.all_equal(kwargs[TransformerKwargs.cu_seqlens_q], cumulative_sequences_q[micro_seq_idx])
         Assert.all_equal(kwargs[TransformerKwargs.cu_seqlens_k], cumulative_sequences_k[micro_seq_idx])