Generate text on multiple GPU

I was wondering how one might go about generating text on a machine with multiple GPU? Here is some example code:

import pandas as pd
from transformers import GPT2Tokenizer, GPT2LMHeadModel

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

def generate_text(primer_sent: str, chunk_sent_length: int, is_finetuning_current_model=None):
    encoded_results = tokenizer(primer_sent, return_tensors='pt')
    tmp_ids = encoded_results['input_ids']
    max_length = tmp_ids.shape[1] + chunk_sent_length
    gen_ids = model.generate(input_ids=tmp_ids,
    decoded_str = tokenizer.decode(gen_ids, skip_special_tokens=True)
    return decoded_str

def generate_text_samples(data_filename: str, chunk_sent_length: int):
    sent_df = pd.read_csv(data_filename)
    sent_df['gen'] = sent_df[0].apply(generate_text, args=(chunk_sent_length,))
    return sent_df

text_df = generate_text_samples('/path/to/data.csv', 50)

When I run this code on my machine, it appears that it is all being done on the CPU. Is it possible to distribute this workflow across the multiple gpu I have? Thanks in advance for your help!


Hi @aclifton314. Have you had any luck with this? I’m also interested in generating text (using a different model - t5) on multiple GPUs to speed up the throughput of my process. Wondering if you were able to do this?

Hi @kmfoda. I was not able to investigate into this further. Since the original posting of this, I do know the methods for generating text were overhauled pretty significantly. I haven’t checked but if there isn’t an input parameter to enable generating text on multiple gpu, then I’m not quite sure how to make it happen. Sorry I don’t have better news! If I discover something, I’ll post it here.