Fine-tuning Llama-2 Model with Intel Gaudi¶

import os
import copy
from typing import Dict

import torch

import datasets
import transformers
from transformers import DataCollatorForLanguageModeling

from tqdm import tqdm

import peft

from optimum.habana import GaudiTrainer, GaudiConfig, GaudiTrainingArguments
from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi

def preprocess_dataset(raw_datasets):

    PROMPT_DICT = {
        "prompt_with_input": (
            "Below is an instruction that describes a task, paired with an input that provides further context. "
            "Write a response that appropriately completes the request.\n\n"
            "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:"
        ),
        "prompt_without_input": (
            "Below is an instruction that describes a task. "
            "Write a response that appropriately completes the request.\n\n"
            "### Instruction:\n{instruction}\n\n### Response:"
        ),
    }

    def create_prompts(examples):
        prompts = {}
        prompts["source"] = []
        prompts["target"] = []
        for example in examples:
            prompt_template = (
                PROMPT_DICT["prompt_with_input"] if example["input"] != "" else PROMPT_DICT["prompt_without_input"]
            )
            source = prompt_template.format_map(example)
            prompts["source"].append(source)
            prompts["target"].append(example["output"])
        return prompts

    # Preprocessing the datasets.
    for key in raw_datasets:
        prompts = create_prompts(raw_datasets[key])
        columns_to_be_removed = list(raw_datasets[key].features.keys())
        raw_datasets[key] = raw_datasets[key].add_column("prompt_sources", prompts["source"])
        raw_datasets[key] = raw_datasets[key].add_column("prompt_targets", prompts["target"])
        raw_datasets[key] = raw_datasets[key].remove_columns(columns_to_be_removed)

