README.md

README.md

@@ -49,34 +49,37 @@ model = transformers.AutoModelForCausalLM.from_pretrained(
 )
 ```
 
-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-storywriter',
-  trust_remote_code=True
-)
+import torch
+import transformers
+
+name = 'mosaicml/mpt-7b-storywriter'
+
+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-storywriter',
+  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 and finetuned with a sequence length of 65536, 
+Although the model was trained with a sequence length of 2048 and finetuned with a sequence length of 65536,
 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-storywriter',
-  trust_remote_code=True
-)
-config.update({"max_seq_len": 83968})
+import transformers
+
+name = 'mosaicml/mpt-7b'
+
+config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
+config.max_seq_len = 83968 # (input + output) tokens can now be up to 83968
+
 model = transformers.AutoModelForCausalLM.from_pretrained(
-  'mosaicml/mpt-7b-storywriter',
+  name,
   config=config,
   trust_remote_code=True
 )
@@ -155,8 +158,8 @@ The data was tokenized using the [EleutherAI/gpt-neox-20b](https://huggingface.c
 
 ### Training Configuration
 
-This model was trained on 8 A100-80GBs for about 2 days using the [MosaicML Platform](https://www.mosaicml.com/platform). 
-The model was trained with sharded data parallelism using [FSDP](https://pytorch.org/docs/stable/fsdp.html) and used the [LION](https://arxiv.org/abs/2302.06675) optimizer. 
+This model was trained on 8 A100-80GBs for about 2 days using the [MosaicML Platform](https://www.mosaicml.com/platform).
+The model was trained with sharded data parallelism using [FSDP](https://pytorch.org/docs/stable/fsdp.html) and used the [LION](https://arxiv.org/abs/2302.06675) optimizer.
 
 ## Limitations and Biases
 
@@ -193,4 +196,4 @@ Please cite this model using the following format:
     note      = {Accessed: 2023-03-28}, % change this date
     urldate   = {2023-03-28} % change this date
 }
-```
+```