abhi-mosaic
commited on
Commit
•
084d672
1
Parent(s):
7c2f422
Update README.md
Browse files
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'] = '
|
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,
|