def preprocess_dataset_to_tokenizer(raw_datasets, tokenizer):
    max_seq_length = 512
    tokenizer.pad_token_id = 0
    tokenizer.eos_token_id = 1
    tokenizer.bos_token_id = 2

    def tokenize(prompt, add_eos_token=True):
        results = tokenizer(
            prompt,
            truncation=True,
            max_length=max_seq_length,
            padding=False,
            return_tensors=None,
        )
        for i in range(len(results["input_ids"])):
            if (
                results["input_ids"][i][-1] != tokenizer.eos_token_id
                and len(results["input_ids"][i]) < max_seq_length
                and add_eos_token
            ):
                results["input_ids"][i].append(tokenizer.eos_token_id)
                results["attention_mask"][i].append(1)

        results["labels"] = copy.deepcopy(results["input_ids"])
        results["input_id_len"] = [len(result) for result in results["input_ids"]]
        return results

    def preprocess_function(examples):
        keys = list(examples.data.keys())
        if len(keys) != 2:
            raise ValueError("Unsupported dataset format")

        st = [s + t for s, t in zip(examples[keys[0]], examples[keys[1]])]

        examples_tokenized = tokenize(st)
        input_ids = examples_tokenized["input_ids"]
        labels = examples_tokenized["labels"]
        return {
            "input_ids": input_ids,
            "labels": labels,
            "attention_mask": examples_tokenized["attention_mask"],
        }

    tokenized_datasets = raw_datasets.map(
        preprocess_function,
        batched=True,
        load_from_cache_file=True,
    )

    def concatenate_data(dataset, max_seq_length):
        concatenated_dataset = {}
        for column in dataset.features:
            concatenated_data = [item for sample in dataset[column] for item in sample]
            reshaped_data = [
                concatenated_data[i * max_seq_length : (i + 1) * max_seq_length]
                for i in range(len(concatenated_data) // max_seq_length)
            ]
            concatenated_dataset[column] = reshaped_data
        return datasets.Dataset.from_dict(concatenated_dataset)

    tokenized_datasets_ = tokenized_datasets["train"].remove_columns(["prompt_sources", "prompt_targets"])
    tokenized_datasets["train"] = concatenate_data(tokenized_datasets_, max_seq_length)

    return tokenized_datasets

def prepare_training_args(config: Dict):
    # prepare execution mode config
    execution_mode = config["execution_mode"]
    use_lazy_mode = True if execution_mode == "lazy" else False
    torch_compile_backend = "hpu_backend" if execution_mode == "eager.compile" else None

    deepspeed = config["deepspeed"] if "deepspeed" in config else None

    return GaudiTrainingArguments(deepspeed=deepspeed,
                                  output_dir=config["output"],
                                  do_train=True,
                                  do_eval=False,
                                  per_device_train_batch_size=config["batch_size_per_worker"],
                                  bf16=True,
                                  learning_rate=config["lr"],
                                  save_strategy="no",
                                  torch_compile_backend=torch_compile_backend,
                                  evaluation_strategy="no",
                                  lr_scheduler_type="cosine",
                                  num_train_epochs=config["epochs"],
                                  use_lazy_mode=use_lazy_mode,
                                  use_habana=True,
                                  pipelining_fwd_bwd=True,
                                  save_only_model=True,
                                  gradient_checkpointing=True,
                                  warmup_ratio=0.03,
                                  throughput_warmup_steps=3,
                                  logging_steps=5)

def prepare_model(config: Dict, device):
    # prepare from pretrained model
    deepspeed = config["deepspeed"] if "deepspeed" in config else None
    if deepspeed is not None:
        auto_config = transformers.AutoConfig.from_pretrained(config["model"], use_cache=False, revision="main", use_auth_token=None, trust_remote_code=None)
        model = transformers.AutoModelForCausalLM.from_pretrained(config["model"], config=auto_config, **config["model_config"])
        model.generation_config.attn_softmax_bf16 = True
        model.generation_config.use_flash_attention = True
    else:
        model = transformers.AutoModelForCausalLM.from_pretrained(config["model"], **config["model_config"])
    model.enable_input_require_grads()

    # convert to peft model for lora training
    peft_config = peft.LoraConfig(**config["lora_config"])
    model = peft.get_peft_model(model, peft_config)

    model.to(dtype=config["model_config"]["torch_dtype"], device=device)

    return model

def train_func_per_worker(config: Dict):
    # adapt transformers to gaudi
    adapt_transformers_to_gaudi()

    # prepare training arguments
    training_args = prepare_training_args(config)

    # prepare datasets
    # here we use dataset "tatsu-lab/alpaca" from huggingface
    raw_datasets = datasets.DatasetDict({"train": datasets.load_dataset("tatsu-lab/alpaca", split='train[0:4096]')})
    preprocess_dataset(raw_datasets)

    # prepare tokenizer
    tokenizer = transformers.AutoTokenizer.from_pretrained(config["model"])
    tokenized_datasets = preprocess_dataset_to_tokenizer(raw_datasets, tokenizer)

    # prepare model
    model = prepare_model(config, training_args.device)

    # prepare data collator
    data_collator = DataCollatorForLanguageModeling(tokenizer, pad_to_multiple_of=8, return_tensors="pt", mlm=False)

    # prepare gaudi config
    gaudi_config = GaudiConfig()
    gaudi_config.use_fused_adam = True
    gaudi_config.use_fused_clip_norm = True

    # instance GaudiTrainer
    trainer = GaudiTrainer(
        model=model,
        gaudi_config=gaudi_config,
        args=training_args,
        train_dataset=tokenized_datasets["train"],
        eval_dataset=None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=None,
        preprocess_logits_for_metrics=None,
    )

    train_result = trainer.train()
    print(f"train_result = {train_result}")
    trainer.save_model()

    return train_result

def train_llama(num_workers, execution_mode, training_method):
    import ray
    from ray.train import ScalingConfig
    from ray.train.torch import TorchTrainer, TorchConfig

    # deepspeed config, can also place it to config file
    deepspeed_config = {
        "steps_per_print": 64,
        "train_batch_size": "auto",
        "train_micro_batch_size_per_gpu": "auto",
        "gradient_accumulation_steps": "auto",
        "bf16": {
            "enabled": True
        },
        "gradient_clipping": 1.0,
        "zero_optimization": {
            "stage": 3,
            "overlap_comm": False,
            "contiguous_gradients": False,
            "stage3_gather_16bit_weights_on_model_save": True
        }
    }

    # Preparing train configurations
    train_config = {
        "execution_mode": execution_mode,
        "model": "/root/models/models--meta-llama--Llama-2-70b-chat-hf/snapshots/e9149a12809580e8602995856f8098ce973d1080/",
        "model_config": {"torch_dtype": torch.bfloat16, "trust_remote_code": False, "use_auth_token": None},
        "lora_config": {"task_type": "CAUSAL_LM", "r": 8, "lora_alpha": 32, "lora_dropout": 0.1, "target_modules": ["q_proj", "v_proj"]},
        "lr": 1e-4,
        "epochs": 2,
        "batch_size_per_worker": 8,
        "output": "/tmp/ray/",
        "deepspeed": deepspeed_config if training_method == "deepspeed" else None,
    }

    # Configure computation resources
    # In ScalingConfig, require an HPU for each worker
    scaling_config = ScalingConfig(num_workers=num_workers, resources_per_worker={"CPU": 1, "HPU": 1})
    # Set backend to hccl in TorchConfig
    torch_config = TorchConfig(backend = "hccl")

    # start your ray cluster
    ray.init()

    # Initialize a Ray TorchTrainer
    trainer = TorchTrainer(
        train_loop_per_worker=train_func_per_worker,
        train_loop_config=train_config,
        torch_config=torch_config,
        scaling_config=scaling_config,
    )

    result = trainer.fit()
    print(f"Training result: {result}")

# set some environment variables
os.environ["RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES"] = "0"
# if using RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES env var
# you must set HABANA_VISIBLE_DEVICES, such as
# os.environ["HABANA_VISIBLE_DEVICES"] = "0,1,2,3"

# execution_mode are ["lazy", "eager", "eager.compile"]
execution_mode = "lazy"
os.environ["PT_HPU_LAZY_MODE"] = "1" if execution_mode == "lazy" else "0"

# training_method are ["ddp", "deepspeed"]
training_method = "deepspeed"
if training_method == "deepspeed":
    os.environ["PT_HPU_MAX_COMPOUND_OP_SIZE"] = "10"
    os.environ["DEEPSPEED_HPU_ZERO3_SYNC_MARK_STEP_REQUIRED"] = "1"

# here use 4 HPUs
train_llama(num_workers=4, execution_mode=execution_mode, training_method=training_method)