abhi-mosaic
commited on
Commit
•
8a9262e
1
Parent(s):
716e2c1
Update README.md
Browse files
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'] = '
|
84 |
-
config.init_device = '
|
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,
|