如何将多GPU添加到人脸模型中

Question

test_dataset = load_dataset("scientific_papers", DATASET_NAME, split="test", cache_dir=CACHE_DIR)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = BigBirdPegasusForConditionalGeneration.from_pretrained(MODEL_ID).to(DEVICE)    
def generate_answer(batch):
inputs_dict = tokenizer(batch["article_text"], padding="max_length", max_length=4096, return_tensors="pt", truncation=True)
inputs_dict = {k: inputs_dict[k].to(DEVICE) for k in inputs_dict}
predicted_abstract_ids = model.generate(**inputs_dict, max_length=256, num_beams=5, length_penalty=0.8)
batch["predicted_abstract"] = tokenizer.decode(predicted_abstract_ids[0], skip_special_tokens=True)
print(batch["predicted_abstract"])
return batch
result = test.map(generate_answer)

我想把muli GPU代码添加到我的代码中，可以只用一行代码吗？

Answer 1

您可以使用Apex。不确定它是否与这个确切的模型兼容，但我一直在与Roberta一起使用它，你应该可以在第3行之后插入这个：

from apex.parallel import DistributedDataParallel as DDP
model = DDP(model)