abhi-mosaic commited on
Commit
084d672
1 Parent(s): 7c2f422

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +2 -2
README.md CHANGED
@@ -98,8 +98,8 @@ import transformers
98
  name = 'mosaicml/mpt-30b-chat'
99
 
100
  config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
101
- config.attn_config['attn_impl'] = 'torch'
102
- config.init_device = 'cuda' # For fast initialization directly on GPU!
103
 
104
  model = transformers.AutoModelForCausalLM.from_pretrained(
105
  name,
 
98
  name = 'mosaicml/mpt-30b-chat'
99
 
100
  config = transformers.AutoConfig.from_pretrained(name, trust_remote_code=True)
101
+ config.attn_config['attn_impl'] = 'triton' # change this to use triton-based FlashAttention
102
+ config.init_device = 'cuda:0' # For fast initialization directly on GPU!
103
 
104
  model = transformers.AutoModelForCausalLM.from_pretrained(
105
  name,