parallelformers icon indicating copy to clipboard operation
parallelformers copied to clipboard

[Feature Request] Add Bloom to the Auto Policy

Open airsplay opened this issue 2 years ago • 2 comments

Add Bloom to the Auto Policy

It would be great to see the recent bloom model from bigscience can be added to the auto policy. The Bloom model is another auto-regressive large language model thus the policy might be inherited from existing policies.

Expected behavior

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-2b5")
tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-2b5")

from parallelformers import parallelize
parallelize(model, num_gpus=2, fp16=True, verbose='detail')

inputs = tokenizer("Parallelformers is", return_tensors="pt")

outputs = model.generate(
    **inputs,
    num_beams=5,
    no_repeat_ngram_size=4,
    max_length=15,
)

print(f"Output: {tokenizer.batch_decode(outputs)[0]}")

airsplay avatar Aug 13 '22 18:08 airsplay

+1!!!

csinva avatar Sep 03 '22 05:09 csinva

This is great.

seopbo avatar Feb 27 '23 10:02 seopbo