Feature/peft compatible models

.gitignore

@@ -120,6 +120,9 @@ ENV/
 env.bak/
 venv.bak/
 
+# python venv installed in the dir, llmfoundry-venv
+*-venv
+
 # Spyder project settings
 .spyderproject
 .spyproject

llmfoundry/__init__.py

@@ -11,7 +11,7 @@
                                  build_finetuning_dataloader,
                                  build_text_denoising_dataloader)
     from llmfoundry.models.hf import (ComposerHFCausalLM, ComposerHFPrefixLM,
-                                      ComposerHFT5)
+                                      ComposerHFT5, ComposerHFCausalLMFromPython)
     from llmfoundry.models.layers.attention import (
         MultiheadAttention, attn_bias_shape, build_alibi_bias, build_attn_bias,
         flash_attn_fn, scaled_multihead_dot_product_attention,
@@ -46,6 +46,7 @@
     'MPTForCausalLM',
     'ComposerMPTCausalLM',
     'ComposerHFCausalLM',
+    'ComposerHFCausalLMFromPython',
     'ComposerHFPrefixLM',
     'ComposerHFT5',
     'COMPOSER_MODEL_REGISTRY',

llmfoundry/models/hf/__init__.py

@@ -1,7 +1,7 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
-from llmfoundry.models.hf.hf_causal_lm import ComposerHFCausalLM
+from llmfoundry.models.hf.hf_causal_lm import ComposerHFCausalLM, ComposerHFCausalLMFromPython
 from llmfoundry.models.hf.hf_fsdp import (prepare_hf_causal_lm_model_for_fsdp,
                                           prepare_hf_enc_dec_model_for_fsdp,
                                           prepare_hf_model_for_fsdp)
@@ -10,6 +10,7 @@
 
 __all__ = [
     'ComposerHFCausalLM',
+    'ComposerHFCausalLMFromPython',
     'ComposerHFPrefixLM',
     'ComposerHFT5',
     'prepare_hf_causal_lm_model_for_fsdp',

llmfoundry/models/hf/hf_causal_lm.py

@@ -21,6 +21,10 @@
 from llmfoundry.models.hf.model_wrapper import HuggingFaceModelWithZLoss
 from llmfoundry.models.utils import init_empty_weights
 
+# required for loading a python model into composer
+import peft
+import transformers
+
 __all__ = ['ComposerHFCausalLM']
 
 Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
@@ -148,3 +152,36 @@ def __init__(self, om_model_config: DictConfig, tokenizer: Tokenizer):
                                           init_device=init_device)
 
         return composer_model
+
+
+class ComposerHFCausalLMFromPython(HuggingFaceModelWithZLoss):
+    """Configures a :class:`.HuggingFaceModel` around a Causal LM that is loaded in memory.
+    Args:
+        model (peft.peft_model.PeftModel or transformers.PreTrainedModel): The HF model loaded into python memory.
+        tokenizer (PreTrainedTokenizer): The tokenizer that the model will use.
+    """
+
+    def __init__(self, model: Union[peft.peft_model.PeftModel,transformers.PreTrainedModel], tokenizer: Tokenizer):
+
+        train_metrics = [
+            LanguageCrossEntropy(),
+            LanguagePerplexity(),
+        ]
+        eval_metrics = [
+            LanguageCrossEntropy(),
+            LanguagePerplexity(),
+            InContextLearningLMAccuracy(),
+            InContextLearningMultipleChoiceAccuracy(),
+            InContextLearningQAAccuracy(),
+            InContextLearningLMExpectedCalibrationError(),
+            InContextLearningMCExpectedCalibrationError()
+        ]
+
+        composer_model = super().__init__(model=model,
+                                          shift_labels=True,
+                                          tokenizer=tokenizer,
+                                          metrics=train_metrics,
+                                          eval_metrics=eval_metrics,
+                                          z_loss=0.0)
+
+        return composer_model