bigscience

/

bloom

like 4.76k

Follow

BigScience Workshop 798

hi @bver !
Thanks for the great summary! Indeed, I suspect that some weights are offloaded to the disk / cpu, which could explain the inference speed that you are observing.
One thing I am curious about is if you can try:

model = BloomForCausalLM.from_pretrained(
   '.',  # the model path (downloaded locally)
   device_map='auto',
)

This should work fine I guess, since it is what we are using, but most importantly the script should not yield you an error as it did I guess when you tried

model = BloomForCausalLM.from_pretrained(
   '.',  # the model path (downloaded locally)
   device_map='balanced_low_0',
)

(The command above should yield you an error Please specify an offload folder...)

If this is the case then I guess we can dig further ;)

PS:
Also another thing that I suspect is that the model might be loaded in fp32, therefore the device maps that are created by accelerate takes into account fp32 weights! In this case make sure you load it in bf16 by doing:

model = BloomForCausalLM.from_pretrained(
   '.',  # the model path (downloaded locally)
   device_map='auto',
   torch_dtype="auto",
)

or

model = BloomForCausalLM.from_pretrained(
   '.',  # the model path (downloaded locally)
   device_map='auto',
   torch_dtype=torch.bfloat16,
)

@ybelkada ,
thank you very much for your help!

I have tried:

model = BloomForCausalLM.from_pretrained(
    '.',
    device_map='auto',
    torch_dtype='auto'
)

and it fixed model initialization issues (i.e. modules offloading to CPU). My guess is that the torch_dtype was the missing piece. The documentation says that the torch_dtype default uses the pytorch default type and not the model's implicit type, doubling device memory requirements in this case.

The total GPUs memory allocated during inference was ~53% of the total available device memory. The first device was 93% full, the last one 24% (this was my motivation for using balanced_low_0 originally).

A typical inference took ~15 secs for small hundreds of tokens which is great. However, the entire model initialization took ~45 mins, which is less than great. I suspect shards unpacking is ineffective.

Additional questions remain:

1. What is the recommended GPU device number for sending input tokens to? I guess this is the device where the transformer input (first BloomBlocks?) is allocated but it is probably controlled by the device map modes such as ‘auto’ or ‘balanced_low_0’. (I did not try to infer_auto_device_map(model) yet.)
2. Is there any way to make the model initialization faster?

However, generally I am glad to write that my issue was solved. Thanks again!

Thanks a lot @bver ! Glad to hear that it worked on your side!

1- I think that you are right here, you should double check the device map and send the input tokens on the first device as you suggested!
2- From our experience the initialization took in average ~2-3 minutes MAX, are the weights still offloaded on the disk?

are the weights still offloaded on the disk?

I do not think so, the offload_folder argument was not present and also no disk nor RAM excessive allocation was seen during the experiment. I need to check possible bottlenecks (io?) later.

Hi,
after a while, I was able to run more experiments and found that:

1. the device map (model.hf_device_map) confirms all model modules are on GPUs (no "cpu" values, just device numbers),
2. model initialization is likely IO-bound -- if the model is on a tmpfs the loading is faster (~4 minutes). There's probably a way to pre-extract tensors from pytorch .bin files and save them to disk (if that's the reason for inefficiency).