tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

device_map auto

Browse Source
This commit is contained in:
Casper
2024-03-17 19:52:56 +01:00
parent 884d81331e
commit d43a79b7bf
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/mistral/mixtral_fused.py
View File

@@ -5,7 +5,7 @@ from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
model_path = "mistralai/Mixtral-8x7B-Instruct-v0.1" model_path = "mistralai/Mixtral-8x7B-Instruct-v0.1"
# Load model # Load model
model = AutoModelForCausalLM.from_pretrained(model_path) model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
for name, module in model.named_modules(): for name, module in model.named_modules():
if isinstance(module, MixtralSparseMoeBlock): if isinstance(module, MixtralSparseMoeBlock):