Officially support naive PP for quantized models + PEFT

[younesbelkada]

What does this PR do?

Fixes #1515

Naive Pipeline Parallelism should be supported by accelerate and should work, if we properly educate users on how to use it.

What is NPP?

It is the simplest paradigm for running a model across multiple GPUs. It tries to evenly fit the model across all available GPUs (e.g. device_map="auto")

When to use it and when not to use it?

Initially I added that check because I was afraid users will train 8bit models that are loaded across multiple GPUs and under multi-GPU distributed regime. In that case the model will be converted to DDP (which is fine if the model fits in a single GPU and duplicated across multiple GPUs (Data Parallelism)) - which can lead to many breaking behaviours such as huggingface/peft#269 (comment) .
The fix is to relax the check constraint and to also check if we are under multi GPU distributed regime (expects to use DDP).

In TRL library, it is possible to use the PPOTrainer (that calls accelerator.prepare under the hood) to apply Naive Pipeline Parallelism: https://huggingface.co/docs/trl/main/en/lora_tuning_peft#naive-pipeline-parallelism-npp-for-large-models-60b-models to train 60B+ scale models using RLHF. The error was never raised there because I forgot to store the attribute hf_device_map inside the model class we use in TRL.

To reproduce (you need PEFT and run this script in a multi-GPU env):

from accelerate import Accelerator
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training

model_id = "facebook/opt-350m"
accelerator = Accelerator()

config = LoraConfig(
    r=16, 
    lora_alpha=32, 
    target_modules=["q_proj", "v_proj"], 
    lora_dropout=0.05, 
    bias="none", 
    task_type="CAUSAL_LM"
)

model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_8bit=True)
model = prepare_model_for_int8_training(model)

print(set(model.hf_device_map.values()))

model = get_peft_model(model, config)

model = accelerator.prepare(model)

cc @sgugger @muellerzr

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

Thanks!

[muellerzr]

Thanks! LG2M :)

[robinsonmhj]

this feature is only in main, is there any plan to put into a new release so that I can use pip to install?