abhi-mosaic commited on
Commit
8a9262e
1 Parent(s): 716e2c1

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +2 -2
README.md CHANGED
@@ -80,8 +80,8 @@ import transformers
80
  name = 'mosaicml/mpt-30b-instruct'
81
 
82
  config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
83
- config.attn_config['attn_impl'] = 'torch' # change this to use triton
84
- config.init_device = 'cpu' # For fast initialization directly on GPU! (if you have enough memory)
85
 
86
  model = transformers.AutoModelForCausalLM.from_pretrained(
87
  name,
 
80
  name = 'mosaicml/mpt-30b-instruct'
81
 
82
  config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
83
+ config.attn_config['attn_impl'] = 'triton' # change this to use triton-based FlashAttention
84
+ config.init_device = 'cuda:0' # For fast initialization directly on GPU!
85
 
86
  model = transformers.AutoModelForCausalLM.from_pretrained(
87
  name,