Add MPTGenerator (#296)

* Add `MPTGenerator`

This is pretty straightforward, except that I had to add a workaround
because we don't support weight tying yet.

* No special JIT handling in causal LMs

The special handling of JIT'ed code has not been necessary since the
model outputs are also tuples.

* Whitespace fix

Co-authored-by: Madeesh Kannan <shadeMe@users.noreply.github.com>

---------

Co-authored-by: Madeesh Kannan <shadeMe@users.noreply.github.com>

Author: danieldk

README.md

@@ -54,6 +54,7 @@ Generator wrappers:
 - Dolly v2
 - Falcon
 - Llama 1/2
+- MPT
 
 All types of models can be loaded from Huggingface Hub.
 

curated_transformers/generation/__init__.py

@@ -15,6 +15,7 @@
     TopPTransform,
     VocabMaskTransform,
 )
+from .mpt import MPTGenerator
 from .stop_conditions import (
     CompoundStopCondition,
     EndOfSequenceCondition,
@@ -39,6 +40,7 @@
     "LlamaGenerator",
     "LogitsTransform",
     "MaxGeneratedPiecesCondition",
+    "MPTGenerator",
     "SampleGeneratorConfig",
     "StopCondition",
     "StringGenerator",

curated_transformers/generation/auto_generator.py

@@ -4,18 +4,20 @@
 
 from ..models.auto_model import AutoModel
 from ..quantization.bnb.config import BitsAndBytesConfig
-from .default_generator import DefaultGenerator
 from .dolly_v2 import DollyV2Generator
 from .falcon import FalconGenerator
 from .generator_wrapper import GeneratorWrapper
 from .hf_hub import FromHFHub
+from .llama import LlamaGenerator
+from .mpt import MPTGenerator
 
 # For the time being, we enable support for a generator on a case-by-case basis.
 # In the future we might defer all unknown generators to DefaultGenerator.
 GENERATOR_MAP: Dict[str, Type[FromHFHub]] = {
     "dolly-v2": DollyV2Generator,
     "falcon": FalconGenerator,
-    "llama": DefaultGenerator,
+    "llama": LlamaGenerator,
+    "mpt": MPTGenerator,
 }
 
 

curated_transformers/generation/mpt.py

@@ -0,0 +1,26 @@
+from typing import TypeVar
+
+from ..models.llama import LlamaCausalLM
+from ..tokenizers.tokenizer import Tokenizer
+from .default_generator import DefaultGenerator
+from .hf_hub import FromHFHub
+
+
+class MPTGenerator(DefaultGenerator, FromHFHub):
+    """
+    Generator for MPT model variants.
+    """
+
+    def __init__(self, tokenizer: Tokenizer, causal_lm: LlamaCausalLM):
+        """
+        Construct an MPT generator.
+
+        :param tokenizer:
+            An MPT tokenizer.
+        :param causal_lm:
+            An MPT causal language model.
+        """
+        super().__init__(
+            tokenizer,
+            causal_lm,
+        )

curated_transformers/models/mpt/_hf.py

@@ -61,6 +61,10 @@ def convert_hf_state_dict(cls, params: Mapping[str, Tensor]) -> Mapping[str, Ten
 
     out = {}
     for name, parameter in stripped_params.items():
+        # Input and output embeddings are tied in MPT.
+        if "lm_head" in name:
+            continue
+
         name = name.replace("transformer", "decoder")
         name = name.replace("blocks", "layers")
 
@@ -80,7 +84,6 @@ def convert_hf_state_dict(cls, params: Mapping[str, Tensor]) -> Mapping[str, Ten
 
         # Embeddings
         name = re.sub(r"wte\.", r"embeddings.piece_embeddings.", name)
-        name = re.sub(r"lm_head\.", r"output_embeddings.", name)
 
         out[name] = parameter
 

curated_transformers/models/mpt/causal_lm.py

@@ -1,11 +1,15 @@
-from typing import Any, Mapping, Optional, Set, Type, TypeVar
+from typing import Any, List, Mapping, Optional, Set, Type, TypeVar
 
 import torch
+import torch.nn.functional as F
 from torch import Tensor
-from torch.nn import Linear
+from torch.nn import Embedding
 
+from ...layers.attention import AttentionMask
+from ...layers.cache import KeyValueCache
 from ...quantization import Quantizable
 from ..hf_hub import FromHFHub
+from ..output import CausalLMOutputWithCache
 from ..transformer import TransformerCausalLM
 from ._hf import convert_hf_config, convert_hf_state_dict
 from .config import MPTConfig
@@ -38,11 +42,45 @@ def __init__(
         super().__init__()
 
         self.decoder = MPTDecoder(config, device=device)
-        self.output_embeddings = Linear(
-            in_features=config.layer.feedforward.hidden_width,
-            out_features=config.embedding.n_pieces,
-            bias=False,
-            device=device,
+
+        # Once we have proper support for tied weights, we will do something like:
+        #
+        # self.output_embeddings = Linear(
+        #    in_features=config.layer.feedforward.hidden_width,
+        #    out_features=config.embedding.n_pieces,
+        #    bias=False,
+        #    device=device,
+        # )
+        # self.output_embeddings.weights = self.decoder.embeddings.piece_embeddings.weights
+        #
+        # For now we'll work around this by using the piece embeddings directly.
+
+    def forward(
+        self,
+        piece_ids: Tensor,
+        attention_mask: AttentionMask,
+        cache: Optional[List[KeyValueCache]] = None,
+        positions: Optional[Tensor] = None,
+        store_cache: bool = False,
+    ) -> CausalLMOutputWithCache[KeyValueCache]:
+        # TODO: remove this forward method once we support weight tying.
+
+        decoder_output = self.decoder(
+            piece_ids,
+            attention_mask,
+            cache=cache,
+            store_cache=store_cache,
+            positions=positions,
+        )
+
+        assert isinstance(self.decoder.embeddings.piece_embeddings, Embedding)
+        output_embeddings = self.decoder.embeddings.piece_embeddings.weight
+
+        logits = F.linear(decoder_output.last_hidden_layer_state, output_embeddings)
+        return CausalLMOutputWithCache(
+            all_outputs=decoder_output.all_outputs,
+            cache=decoder_output.cache,
+            logits=logits,
         )
 
     @classmethod

curated_transformers/models/transformer.py

@@ -92,16 +92,12 @@ def forward(
             store_cache=store_cache,
             positions=positions,
         )
-        if torch.jit.is_tracing():
-            logits = self.output_embeddings(decoder_output[0][-1])
-            return decoder_output + (logits,)  # type: ignore[return-value]
-        else:
-            logits = self.output_embeddings(decoder_output.last_hidden_layer_state)
-            return CausalLMOutputWithCache(
-                all_outputs=decoder_output.all_outputs,
-                cache=decoder_output.cache,
-                logits=logits,
-            )
+        logits = self.output_embeddings(decoder_output.last_hidden_layer_state)
+        return CausalLMOutputWithCache(
+            all_outputs=decoder_output.all_outputs,
+            cache=decoder_output.cache,
+            logits=logits,
+        )
 
 
 class TransformerEncoder(EncoderModule):

curated_transformers/tests/generation/test_auto_generator.py

@@ -1,17 +1,33 @@
 import pytest
 
 from curated_transformers.generation import AutoGenerator
-from curated_transformers.generation.default_generator import DefaultGenerator
 from curated_transformers.generation.dolly_v2 import DollyV2Generator
 from curated_transformers.generation.falcon import FalconGenerator
+from curated_transformers.generation.llama import LlamaGenerator
+from curated_transformers.generation.mpt import MPTGenerator
 
 
 @pytest.mark.slow
 def test_auto_generator():
     model_causallm_map = {
         "databricks/dolly-v2-3b": DollyV2Generator,
         "tiiuae/falcon-7b": FalconGenerator,
-        "openlm-research/open_llama_3b": DefaultGenerator,
+        "openlm-research/open_llama_3b": LlamaGenerator,
+    }
+
+    for name, generator_cls in model_causallm_map.items():
+        generator = AutoGenerator.from_hf_hub(name=name)
+        assert isinstance(generator, generator_cls)
+
+    with pytest.raises(ValueError, match="Unsupported generator"):
+        AutoGenerator.from_hf_hub(name="trl-internal-testing/tiny-random-GPT2Model")
+
+
+@pytest.mark.hf_head
+@pytest.mark.slow
+def test_auto_generator_hf_head():
+    model_causallm_map = {
+        "mosaicml/mpt-7b": MPTGenerator,
     }
 
     for name, generator_cls in model_causallm_map.items():

docs/source/generation.rst

@@ -51,6 +51,11 @@ These classes provide the interface for performing text generation using causal
    :special-members: __call__
    :show-inheritance:   
 
+.. autoclass:: curated_transformers.generation.MPTGenerator
+   :members:
+   :inherited-members:
+   :special-members: __call__
+   :show-inheritance:
 
 Downloading
 -----------

docs/source/index.rst

@@ -60,6 +60,7 @@ Generator wrappers:
 - Dolly v2
 - Falcon
 - Llama 1/2
+- MPT
 
 All types of models can be loaded from Hugging Face Hub.