面子变形金刚gpt2生成多个GPU

Question

我使用的是拥抱变压器gpt-xl模型来生成多个响应。我正在尝试在多个GPU上运行它，因为GPU内存会因多个较大的响应而最大化。我试过使用dataparallel来做到这一点，但是，在nvidia-smi看来，似乎没有使用过第二代GPU。这是我的代码：

import numpy as np
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch
n_gpu=torch.cuda.device_count()
#device = xm.xla_device()
device=torch.device("cuda:0")
tokenizer = GPT2Tokenizer.from_pretrained('/spell/GPT2Model/GPT2Model/') #downloaded pre-trained model and tokenizer earlier
model = GPT2LMHeadModel.from_pretrained('/spell/GPT2Model/GPT2Model/')
model.to(device)
model = torch.nn.DataParallel(model, device_ids=[0,1])
encoded_prompt=tokenizer.encode(prompt_text, add_special_tokens=True,return_tensors="pt")
encoded_prompt = encoded_prompt.to(device)
outputs = model.module.generate(encoded_prompt,response_length,temperature=.8,num_return_sequences=num_of_responses,repetition_penalty=85,do_sample=True,top_k=80,top_p=.85 )

程序在双T4上获得oom，第二个GPU的内存永远不会超过11M。