The Falcon 180B is a monumental language model pioneered by Abu Dhabi's Technology Innovation Institute (TII). As an extension of the preceding ‘Falcon’ series, particularly Falcon 40B, Falcon 180B amplifies its predecessor's capabilities. This auto-regressive language model employs a refined transformer architecture and was educated using a vast dataset of 3.5 trillion tokens. Boasting a staggering 180 billion parameters, Falcon 180B claims the title of the most expansive open-source Large Language Model in existence, catering to both research and business applications. It proudly occupies the top position on the Hugging Face Leaderboard for open-source LLMs.
Now, let's delve into the process of fine-tuning Falcon 180B. We'll harness the power of DeepSpeed, Hugging Face Transformers, LoRA, and Flash Attention on a set-up with multiple GPUs.
First, What Exactly Is DeepSpeed ZeRO?
Originating from Microsoft Research, ZeRO, which stands for Zero Redundancy Optimizer, is a groundbreaking memory optimization solution tailored for expansive distributed deep learning. It introduces a novel optimizer that drastically curtails the resources demanded by model and data parallelism, while simultaneously amplifying the trainable parameter count. ZeRO has been engineered to eradicate memory overlaps in both data and model parallel training, ensuring minimized communication overhead and maximized computational precision. This ensures a scalable model size that aligns with the number of available devices, all while maintaining peak efficiency.
Incorporated within DeepSpeed, an open-source toolkit from Microsoft Research, ZeRO propels large model training by enhancing scalability, speed, affordability, and user experience. This makes a 100-billion-parameter model training achievable. Designed to complement PyTorch, DeepSpeed provides comprehensive support for all stages of ZeRO, namely stages 1, 2, and 3, and also facilitates CPU and Disk offloading for optimizer states, gradients, and parameters.
What Is LoRA?
LoRA, or Low-Rank Adaptation, is a specialized training method applied in various domains, including the generation of AI-driven art and the optimization of language models. In the realm of AI art, LoRA aids in the fine-tuning of Stable Diffusion models. It allows the system to be trained on distinct concepts, whether they be specific characters or artistic styles, by enacting minimal alterations to the related model file. Notably, models developed using LoRA are much more compact than standard checkpoint models, offering a convenient solution for those limited by storage capacity.
When applied to language models, particularly those as expansive as GPT-3 with its billions of parameters, LoRA offers a refined approach to fine-tuning. Instead of altering the vast number of pre-trained weights, LoRA introduces trainable layers within each transformer block, thereby cutting down on the number of adjustable parameters and the GPU memory they consume. This strategy bypasses the need to compute gradients for the majority of model weights, which expedites the fine-tuning process. With LoRA, models can be fine-tuned using a fraction of the parameters typically required, marking a significant improvement over conventional methods.
In essence, LoRA stands out as a pivotal training approach, beneficial for fine-tuning both AI-generated art models and mammoth language models. It introduces efficiencies by modifying model files with slight changes and providing compact models that are storage-friendly. Furthermore, in the landscape of language models, LoRA enhances the fine-tuning process, making it both faster and more resource-efficient.
What Is Flash Attention?
Flash Attention is an optimization algorithm designed to accelerate the attention mechanism within Transformer-based language models. The attention mechanism, a vital component in Transformers, allows the model to focus on different parts of the input sequence when making predictions. However, this mechanism can be computationally expensive, particularly when processing longer text sequences, leading to high memory costs and slower processing times.
The Flash Attention algorithm addresses this issue through a series of techniques. It utilizes methods such as tiling, which involves dividing the input sequence into smaller segments, and recomputation, which involves recalculating certain values to reduce memory usage. These techniques help to streamline the computations involved in the attention mechanism, thereby enhancing the model's efficiency in handling longer text sequences.
Flash Attention-2 is an improved version of the original algorithm that further optimizes the parallelism and work partitioning of computations. By refining the distribution of tasks and enhancing parallel processing capabilities, Flash Attention-2 achieves a notable 2x speedup compared to its predecessor. This enhancement is especially significant as it enables the algorithm to reach a high processing speed of 230 TFLOPS/s (tera floating-point operations per second) on A100 GPUs, indicating a substantial improvement in computational performance and efficiency.
Steps for Fine-Tuning Falcon 180B with DeepSpeed ZeRO, LoRA & Flash Attention on E2E Cloud
1. We have to first make sure that we have accepted the license tiiuae/falcon-180B to be able to use it. You can accept the license by clicking on the Agree and Access Repository button on the model page at https://huggingface.co/tiiuae/falcon-180B.
2. Launch a GPU node on the E2E Cloud Computing platform.
For this article, we will be using the NVIDIA-A-100 node.
3. Set up the environment.
conda create --name hf python=3.10 -c conda-forge
When you execute this command, Conda will set up a new isolated environment named ‘hf’ with Python 3.10 installed within it. This will allow you to work on projects that specifically require this environment without affecting other Python projects or environments on your system. Additionally, the use of the ‘conda-forge’ channel ensures that packages from this community-driven repository are used for the environment.
# install torch with the correct cuda version, check nvcc --version
!pip install torch --extra-index-url https://download.pytorch.org/whl/cu118 --upgrade
# install Hugging Face Libraries and additional dependencies
!pip install "transformers==4.33.1" "datasets==2.14.5" "accelerate==0.22.0" "evaluate==0.4.0" "peft==0.5.0" tensorboard packaging --upgrade
# install deepspeed and ninja for jit compilations of kernels
!pip install "deepspeed==0.10.3" ninja --upgrade
# install additional Flash Attention
!pip install flash-attn --no-build-isolation --upgrade
Let’s log in into our Hugging Face account by running the following command:
!huggingface-cli login --token YOUR_TOKEN
The purpose of this command is to authenticate the user's identity on the Hugging Face Hub and enable them to perform various actions on the Hub.
4. Load and pre-process the dataset.
We'll work with ‘Dolly’, an open-source dataset containing instruction-following records produced by numerous Databricks employees. The dataset falls into various behavioral categories detailed in the InstructGPT paper, such as brainstorming, classification, closed QA, generation, information extraction, open QA, and summarization.
{
"instruction": "When was the wheel invented?",
"context": "",
"response": "No one knows when the wheel was first made. Archaeological evidence seems to show that the wheel was first used to make pottery in Mesopotamia about 7,000 years ago, during the Neolithic period. Later, it was used to make chariots. Chariot wheels have been found in Sumer as early as 3000 BC. The wheel is one of the most important inventions in human history. It was very useful for making pottery and getting around on land. It is also a key part of machines."
}
We use the load_dataset() method from the Hugging Face Datasets library.
from datasets import load_dataset
from random import randrange
# Load dataset from the hub
dataset = load_dataset("databricks/databricks-dolly-15k", split="train")
print(f"dataset size: {len(dataset)}")
print(dataset[randrange(len(dataset))])
# dataset size: 15011
To fine-tune our model effectively, we are required to transform our organized examples into a set of tasks specified by instructions. We establish a formatting function designed to accept a sample and produce a string containing our formatted instruction.
def format_dolly(sample):
instruction = f"### Instruction\n{sample['instruction']}"
context = f"### Context\n{sample['context']}" if len(sample["context"]) > 0 else None
response = f"### Answer\n{sample['response']}"
# join all the parts together
prompt = "\n\n".join([i for i in [instruction, context, response] if i is not None])
return prompt
Let's test our formatting function on a random example.
from random import randrange
print(format_dolly(dataset[randrange(len(dataset))]))
Also, apart from formatting our samples, we aim to consolidate multiple samples into one sequence to enhance the efficiency of our training.
from transformers import AutoTokenizer
model_id = "tiiuae/falcon-180B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token
We create several auxiliary functions to bundle our samples into sequences of a specified length, and subsequently tokenize them.
# Print random sample
print(dataset[randint(0, len(dataset))]["text"])
# Empty list to save remainder from batches to use in next batch
remainder = {"input_ids": [], "attention_mask": [], "token_type_ids": []}
def chunk(sample, chunk_length=2048):
# define global remainder variable to save remainder from batches to use in next batch
global remainder
# Concatenate all texts and add remainder from previous batch
concatenated_examples = {k: list(chain(*sample[k])) for k in sample.keys()}
concatenated_examples = {k: remainder[k] + concatenated_examples[k] for k in concatenated_examples.keys()}
# get total number of tokens for batch
batch_total_length = len(concatenated_examples[list(sample.keys())[0]])
# get max number of chunks for batch
if batch_total_length >= chunk_length:
batch_chunk_length = (batch_total_length // chunk_length) * chunk_length
# Split by chunks of max_len.
result = {
k: [t[i : i + chunk_length] for i in range(0, batch_chunk_length, chunk_length)]
for k, t in concatenated_examples.items()
}
# add remainder to global variable for next batch
remainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()}
# prepare labels
result["labels"] = result["input_ids"].copy()
return result
# tokenize and chunk dataset
lm_dataset = dataset.map(
lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(dataset.features)
).map(
partial(chunk, chunk_length=2048),
batched=True,
)
# Print total number of samples
print(f"Total number of samples: {len(lm_dataset)}")
Once we have processed the datasets, our intention is to store them on disk so that we can utilize the processed dataset later for training purposes.
lm_dataset.save_to_disk("dolly-processed")
5. Fine-tune Falcon 180B using DeepSpeed, Hugging Face Transformers, and LoRA with Flash Attention.
The Hugging Face Transformers Trainer seamlessly incorporates DeepSpeed ZeRO, allowing for easy utilization with the provision of a DeepSpeed config file, with the Trainer managing the process. Two DeepSpeed configurations, including CPU offloading, were generated for our experiments.
ds_falcon_180b_z3.json
{
"bf16": {
"enabled": "auto"
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto"
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto"
}
},
"zero_optimization": {
"stage": 3,
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
ds_falcon_180b_z3_offload.json
{
"bf16": {
"enabled": "auto"
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto"
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto"
}
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
ds_falcon_180b_z3_offload.json
{
"bf16": {
"enabled": "auto"
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"betas": "auto",
"eps": "auto",
"weight_decay": "auto"
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto"
}
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": true
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false
}
Alongside the DeepSpeed configuration, we require a training script that incorporates LoRA and integrates flash-attention into our model via the falcon_patch.py utilities. Our run_ds_lora.py script has been designed for this purpose, incorporating the implementation of LoRA with the assistance of peft_utils.py.
run_ds_lora.py
from dataclasses import dataclass, field
from typing import cast
import os
import subprocess
from typing import Optional
import torch
from transformers import HfArgumentParser, TrainingArguments, Trainer
from object_detection.utils.peft_utils import SaveDeepSpeedPeftModelCallback, create_and_prepare_model
from datasets import load_from_disk
# Define and parse arguments.
@dataclass
class ScriptArguments:
"""
Additional arguments for training, which are not part of TrainingArguments.
"""
model_id: str = field(
metadata={
"help": "The model that you want to train from the Hugging Face hub. E.g. gpt2, gpt2-xl, bert, etc."
},
)
dataset_path: Optional[str] = field(
default="timdettmers/openassistant-guanaco",
metadata={"help": "The preference dataset to use."},
)
lora_alpha: Optional[int] = field(default=16)
lora_dropout: Optional[float] = field(default=0.1)
lora_r: Optional[int] = field(default=64)
use_flash_attn: Optional[bool] = field(
default=False,
metadata={"help": "Enables Flash attention for training."},
)
merge_adapters: bool = field(
metadata={"help": "Whether to merge weights for LoRA."},
default=False,
)
def training_function(script_args:ScriptArguments, training_args:TrainingArguments):
# Load processed dataset from disk
dataset = load_from_disk(script_args.dataset_path)
# Load and create peft model
model, peft_config, tokenizer = create_and_prepare_model(script_args.model_id,training_args, script_args)
model.config.use_cache = False
# Create trainer and add callbacks
trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.accelerator.print(f"{trainer.model}")
trainer.model.print_trainable_parameters()
trainer.add_callback(SaveDeepSpeedPeftModelCallback(trainer, save_steps=training_args.save_steps))
# Start training
trainer.train()
# Save model on main process
trainer.accelerator.wait_for_everyone()
state_dict = trainer.accelerator.get_state_dict(trainer.deepspeed)
unwrapped_model = trainer.accelerator.unwrap_model(trainer.deepspeed)
if trainer.accelerator.is_main_process:
unwrapped_model.save_pretrained(training_args.output_dir, state_dict=state_dict)
trainer.accelerator.wait_for_everyone()
# TODO: add merge adapters
# Save everything else on main process
if trainer.args.process_index == 0:
if script_args.merge_adapters:
# merge adapter weights with base model and save
# save int 4 model
trainer.model.save_pretrained(training_args.output_dir, safe_serialization=False)
# clear memory
del model
del trainer
torch.cuda.empty_cache()
from peft import AutoPeftModelForCausalLM
# load PEFT model in fp16
model = AutoPeftModelForCausalLM.from_pretrained(
training_args.output_dir,
low_cpu_mem_usage=True,
torch_dtype=torch.float16,
)
# Merge LoRA and base model and save
model = model.merge_and_unload()
model.save_pretrained(
training_args.output_dir, safe_serialization=True, max_shard_size="8GB"
)
else:
trainer.model.save_pretrained(
training_args.output_dir, safe_serialization=True
)
# save tokenizer
tokenizer.save_pretrained(training_args.output_dir)
def main():
parser = HfArgumentParser([ScriptArguments,TrainingArguments])
script_args, training_args = parser.parse_args_into_dataclasses()
script_args = cast(ScriptArguments, script_args)
training_args = cast(TrainingArguments, training_args)
training_function(script_args, training_args)
if __name__ == "__main__":
main()
peft_utils.py
import torch
from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl
from peft import LoraConfig, get_peft_model
from peft.tuners.lora import LoraLayer
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
AutoTokenizer,
TrainingArguments,
)
from utils.falcon_patch import replace_attn_with_flash_attn as replace_falcon_attn_with_flash_attn
from utils.llama_patch import replace_attn_with_flash_attn as replace_llama_attn_with_flash_attn
class SaveDeepSpeedPeftModelCallback(TrainerCallback):
def __init__(self, trainer, save_steps=500):
self.trainer = trainer
self.save_steps = save_steps
def on_step_end(
self,
args: TrainingArguments,
state: TrainerState,
control: TrainerControl,
**kwargs,
):
if (state.global_step + 1) % self.save_steps == 0:
self.trainer.accelerator.wait_for_everyone()
state_dict = self.trainer.accelerator.get_state_dict(self.trainer.deepspeed)
unwrapped_model = self.trainer.accelerator.unwrap_model(self.trainer.deepspeed)
if self.trainer.accelerator.is_main_process:
unwrapped_model.save_pretrained(args.output_dir, state_dict=state_dict)
self.trainer.accelerator.wait_for_everyone()
return control
def create_and_prepare_model(model_id: str, training_args: TrainingArguments, script_args):
model = AutoModelForCausalLM.from_pretrained(
model_id,
use_cache=not training_args.gradient_checkpointing,
use_flash_attention_2=script_args.use_flash_attn,
)
print("model loaded")
# find all linear modules in model for lora
target_modules = find_all_linear_names(model)
# create lora config
peft_config = LoraConfig(
lora_alpha=script_args.lora_alpha,
lora_dropout=script_args.lora_dropout,
r=script_args.lora_r,
bias="none",
task_type="CAUSAL_LM",
target_modules=target_modules,
)
# enable gradient checkpointing
if training_args.gradient_checkpointing:
model.gradient_checkpointing_enable()
# pre-process the model by upcasting the layer norms in float 32 for
# Adapted from https://github.com/tmm1/axolotl/blob/2eda9e02a9d15a7a3f92b41f257d9844d72fc220/src/axolotl/utils/models.py#L338
print("pre-processing model for peft")
for name, module in model.named_modules():
if isinstance(module, LoraLayer):
module = module.to(torch.bfloat16)
if "norm" in name:
module = module.to(torch.bfloat16)
if any(x in name for x in ["lm_head", "embed_tokens", "wte", "wpe"]):
if hasattr(module, "weight"):
module = module.to(torch.bfloat16)
# initialize peft model
print("initializing peft model")
model = get_peft_model(model, peft_config)
# logger.info parameters
model.print_trainable_parameters()
# tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token
return model, peft_config, tokenizer
def find_all_linear_names(model):
cls = torch.nn.Linear
lora_module_names = set()
for name, module in model.named_modules():
if isinstance(module, cls):
names = name.split(".")
lora_module_names.add(names[0] if len(names) == 1 else names[-1])
if "lm_head" in lora_module_names: # needed for 16-bit
lora_module_names.remove("lm_head")
return list(lora_module_names)
falcon_patch.py
from typing import List, Optional, Tuple
import torch
import transformers
from peft.tuners.lora import LoraLayer
try:
from flash_attn import flash_attn_func
except Exception:
raise ModuleNotFoundError(
"Please install FlashAttention first, e.g., with pip install flash-attn --no-build-isolation, Learn more at https://github.com/Dao-AILab/flash-attention#installation-and-features"
)
try:
from einops import rearrange
except Exception:
raise ModuleNotFoundError("Please install einops first, e.g., with pip install einops")
# ADAPTED https://github.com/pacman100/DHS-LLM-Workshop/blob/main/chat_assistant/training/falcon_flash_attn_monkey_patch.py
def forward(
self,
hidden_states: torch.Tensor,
alibi: Optional[torch.Tensor],
attention_mask: torch.Tensor,
layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
head_mask: Optional[torch.Tensor] = None,
use_cache: bool = False,
output_attentions: bool = False,
):
fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size]
num_kv_heads = self.num_heads if self.new_decoder_architecture else self.num_kv_heads
# 3 x [batch_size, seq_length, num_heads, head_dim]
(query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)
batch_size, query_length, _, _ = query_layer.shape
query_layer = query_layer.transpose(1, 2).reshape(batch_size * self.num_heads, query_length, self.head_dim)
key_layer = key_layer.transpose(1, 2).reshape(
batch_size * num_kv_heads,
query_length,
self.head_dim,
)
value_layer = value_layer.transpose(1, 2).reshape(batch_size * num_kv_heads, query_length, self.head_dim)
past_kv_length = 0 if layer_past is None else layer_past[0].shape[1]
query_layer, key_layer = self.maybe_rotary(query_layer, key_layer, past_kv_length)
if layer_past is not None:
past_key, past_value = layer_past
# concatenate along seq_length dimension:
# - key: [batch_size * self.num_heads, kv_length, head_dim]
# - value: [batch_size * self.num_heads, kv_length, head_dim]
key_layer = torch.cat((past_key, key_layer), dim=1)
value_layer = torch.cat((past_value, value_layer), dim=1)
_, kv_length, _ = key_layer.shape
if use_cache:
present = (key_layer, value_layer)
else:
present = None
attention_mask_float = (attention_mask * 1.0).masked_fill(attention_mask, float("-1e9")).to(query_layer.dtype)
query_layer_ = (
query_layer.reshape(batch_size, self.num_heads, -1, self.head_dim).transpose(1, 2).to(torch.bfloat16)
)
key_layer_ = key_layer.reshape(batch_size, num_kv_heads, -1, self.head_dim).transpose(1, 2).to(torch.bfloat16)
value_layer_ = value_layer.reshape(batch_size, num_kv_heads, -1, self.head_dim).transpose(1, 2).to(torch.bfloat16)
if alibi is not None:
raise ValueError("`alibi` is not supported when `use_flash_attn` is True")
# below output will have shape (batch_size, seqlen, nheads, headdim)
attn_output = flash_attn_func(query_layer_, key_layer_, value_layer_, causal=True)
attn_output = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
output_tensor = self.dense(attn_output)
return output_tensor, present
def replace_attn_with_flash_attn():
cuda_major, cuda_minor = torch.cuda.get_device_capability()
if cuda_major < 8:
print(
"Flash attention is only supported on Ampere or Hopper GPU during training due to head dim > 64 backward."
"ref: https://github.com/HazyResearch/flash-attention/issues/190#issuecomment-1523359593"
)
transformers.models.falcon.modeling_falcon.FalconAttention.forward = forward
def unplace_flash_attn_with_attn():
import importlib
import transformers
print("Reloading falcon model, unpatching flash attention")
importlib.reload(transformers.models.falcon.modeling_falcon)
# Adapted from https://github.com/tmm1/axolotl/blob/2eda9e02a9d15a7a3f92b41f257d9844d72fc220/src/axolotl/utils/models.py#L338
def upcast_layer_for_flash_attention(model, torch_dtype):
# LlamaRMSNorm layers are in fp32 after kbit_training, so we need to
# convert them back to fp16/bf16 for flash-attn compatibility.
for name, module in model.named_modules():
if isinstance(module, LoraLayer):
module.to(torch_dtype)
if "norm" in name:
module.to(torch_dtype)
if "lm_head" in name or "embed_tokens" in name:
if hasattr(module, "weight"):
module.to(torch_dtype)
return model
Upon confirming the appropriate configuration and training script, we can initiate the training process using torchrun.
!torchrun --nproc_per_node 8 run_ds_lora.py \
--model_id tiiuae/falcon-180B \
--dataset_path dolly-processed \
--output_dir falcon-180b-lora-fa \
--num_train_epochs 3 \
--per_device_train_batch_size 1 \
--learning_rate 4e-3 \
--gradient_checkpointing True \
--gradient_accumulation_steps 8 \
--bf16 True \
--tf32 True \
--use_flash_attn True \
--lr_scheduler_type "constant_with_warmup" \
--logging_steps 25 \
--save_steps 100 \
--save_total_limit 3 \
--deepspeed configs/ds_falcon_180b_z3.json
Please note that due to our usage of LoRA, we are solely saving the ‘trained’ adapter weights to conserve storage space. If you intend to consolidate the adapters back into the base model and preserve the merged model, you can either include ‘--merge_adapters True’ or employ the merge_adapter_weights.py script.
merge_adapter_weights.py
from dataclasses import dataclass, field
from typing import Optional
import torch
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer, HfArgumentParser
# execute script
# python scripts/merge_peft_into_model.py --peft_model_id falcon-180b-lora-fa --output_dir merged-weights --save_tokenizer True
@dataclass
class ScriptArguments:
peft_model_id: str = field(metadata={"help": "model id or path to model"})
output_dir: Optional[str] = field(default="merged-weights", metadata={"help": "where the merged model should be saved"})
save_tokenizer: Optional[bool] = field(default=True, metadata={"help": "whether to save the tokenizer"})
push_to_hub: Optional[bool] = field(default=False, metadata={"help": "whether to push the model to the hub"})
repository_id: Optional[str] = field(default=None, metadata={"help": "the model name"})
parser = HfArgumentParser(ScriptArguments)
args = parser.parse_args_into_dataclasses()[0]
model = AutoPeftModelForCausalLM.from_pretrained(
args.peft_model_id,
low_cpu_mem_usage=True,
torch_dtype=torch.float16,
)
# Merge LoRA and base model and save
model = model.merge_and_unload()
model.save_pret
