update README.md

README.md

@@ -12,7 +12,7 @@ inference: false
 # MPT-7B-Instruct
 
 MPT-7B-Instruct is a model for short-form instruction following.
-It is built by finetuning [MPT-7B](https://huggingface.co/spaces/mosaicml/mpt-7b) on a [dataset](https://huggingface.co/datasets/sam-mosaic/dolly_hhrlhf) derived from the [Databricks Dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k) and the [Anthropic Helpful and Harmless (HH-RLHF)](https://huggingface.co/datasets/Anthropic/hh-rlhf) datasets. 
+It is built by finetuning [MPT-7B](https://huggingface.co/spaces/mosaicml/mpt-7b) on a [dataset](https://huggingface.co/datasets/sam-mosaic/dolly_hhrlhf) derived from the [Databricks Dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k) and the [Anthropic Helpful and Harmless (HH-RLHF)](https://huggingface.co/datasets/Anthropic/hh-rlhf) datasets.
   * License: _CC-By-SA-3.0_
   * [Demo on Hugging Face Spaces](https://huggingface.co/spaces/mosaicml/mpt-7b-instruct)
 
@@ -55,37 +55,41 @@ model = transformers.AutoModelForCausalLM.from_pretrained(
   trust_remote_code=True
 )
 ```
-Note: This model requires that `trust_remote_code=True` be passed to the `from_pretrained` method. 
+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](https://github.com/openai/triton) of FlashAttention, you can load the model with `attn_impl='triton'` and move the model to `bfloat16`:
+To use the optimized [triton implementation](https://github.com/openai/triton) of FlashAttention, you can load the model on GPU (`cuda:0`) with `attn_impl='triton'` and with `bfloat16` precision:
 ```python
-config = transformers.AutoConfig.from_pretrained(
-  'mosaicml/mpt-7b-instruct',
-  trust_remote_code=True
-)
+import torch
+import transformers
+
+name = 'mosaicml/mpt-7b-instruct'
+
+config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
 config.attn_config['attn_impl'] = 'triton'
+config.init_device = 'cuda:0' # For fast initialization directly on GPU!
 
 model = transformers.AutoModelForCausalLM.from_pretrained(
-  'mosaicml/mpt-7b-instruct',
+  name,
   config=config,
-  torch_dtype=torch.bfloat16,
+  torch_dtype=torch.bfloat16, # Load model weights in 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:
 
 ```python
-config = transformers.AutoConfig.from_pretrained(
-  'mosaicml/mpt-7b-instruct',
-  trust_remote_code=True
-)
-config.update({"max_seq_len": 4096})
+import transformers
+
+name = 'mosaicml/mpt-7b-instruct'
+
+config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
+config.max_seq_len = 4096 # (input + output) tokens can now be up to 4096
+
 model = transformers.AutoModelForCausalLM.from_pretrained(
-  'mosaicml/mpt-7b-instruct',
+  name,
   config=config,
   trust_remote_code=True
 )
@@ -182,4 +186,4 @@ Please cite this model using the following format:
     note      = {Accessed: 2023-03-28}, % change this date
     urldate   = {2023-03-28} % change this date
 }
-```
+```