From 01e8902aea089cf29faf6008998660db3e6490cf Mon Sep 17 00:00:00 2001 From: Madeesh Kannan Date: Mon, 18 Sep 2023 11:19:38 +0200 Subject: [PATCH] Persist model config to aid serialization (#328) * Persist model config and expose it to the API * Add test * Exclude tests with JIT * Add protocol as a config typevar bound --- .../generation/default_generator.py | 7 ++-- curated_transformers/generation/generator.py | 6 +-- curated_transformers/models/albert/encoder.py | 4 +- curated_transformers/models/auto_model.py | 19 ++++----- curated_transformers/models/bert/encoder.py | 4 +- curated_transformers/models/config.py | 10 ++++- .../models/falcon/causal_lm.py | 4 +- curated_transformers/models/falcon/decoder.py | 4 +- .../models/gpt_neox/causal_lm.py | 4 +- .../models/gpt_neox/decoder.py | 4 +- .../models/llama/causal_lm.py | 4 +- curated_transformers/models/llama/decoder.py | 4 +- curated_transformers/models/module.py | 40 +++++++++++++++++-- curated_transformers/models/mpt/causal_lm.py | 4 +- curated_transformers/models/mpt/decoder.py | 4 +- .../models/roberta/encoder.py | 4 +- curated_transformers/models/transformer.py | 10 ++--- curated_transformers/tests/models/util.py | 17 ++++++++ docs/source/causal-lm.rst | 2 + docs/source/decoders.rst | 2 + docs/source/encoders.rst | 2 + 21 files changed, 112 insertions(+), 47 deletions(-) diff --git a/curated_transformers/generation/default_generator.py b/curated_transformers/generation/default_generator.py index 16d6c9b6..358161d3 100644 --- a/curated_transformers/generation/default_generator.py +++ b/curated_transformers/generation/default_generator.py @@ -1,10 +1,11 @@ import dataclasses -from typing import List, Optional, Type, TypeVar +from typing import Any, Generic, List, Optional, Type, TypeVar import torch from ..models.auto_model import AutoCausalLM from ..models.module import CausalLMModule +from ..models.output import CacheT from ..quantization.bnb.config import BitsAndBytesConfig from ..tokenizers.auto_tokenizer import AutoTokenizer from ..tokenizers.chunks import InputChunks, TextChunk @@ -19,7 +20,7 @@ Self = TypeVar("Self", bound="DefaultGenerator") -class DefaultGenerator(GeneratorWrapper, FromHFHub): +class DefaultGenerator(Generic[CacheT], GeneratorWrapper, FromHFHub): """ Generator wrapper for models that do not need specific prompting. """ @@ -27,7 +28,7 @@ class DefaultGenerator(GeneratorWrapper, FromHFHub): def __init__( self, tokenizer: TokenizerBase, - causal_lm: CausalLMModule, + causal_lm: CausalLMModule[Any, CacheT], default_config: Optional[GeneratorConfig] = None, ): """ diff --git a/curated_transformers/generation/generator.py b/curated_transformers/generation/generator.py index 3d1b2d7b..c7bbd927 100644 --- a/curated_transformers/generation/generator.py +++ b/curated_transformers/generation/generator.py @@ -1,4 +1,4 @@ -from typing import Generic, Iterator, List, Optional, Tuple +from typing import Any, Generic, Iterator, List, Optional, Tuple import torch from torch import Tensor @@ -17,9 +17,9 @@ class Generator(Generic[CacheT]): Generator base class for causal language models. """ - model: CausalLMModule[CacheT] + model: CausalLMModule[Any, CacheT] - def __init__(self, model: CausalLMModule[CacheT]): + def __init__(self, model: CausalLMModule[Any, CacheT]): """ Construct a generator. diff --git a/curated_transformers/models/albert/encoder.py b/curated_transformers/models/albert/encoder.py index 51c801fa..ffaf6489 100644 --- a/curated_transformers/models/albert/encoder.py +++ b/curated_transformers/models/albert/encoder.py @@ -21,7 +21,7 @@ Self = TypeVar("Self", bound="ALBERTEncoder") -class ALBERTEncoder(EncoderModule, FromHFHub): +class ALBERTEncoder(EncoderModule[ALBERTConfig], FromHFHub): """ ALBERT (`Lan et al., 2022`_) encoder. @@ -39,7 +39,7 @@ def __init__(self, config: ALBERTConfig, *, device: Optional[torch.device] = Non :returns: The encoder. """ - super().__init__() + super().__init__(config) self.max_seq_len = config.model_max_length self.n_hidden_layers = config.layer.n_hidden_layers diff --git a/curated_transformers/models/auto_model.py b/curated_transformers/models/auto_model.py index b57f5181..ef7d3be6 100644 --- a/curated_transformers/models/auto_model.py +++ b/curated_transformers/models/auto_model.py @@ -11,6 +11,7 @@ from .albert import ALBERTEncoder from .bert import BERTEncoder from .camembert import CamemBERTEncoder +from .config import ConfigDataclass from .falcon import FalconCausalLM, FalconDecoder from .gpt_neox import GPTNeoXCausalLM, GPTNeoXDecoder from .hf_hub import FromHFHub @@ -184,7 +185,7 @@ def from_hf_hub_to_cache( module_cls.from_hf_hub_to_cache(name=name, revision=revision) -class AutoEncoder(AutoModel[EncoderModule]): +class AutoEncoder(AutoModel[EncoderModule[ConfigDataclass]]): """ Encoder model loaded from the Hugging Face Model Hub. """ @@ -206,7 +207,7 @@ def from_fsspec( fsspec_args: Optional[Dict[str, Any]] = None, device: Optional[torch.device] = None, quantization_config: Optional[BitsAndBytesConfig] = None, - ) -> EncoderModule: + ) -> EncoderModule[ConfigDataclass]: encoder = cls._instantiate_model_from_fsspec( fs, model_path, fsspec_args, device, quantization_config ) @@ -221,7 +222,7 @@ def from_hf_hub( revision: str = "main", device: Optional[torch.device] = None, quantization_config: Optional[BitsAndBytesConfig] = None, - ) -> EncoderModule: + ) -> EncoderModule[ConfigDataclass]: encoder = cls._instantiate_model_from_hf_hub( name, revision, device, quantization_config ) @@ -229,7 +230,7 @@ def from_hf_hub( return encoder -class AutoDecoder(AutoModel[DecoderModule]): +class AutoDecoder(AutoModel[DecoderModule[ConfigDataclass, KeyValueCache]]): """ Decoder module loaded from the Hugging Face Model Hub. """ @@ -252,7 +253,7 @@ def from_fsspec( fsspec_args: Optional[Dict[str, Any]] = None, device: Optional[torch.device] = None, quantization_config: Optional[BitsAndBytesConfig] = None, - ) -> DecoderModule: + ) -> DecoderModule[ConfigDataclass, KeyValueCache]: decoder = cls._instantiate_model_from_fsspec( fs, model_path, fsspec_args, device, quantization_config ) @@ -267,7 +268,7 @@ def from_hf_hub( revision: str = "main", device: Optional[torch.device] = None, quantization_config: Optional[BitsAndBytesConfig] = None, - ) -> DecoderModule: + ) -> DecoderModule[ConfigDataclass, KeyValueCache]: decoder = cls._instantiate_model_from_hf_hub( name, revision, device, quantization_config ) @@ -275,7 +276,7 @@ def from_hf_hub( return decoder -class AutoCausalLM(AutoModel[CausalLMModule[KeyValueCache]]): +class AutoCausalLM(AutoModel[CausalLMModule[ConfigDataclass, KeyValueCache]]): """ Causal LM model loaded from the Hugging Face Model Hub. """ @@ -298,7 +299,7 @@ def from_fsspec( fsspec_args: Optional[Dict[str, Any]] = None, device: Optional[torch.device] = None, quantization_config: Optional[BitsAndBytesConfig] = None, - ) -> CausalLMModule[KeyValueCache]: + ) -> CausalLMModule[ConfigDataclass, KeyValueCache]: causal_lm = cls._instantiate_model_from_fsspec( fs, model_path, fsspec_args, device, quantization_config ) @@ -313,7 +314,7 @@ def from_hf_hub( revision: str = "main", device: Optional[torch.device] = None, quantization_config: Optional[BitsAndBytesConfig] = None, - ) -> CausalLMModule[KeyValueCache]: + ) -> CausalLMModule[ConfigDataclass, KeyValueCache]: causal_lm = cls._instantiate_model_from_hf_hub( name, revision, device, quantization_config ) diff --git a/curated_transformers/models/bert/encoder.py b/curated_transformers/models/bert/encoder.py index c8048f93..0dc55d1b 100644 --- a/curated_transformers/models/bert/encoder.py +++ b/curated_transformers/models/bert/encoder.py @@ -24,7 +24,7 @@ Self = TypeVar("Self", bound="BERTEncoder") -class BERTEncoder(TransformerEncoder, FromHFHub): +class BERTEncoder(TransformerEncoder[BERTConfig], FromHFHub): """ BERT (`Devlin et al., 2018`_) encoder. @@ -42,7 +42,7 @@ def __init__(self, config: BERTConfig, *, device: Optional[torch.device] = None) :returns: The encoder. """ - super().__init__() + super().__init__(config) self.embeddings = TransformerEmbeddings( dropouts=EmbeddingDropouts( diff --git a/curated_transformers/models/config.py b/curated_transformers/models/config.py index 4990e19e..189ab254 100644 --- a/curated_transformers/models/config.py +++ b/curated_transformers/models/config.py @@ -1,9 +1,17 @@ from dataclasses import dataclass -from typing import Optional +from typing import ClassVar, Optional, Protocol from ..layers.activations import Activation +class ConfigDataclass(Protocol): + """ + Protocol that describes a config data class. + """ + + __dataclass_fields__: ClassVar[dict] + + @dataclass class RotaryEmbeddingConfig: """ diff --git a/curated_transformers/models/falcon/causal_lm.py b/curated_transformers/models/falcon/causal_lm.py index 34228f06..ae570d48 100644 --- a/curated_transformers/models/falcon/causal_lm.py +++ b/curated_transformers/models/falcon/causal_lm.py @@ -15,7 +15,7 @@ Self = TypeVar("Self", bound="FalconCausalLM") -class FalconCausalLM(TransformerCausalLM, FromHFHub, Quantizable): +class FalconCausalLM(TransformerCausalLM[FalconConfig], FromHFHub, Quantizable): """ Falcon (`Penedo et al., 2019`_) causal language model. @@ -35,7 +35,7 @@ def __init__( :returns: The causal LM. """ - super().__init__() + super().__init__(config) self.decoder = FalconDecoder(config, device=device) self.output_embeddings = Linear( diff --git a/curated_transformers/models/falcon/decoder.py b/curated_transformers/models/falcon/decoder.py index cec6f5b9..0a16154e 100644 --- a/curated_transformers/models/falcon/decoder.py +++ b/curated_transformers/models/falcon/decoder.py @@ -31,7 +31,7 @@ Self = TypeVar("Self", bound="FalconDecoder") -class FalconDecoder(TransformerDecoder, FromHFHub): +class FalconDecoder(TransformerDecoder[FalconConfig], FromHFHub): """ Falcon (`Penedo et al., 2019`_) decoder. @@ -51,7 +51,7 @@ def __init__( :returns: The decoder. """ - super().__init__() + super().__init__(config) self.embeddings = TransformerEmbeddings( dropouts=EmbeddingDropouts( diff --git a/curated_transformers/models/gpt_neox/causal_lm.py b/curated_transformers/models/gpt_neox/causal_lm.py index 7a1b9c2f..cc6f9bb7 100644 --- a/curated_transformers/models/gpt_neox/causal_lm.py +++ b/curated_transformers/models/gpt_neox/causal_lm.py @@ -15,7 +15,7 @@ Self = TypeVar("Self", bound="GPTNeoXCausalLM") -class GPTNeoXCausalLM(TransformerCausalLM, FromHFHub, Quantizable): +class GPTNeoXCausalLM(TransformerCausalLM[GPTNeoXConfig], FromHFHub, Quantizable): """ GPT-NeoX (`Black et al., 2022`_) causal language model. @@ -35,7 +35,7 @@ def __init__( :returns: The causal LM. """ - super().__init__() + super().__init__(config) self.decoder = GPTNeoXDecoder(config, device=device) self.output_embeddings = Linear( diff --git a/curated_transformers/models/gpt_neox/decoder.py b/curated_transformers/models/gpt_neox/decoder.py index 39c06eea..d04189ca 100644 --- a/curated_transformers/models/gpt_neox/decoder.py +++ b/curated_transformers/models/gpt_neox/decoder.py @@ -25,7 +25,7 @@ Self = TypeVar("Self", bound="GPTNeoXDecoder") -class GPTNeoXDecoder(TransformerDecoder, FromHFHub): +class GPTNeoXDecoder(TransformerDecoder[GPTNeoXConfig], FromHFHub): """ GPT-NeoX (`Black et al., 2022`_) decoder. @@ -45,7 +45,7 @@ def __init__( :returns: The decoder. """ - super().__init__() + super().__init__(config) self.embeddings = TransformerEmbeddings( dropouts=EmbeddingDropouts( diff --git a/curated_transformers/models/llama/causal_lm.py b/curated_transformers/models/llama/causal_lm.py index 58e971e1..135178cc 100644 --- a/curated_transformers/models/llama/causal_lm.py +++ b/curated_transformers/models/llama/causal_lm.py @@ -15,7 +15,7 @@ Self = TypeVar("Self", bound="LlamaCausalLM") -class LlamaCausalLM(TransformerCausalLM, FromHFHub, Quantizable): +class LlamaCausalLM(TransformerCausalLM[LlamaConfig], FromHFHub, Quantizable): """ Llama (`Touvron et al., 2023 [a]`_, `Touvron et al., 2023 [b]`_) causal language model. @@ -36,7 +36,7 @@ def __init__( :returns: The causal LM. """ - super().__init__() + super().__init__(config) self.decoder = LlamaDecoder(config, device=device) self.output_embeddings = Linear( diff --git a/curated_transformers/models/llama/decoder.py b/curated_transformers/models/llama/decoder.py index d26c9950..072a2cc3 100644 --- a/curated_transformers/models/llama/decoder.py +++ b/curated_transformers/models/llama/decoder.py @@ -26,7 +26,7 @@ Self = TypeVar("Self", bound="LlamaDecoder") -class LlamaDecoder(TransformerDecoder, FromHFHub): +class LlamaDecoder(TransformerDecoder[LlamaConfig], FromHFHub): """ Llama (`Touvron et al., 2023 [a]`_, `Touvron et al., 2023 [b]`_) decoder. @@ -47,7 +47,7 @@ def __init__( :returns: The decoder. """ - super().__init__() + super().__init__(config) self.embeddings = TransformerEmbeddings( dropouts=EmbeddingDropouts( diff --git a/curated_transformers/models/module.py b/curated_transformers/models/module.py index eaaa2bb9..b5bcf91d 100644 --- a/curated_transformers/models/module.py +++ b/curated_transformers/models/module.py @@ -1,18 +1,44 @@ from abc import abstractmethod -from typing import Generic, List, Optional +from typing import Generic, List, Optional, TypeVar from torch import Tensor from torch.nn import Module from ..layers.attention import AttentionMask +from .config import ConfigDataclass from .output import CacheT, CausalLMOutputWithCache, ModelOutput, ModelOutputWithCache +ConfigT = TypeVar("ConfigT", bound=ConfigDataclass) -class CausalLMModule(Generic[CacheT], Module): + +class TransformerModule(Generic[ConfigT], Module): + """ + Base class for transformer modules. + """ + + _config: ConfigT + + def __init__(self, config: ConfigT): + super().__init__() + + self._config = config + + @property + def config(self) -> ConfigT: + """ + Returns the model's configuration. + """ + return self._config + + +class CausalLMModule(Generic[ConfigT, CacheT], TransformerModule[ConfigT]): """ Base class for causal language model modules. """ + def __init__(self, config: ConfigT): + super().__init__(config) + @abstractmethod def forward( self, @@ -51,11 +77,14 @@ def forward( raise NotImplementedError -class DecoderModule(Generic[CacheT], Module): +class DecoderModule(Generic[ConfigT, CacheT], TransformerModule[ConfigT]): """ Base class for decoder modules. """ + def __init__(self, config: ConfigT): + super().__init__(config) + @abstractmethod def forward( self, @@ -94,11 +123,14 @@ def forward( raise NotImplementedError -class EncoderModule(Module): +class EncoderModule(Generic[ConfigT], TransformerModule[ConfigT]): """ Base class for encoder modules. """ + def __init__(self, config: ConfigT): + super().__init__(config) + @abstractmethod def forward( self, diff --git a/curated_transformers/models/mpt/causal_lm.py b/curated_transformers/models/mpt/causal_lm.py index 8e850420..0006c064 100644 --- a/curated_transformers/models/mpt/causal_lm.py +++ b/curated_transformers/models/mpt/causal_lm.py @@ -19,7 +19,7 @@ Self = TypeVar("Self", bound="MPTCausalLM") -class MPTCausalLM(TransformerCausalLM, FromHFHub, Quantizable): +class MPTCausalLM(TransformerCausalLM[MPTConfig], FromHFHub, Quantizable): """ `MosaicML MPT`_ causal language model. @@ -39,7 +39,7 @@ def __init__( :returns: The causal LM. """ - super().__init__() + super().__init__(config) self.decoder = MPTDecoder(config, device=device) diff --git a/curated_transformers/models/mpt/decoder.py b/curated_transformers/models/mpt/decoder.py index 5a535065..1207dc5c 100644 --- a/curated_transformers/models/mpt/decoder.py +++ b/curated_transformers/models/mpt/decoder.py @@ -29,7 +29,7 @@ Self = TypeVar("Self", bound="MPTDecoder") -class MPTDecoder(TransformerDecoder, FromHFHub): +class MPTDecoder(TransformerDecoder[MPTConfig], FromHFHub): """ `MosaicML MPT`_ decoder. @@ -49,7 +49,7 @@ def __init__( :returns: The decoder. """ - super().__init__() + super().__init__(config) self.embeddings = TransformerEmbeddings( dropouts=EmbeddingDropouts( diff --git a/curated_transformers/models/roberta/encoder.py b/curated_transformers/models/roberta/encoder.py index dad33a59..11b06c1d 100644 --- a/curated_transformers/models/roberta/encoder.py +++ b/curated_transformers/models/roberta/encoder.py @@ -24,7 +24,7 @@ Self = TypeVar("Self", bound="RoBERTaEncoder") -class RoBERTaEncoder(TransformerEncoder, FromHFHub): +class RoBERTaEncoder(TransformerEncoder[RoBERTaConfig], FromHFHub): """ RoBERTa (`Liu et al., 2019`_) encoder. @@ -42,7 +42,7 @@ def __init__(self, config: RoBERTaConfig, *, device: Optional[torch.device] = No :returns: The encoder. """ - super().__init__() + super().__init__(config) self.embeddings = RoBERTaEmbeddings( dropouts=EmbeddingDropouts( diff --git a/curated_transformers/models/transformer.py b/curated_transformers/models/transformer.py index 44fcf1ad..54d635fa 100644 --- a/curated_transformers/models/transformer.py +++ b/curated_transformers/models/transformer.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from typing import Generic, List, Optional import torch from torch import Tensor @@ -6,11 +6,11 @@ from ..layers.attention import AttentionMask from ..layers.cache import KeyValueCache -from .module import CausalLMModule, DecoderModule, EncoderModule +from .module import CausalLMModule, ConfigT, DecoderModule, EncoderModule from .output import CausalLMOutputWithCache, ModelOutput, ModelOutputWithCache -class TransformerDecoder(DecoderModule): +class TransformerDecoder(Generic[ConfigT], DecoderModule[ConfigT, KeyValueCache]): """ Transformer decoder (`Vaswani et al., 2017`_) base class. @@ -64,7 +64,7 @@ def forward( ) -class TransformerCausalLM(CausalLMModule[KeyValueCache]): +class TransformerCausalLM(Generic[ConfigT], CausalLMModule[ConfigT, KeyValueCache]): """ Transformer causal LM (`Vaswani et al., 2017`_) base class. @@ -100,7 +100,7 @@ def forward( ) -class TransformerEncoder(EncoderModule): +class TransformerEncoder(Generic[ConfigT], EncoderModule[ConfigT]): """ Transformer encoder (`Vaswani et al., 2017`_) base class. diff --git a/curated_transformers/tests/models/util.py b/curated_transformers/tests/models/util.py index aaf4adbd..d9a12926 100644 --- a/curated_transformers/tests/models/util.py +++ b/curated_transformers/tests/models/util.py @@ -13,6 +13,7 @@ CausalLMModule, DecoderModule, EncoderModule, + TransformerModule, ) from curated_transformers.models.output import ModelOutput, ModelOutputWithCache @@ -147,6 +148,7 @@ def assert_decoder_output_equals_hf( with_mask: bool = True, jit_method: JITMethod = JITMethod.Disable, with_torch_sdp=False, + check_config=True, ): orig_model = model_class.from_hf_hub( name=model_name, revision=model_revision, device=torch_device @@ -193,6 +195,9 @@ def assert_decoder_output_equals_hf( orig_model, hf_model, torch_device, atol, rtol, jit_method ) + if check_config and jit_method == JITMethod.Disable: + assert_model_config(model, Y) + def assert_encoder_output_equals_hf( model_class: Type[FromHFHub], @@ -204,6 +209,7 @@ def assert_encoder_output_equals_hf( jit_method: JITMethod = JITMethod.Disable, with_fsspec: bool = False, with_torch_sdp: bool = False, + check_config=True, ): if with_fsspec: orig_model = model_class.from_fsspec( @@ -240,6 +246,9 @@ def assert_encoder_output_equals_hf( orig_model, hf_model, torch_device, atol, rtol, jit_method ) + if check_config and jit_method == JITMethod.Disable: + assert_model_config(model, Y) + def assert_decoder_with_cache_output_equals_hf( orig_model: DecoderModule, @@ -347,3 +356,11 @@ def assert_decoder_with_positions_equals_hf( Y_hf = hf_model(X, position_ids=positions).last_hidden_state torch_assertclose(Y, Y_hf, atol=atol, rtol=rtol) + + +def assert_model_config(model: TransformerModule, model_output: Tensor): + assert isinstance(model, TransformerModule) + config = model.config + + hidden_width = model_output.size(-1) + assert config.layer.feedforward.hidden_width == hidden_width diff --git a/docs/source/causal-lm.rst b/docs/source/causal-lm.rst index def1f1de..f9507669 100644 --- a/docs/source/causal-lm.rst +++ b/docs/source/causal-lm.rst @@ -7,10 +7,12 @@ Base Classes .. autoclass:: curated_transformers.models.CausalLMModule :members: :show-inheritance: + :inherited-members: Module .. autoclass:: curated_transformers.models.TransformerCausalLM :members: :show-inheritance: + :inherited-members: Module Architectures ------------- diff --git a/docs/source/decoders.rst b/docs/source/decoders.rst index 49019545..fdfd2a5e 100644 --- a/docs/source/decoders.rst +++ b/docs/source/decoders.rst @@ -7,10 +7,12 @@ Base Classes .. autoclass:: curated_transformers.models.DecoderModule :members: :show-inheritance: + :inherited-members: Module .. autoclass:: curated_transformers.models.TransformerDecoder :members: :show-inheritance: + :inherited-members: Module Architectures ------------- diff --git a/docs/source/encoders.rst b/docs/source/encoders.rst index ea9860bb..31b305c4 100644 --- a/docs/source/encoders.rst +++ b/docs/source/encoders.rst @@ -7,10 +7,12 @@ Base Classes .. autoclass:: curated_transformers.models.EncoderModule :members: :show-inheritance: + :inherited-members: Module .. autoclass:: curated_transformers.models.TransformerEncoder :members: :show-inheritance: + :inherited-members: Module Architectures -------------