Does BFloat16 model training supported on ARC?

[XinyuYe-Intel]

Describe the issue

Hi, I tried to finetune Llama2 on ARC with BFloat16 data type and AdamW optimizer using transformers.Trainer, but I met RuntimeError: parameter in optimizer(Adamw) is not FP32, need check from intel_extension_for_pytorch/optim/_functional.py", line 1256, in adamw_step, so I want to confirm that does BFloat16 model training supported on ARC now? If not, could you please add this feature?
PS: I'm using dev/QLLM branch commit id d600370d44882ae8e15949452fe1d9f324cf6900

[kminhta]

Hi, Arc should support bfloat16 training. Can you share a simple reproducer, I can check on my end

[XinyuYe-Intel]

Hi @kta-intel , below code can reproduce the issue.

import torch
from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM
from transformers import TrainingArguments, Trainer

dataset_name = "yahma/alpaca-cleaned"
model_path = "facebook/opt-125m"

dataset = load_dataset(dataset_name)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path).to(torch.bfloat16)

def tokenize_function(examples):
    res = tokenizer(examples["instruction"], padding="max_length", truncation=True)
    res.update({"labels": res["input_ids"]})
    return res

tokenized_datasets = dataset.map(tokenize_function, batched=True)
training_args = TrainingArguments(output_dir="tmp", per_device_train_batch_size=1)

trainer = Trainer(model=model, args=training_args, train_dataset=tokenized_datasets['train'])
trainer.train()

[Jasonzzt]

Hi @XinyuYe-Intel , do you solve this problem later? I also encounted the same error when running LISA fintuning on ARC770.

[kminhta]

@XinyuYe-Intel @Jasonzzt sorry for the late response. I was able to reproduce the error with your script, too. Instead of sending model directly to bfloat16, try setting bf16=True in the training_args. You may also need to import ipex to offload properly to the xpu

Btw, since you're asking specifically for Arc, did your script work as intended for other Intel GPUs or have you only tried on Arc? Want to make sure there isn't some major difference in functionality

[Jasonzzt]

@XinyuYe-Intel @Jasonzzt sorry for the late response. I was able to reproduce the error with your script, too. Instead of sending model directly to bfloat16, try setting bf16=True in the training_args. You may also need to import ipex to offload properly to the xpu

Btw, since you're asking specifically for Arc, did your script work as intended for other Intel GPUs or have you only tried on Arc? Want to make sure there isn't some major difference in functionality

Hi, thanks for your reply!
I tried setting bf16=True in training_args but the error persists. I also get the same error on the MAX1550.

These problems all occurred when the optimizer was adamw_torch. When I changed the optimizer to adamw_hf, I no longer encountered this problem.

[kminhta]

Hmm, I can't seem to reproduce the issue with adamw_torch. It's the same runtime error as originally reported?
Here's the code snippet I ran (modified from the snippet earlier in the thread), if you want to give this a try:

import torch
import intel_extension_for_pytorch as ipex # import ipex
from datasets import load_dataset
from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM
from transformers import TrainingArguments, Trainer
import os
os.environ['CURL_CA_BUNDLE'] = ''

dataset_name = "yahma/alpaca-cleaned"
model_path = "facebook/opt-125m"

dataset = load_dataset(dataset_name)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path).to('xpu') # send model to xpu

def tokenize_function(examples):
    res = tokenizer(examples["instruction"], padding="max_length", truncation=True)
    res.update({"labels": res["input_ids"]})
    return res

tokenized_datasets = dataset.map(tokenize_function, batched=True)
training_args = TrainingArguments(output_dir="tmp", per_device_train_batch_size=1, bf16=True, optim = "adamw_torch") # set bf16 and set optim to adamw_torch

trainer = Trainer(model=model, args=training_args, train_dataset=tokenized_datasets['train'])
trainer.train()

[Jasonzzt]

I run your code on ARC, it seems that there's no error RuntimeError: parameter in optimizer(Adamw) is not FP32, need check.
Actually, I got this error when I tried to do lisa fintuning with IPEX-LLM on ARC/MAX1550

[kminhta]

IPEX-LLM as this in repo? https://github.com/intel-analytics/ipex-llm

[kminhta]

@Jasonzzt @XinyuYe-Intel I just wanted to follow up. Are you still running into this issue or may we close?

[Jasonzzt]

Hi @kta-intel , sorry to reply so late. Actually the original problem still remains when trying to do lisa fintuning with IPEX-LLM on ARC/MAX1550.
Optimizer: adamw_torch

[kminhta]

Can you open a ticket directly in the ipex-llm repository. It is maintained by a different dev team, so they may be able to help you better there

[Jasonzzt]

Can you open a ticket directly in the ipex-llm repository. It is maintained by a different dev team, so they may be able to help you better there

Sure, thank you. You can close this issue