Training NMT models?

Hello! Thanks, Tim! I tried bitsandbytes for language models like BLOOM, and it works well.

I have a question about NMT models like NLLB, M2M, mBART, or OPUS. I tried inference for NLLB, and apparently it is not supported. Are any of these models supported for inference, and especially for fine-tuning?

Many thanks!

[younesbelkada]

Hi @ymoslem Thanks a lot for your message! Indeed, it is not possible for now to train any 8bit model using transformers - we are currently seeing if we can apply LoRA (Low Rank Adaptators) on 8-bit models using transformers but it is under discussion. We’ll keep you posted

[younesbelkada]

Hi @ymoslem Thanks a lot for your message

1- It is not faster than float16. When loading nllb-200-distilled-600M, with load_in_8bit=True, it takes 17.1 seconds, while with torch_dtype=torch.float16, it takes 17.9 seconds.

Yes this is expected, the 8-bit is currently slower than the fp16 model because the 8-bit quantization is done in two stages. You can check out more about that on the 8-bit integration blogpost.

2- When adding int8_threshold=2.0, I got “an unexpected keyword argument” error. It seems that AutoModelForSeq2SeqLM does not support it.

Yes, please use load_in_8bit_threshold instead. Could you point me to the place you have read that says to use int8_threshold? Maybe the documentation has not been updated

3- GPU consumption seems the similar in both cases; 4871MB with 8bit and 4151MB with float16.

Could you share with me how do you measure that? Note that the memory optimization between fp16 and int8 model really depends on the model size, for nllb-600M you get a memory footpint saving of 1.18, for 3.3B you get a saving of 1.41, etc and it linearly grows with the size of the model. You can check that with this snippet:

import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

src_lang = "eng_Latn"
tgt_lang = "spa_Latn"

model_id = "facebook/nllb-200-3.3B"

model = AutoModelForSeq2SeqLM.from_pretrained(model_id,device_map= "auto",load_in_8bit=False, torch_dtype=torch.float16)
model_8bit = AutoModelForSeq2SeqLM.from_pretrained(model_id,device_map= "auto",load_in_8bit=True)
print(model.get_memory_footprint() / model_8bit.get_memory_footprint())