Flan T5 Parallel Usage

parallel_t5.py

from transformers import AutoTokenizer, T5ForConditionalGeneration

# Model Init
n_gpu = 8
tokenizer = AutoTokenizer.from_pretrained("google/flan-ul2")
model = T5ForConditionalGeneration.from_pretrained("google/flan-ul2")
heads_per_gpu = len(model.encoder.block) // n_gpu 
device_map = {
    gpu: list(
        range(
            0 + (gpu * heads_per_gpu),
            (0 + (gpu * heads_per_gpu)) + heads_per_gpu,
        )
    )
    for gpu in range(n_gpu)
}
model.parallelize(device_map)

# training
prompt: "{}\nWhat is the sentiment of this review?\npositive\nnegative\nneutral\n".format("I hate this movie!")
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
labels = tokenizer("positive", return_tensors="pt").input_ids
outputs = model(input_ids=input_ids, labels=labels)
loss = outputs.loss
logits = outputs.logits

loss.backward()