[Trainer] Force is_model_parallel when model is loaded in multiple GPUs using accelerate

[younesbelkada]

What does this PR do?

When using the Trainer on a multi-GPU environment, users currently apply a patch that leads to some bugs.
Before running a training they need to call:

setattr(model, 'model_parallel', True)
setattr(model, 'is_parallelizable', True)

Which can lead to unexpected bugs on some models, such as T5, that has the parallelize API that is still in place, thus when forcing model_parallel to be True, calls that API, which is deprecated and should not be maintained.

Script to reproduce:

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling

from peft import prepare_model_for_int8_training,LoraConfig, get_peft_model

causal_lm_model_id = "facebook/opt-350m"

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

tokenizer = AutoTokenizer.from_pretrained(causal_lm_model_id)
model = prepare_model_for_int8_training(model)

# setattr(model, 'model_parallel', True)
# setattr(model, 'is_parallelizable', True)  

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

model = get_peft_model(model, config)

data = load_dataset("Abirate/english_quotes")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)

trainer = Trainer(
    model=model,
    train_dataset=data["train"],
    args=TrainingArguments(
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        warmup_steps=2,
        max_steps=3,
        learning_rate=2e-4,
        fp16=True,
        logging_steps=1,
        output_dir="outputs",
    ),
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

trainer.is_model_parallel = True

model.config.use_cache = False
trainer.train()

cc @sgugger

Related: huggingface/peft#205

[sgugger]

Could you elaborate why is such a patch needed and what is the goal of your PR? Cause all of this seems very hacky.

[younesbelkada]

These hacks were needed because self.place_model_on_device needs to be set to True in order for the Trainer to work correctly on a multi-GPU environment, i.e. with a model that has been loaded across multiple GPUs (so we're talking about Naive PP here). Otherwise users will encounter device mismatch between model's input/output.

Moreover, modifying place_model_on_device directly on TrainingArguments seems to not work, as this argument seems to not be on the __init__ of that class, and also it seems to me that it is better to not touch this attribute as it is a property method:

transformers/src/transformers/training_args.py

Line 1801 in 9419f14

def place_model_on_device(self):

That is why I preferred to introduce a new argument to avoid modifying what is already in place and modify directly what is needed to be edited, without having to modify the model's internals (forcing model_parallel to True on T5 models will call the deprecated parallelize API that leads to some bugs)

[HuggingFaceDocBuilderDev]

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

[sgugger]

Or you could just analyze the device map of the model and determine if there are several GPUs used. It would be cleaner and not require the user to learn the 97th training argument.

[younesbelkada]

Ahh yes good point!

[sgugger]

Maybe actually check the number of GPUs, cause this could be one GPU and CPU here.

[younesbelkada]

Makes sense! Fixed in 5eb72b4

[pacman100]

Multi-device placement should only be on GPUs for naive pipelining to work, right? Offloading to CPU/disk won't work, isn't it the case?

[younesbelkada]

I think offloading to CPU/disk won't work yes, I am also unsure if CPU/disk offload training works out of the box with accelerate (without DeepSpeed)

[sgugger]

Thanks!