Update README.md

README.md

@@ -49,14 +49,19 @@ It includes options for many training efficiency features such as [FlashAttentio
 
 ```python
 import transformers
-model = transformers.AutoModelForCausalLM.from_pretrained('mosaicml/mpt-7b-instruct', trust_remote_code=True, torch_dtype=torch.bfloat16)
+model = transformers.AutoModelForCausalLM.from_pretrained('mosaicml/mpt-7b-instruct', trust_remote_code=True)
 ```
+Note: This model requires that `trust_remote_code=True` be passed to the `from_pretrained` method. 
+This is because we use a custom `MPT` model architecture that is not yet part of the Hugging Face `transformers` package.
+`MPT` includes options for many training efficiency features such as [FlashAttention](https://arxiv.org/pdf/2205.14135.pdf), [ALiBi](https://arxiv.org/abs/2108.12409), [QK LayerNorm](https://arxiv.org/abs/2010.04245), and more.
 
-To use the optimized triton implementation of FlashAttention, you can load with `attn_impl='triton'` and move the model to `bfloat16` like so:
-
+To use the optimized [triton implementation](https://github.com/openai/triton) of FlashAttention, you can load the model with `attn_impl='triton'` and move the model to `bfloat16`:
 ```python
-model = transformers.AutoModelForCausalLM.from_pretrained('mosaicml/mpt-7b-instruct', trust_remote_code=True, torch_dtype=torch.bfloat16, attn_impl='triton')
-model.to(device='cuda:0', dtype=torch.bfloat16)
+config = transformers.AutoConfig.from_pretrained('mosaicml/mpt-7b-instruct', trust_remote_code=True)
+config.attn_config['attn_impl'] = 'triton'
+
+model = transformers.AutoModelForCausalLM.from_pretrained('mosaicml/mpt-7b-instruct', config=config, torch_dtype=torch.bfloat16, trust_remote_code=True)
+model.to(device='cuda:0')
 ```
 
 Although the model was trained with a sequence length of 2048, ALiBi enables users to increase the maximum sequence length during finetuning and/or inference. For example:
@@ -65,7 +70,6 @@ Although the model was trained with a sequence length of 2048, ALiBi enables use
 config = transformers.AutoConfig.from_pretrained('mosaicml/mpt-7b', trust_remote_code=True)
 config.update({"max_seq_len": 4096})
 model = transformers.AutoModelForCausalLM.from_pretrained('mosaicml/mpt-7b', config=config, trust_remote_code=True)
-```
 
 This model was trained with the [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b) tokenizer.