In [1]:
from datasets import Dataset, load_dataset
# TODO: I hate using globals like this, but I cannot come up with a cleaner alternative right now
In [2]:
%%capture
import torch
major_version, minor_version = torch.cuda.get_device_capability()
# Must install separately since Colab has torch 2.2.1, which breaks packages
!pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
if major_version >= 8:
# Use this for new GPUs like Ampere, Hopper GPUs (RTX 30xx, RTX 40xx, A100, H100, L40)
!pip install --no-deps packaging ninja einops flash-attn xformers trl peft accelerate bitsandbytes
else:
# Use this for older GPUs (V100, Tesla T4, RTX 20xx)
!pip install --no-deps xformers trl peft accelerate bitsandbytes
pass
In [3]:
from unsloth import FastLanguageModel
import torch
max_seq_length = 25536 # Choose any! We auto support RoPE Scaling internally!
dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
load_in_4bit = False # Use 4bit quantization to reduce memory usage. Can be False.
# 4bit pre quantized models we support for 4x faster downloading + no OOMs.
fourbit_models = [
"unsloth/mistral-7b-bnb-4bit",
"unsloth/mistral-7b-instruct-v0.2-bnb-4bit",
"unsloth/llama-2-7b-bnb-4bit",
"unsloth/llama-2-13b-bnb-4bit",
"unsloth/codellama-34b-bnb-4bit",
"unsloth/tinyllama-bnb-4bit",
"unsloth/gemma-7b-bnb-4bit", # New Google 6 trillion tokens model 2.5x faster!
"unsloth/gemma-2b-bnb-4bit",
] # More models at https://huggingface.co/unsloth
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = "unsloth/llama-2-13b-bnb-4bit", # Choose ANY! eg teknium/OpenHermes-2.5-Mistral-7B
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = load_in_4bit,
)
🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning. ==((====))== Unsloth 2024.9.post4: Fast Llama patching. Transformers = 4.44.2. \\ /| GPU: NVIDIA A100-SXM4-40GB. Max memory: 39.394 GB. Platform = Linux. O^O/ \_/ \ Pytorch: 2.4.0+cu121. CUDA = 8.0. CUDA Toolkit = 12.1. \ / Bfloat16 = TRUE. FA [Xformers = 0.0.27.post2. FA2 = True] "-____-" Free Apache license: http://github.com/unslothai/unsloth
Unsloth: unsloth/llama-2-13b can only handle sequence lengths of at most 4096. But with kaiokendev's RoPE scaling of 6.234, it can be magically be extended to 25536!
Loading checkpoint shards: 0%| | 0/3 [00:00<?, ?it/s]
In [4]:
model = FastLanguageModel.get_peft_model(
model,
r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj",],
lora_alpha = 16,
lora_dropout = 0, # Supports any, but = 0 is optimized
bias = "none", # Supports any, but = "none" is optimized
use_gradient_checkpointing = "unsloth", # True or "unsloth" for long context
random_state = 3407,
use_rslora = False, # We support rank stabilized LoRA
loftq_config = None, # And LoftQ
)
Unsloth 2024.9.post4 patched 40 layers with 40 QKV layers, 40 O layers and 40 MLP layers.
In [5]:
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
### Instruction:
{}
### Input:
{}
### Response:
{}"""
EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN
def formatting_prompts_func(examples):
instructions = examples["instruction"]
inputs = examples["input"]
outputs = examples["output"]
texts = []
for instruction, input, output in zip(instructions, inputs, outputs):
# Must add EOS_TOKEN, otherwise your generation will go on forever!
text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN
texts.append(text)
return { "text" : texts, }
pass
from datasets import load_dataset
dataset = load_dataset("Yukang/LongAlpaca-12k", split = "train")
dataset = dataset.map(formatting_prompts_func, batched = True,)
Repo card metadata block was not found. Setting CardData to empty.
In [6]:
print(model)
PeftModelForCausalLM(
(base_model): LoraModel(
(model): LlamaForCausalLM(
(model): LlamaModel(
(embed_tokens): Embedding(32000, 5120)
(layers): ModuleList(
(0-39): 40 x LlamaDecoderLayer(
(self_attn): LlamaAttention(
(q_proj): lora.Linear(
(base_layer): Linear(in_features=5120, out_features=5120, bias=False)
(lora_dropout): ModuleDict(
(default): Identity()
)
(lora_A): ModuleDict(
(default): Linear(in_features=5120, out_features=16, bias=False)
)
(lora_B): ModuleDict(
(default): Linear(in_features=16, out_features=5120, bias=False)
)
(lora_embedding_A): ParameterDict()
(lora_embedding_B): ParameterDict()
(lora_magnitude_vector): ModuleDict()
)
(k_proj): lora.Linear(
(base_layer): Linear(in_features=5120, out_features=5120, bias=False)
(lora_dropout): ModuleDict(
(default): Identity()
)
(lora_A): ModuleDict(
(default): Linear(in_features=5120, out_features=16, bias=False)
)
(lora_B): ModuleDict(
(default): Linear(in_features=16, out_features=5120, bias=False)
)
(lora_embedding_A): ParameterDict()
(lora_embedding_B): ParameterDict()
(lora_magnitude_vector): ModuleDict()
)
(v_proj): lora.Linear(
(base_layer): Linear(in_features=5120, out_features=5120, bias=False)
(lora_dropout): ModuleDict(
(default): Identity()
)
(lora_A): ModuleDict(
(default): Linear(in_features=5120, out_features=16, bias=False)
)
(lora_B): ModuleDict(
(default): Linear(in_features=16, out_features=5120, bias=False)
)
(lora_embedding_A): ParameterDict()
(lora_embedding_B): ParameterDict()
(lora_magnitude_vector): ModuleDict()
)
(o_proj): lora.Linear(
(base_layer): Linear(in_features=5120, out_features=5120, bias=False)
(lora_dropout): ModuleDict(
(default): Identity()
)
(lora_A): ModuleDict(
(default): Linear(in_features=5120, out_features=16, bias=False)
)
(lora_B): ModuleDict(
(default): Linear(in_features=16, out_features=5120, bias=False)
)
(lora_embedding_A): ParameterDict()
(lora_embedding_B): ParameterDict()
(lora_magnitude_vector): ModuleDict()
)
(rotary_emb): LlamaLinearScalingRotaryEmbedding()
)
(mlp): LlamaMLP(
(gate_proj): lora.Linear(
(base_layer): Linear(in_features=5120, out_features=13824, bias=False)
(lora_dropout): ModuleDict(
(default): Identity()
)
(lora_A): ModuleDict(
(default): Linear(in_features=5120, out_features=16, bias=False)
)
(lora_B): ModuleDict(
(default): Linear(in_features=16, out_features=13824, bias=False)
)
(lora_embedding_A): ParameterDict()
(lora_embedding_B): ParameterDict()
(lora_magnitude_vector): ModuleDict()
)
(up_proj): lora.Linear(
(base_layer): Linear(in_features=5120, out_features=13824, bias=False)
(lora_dropout): ModuleDict(
(default): Identity()
)
(lora_A): ModuleDict(
(default): Linear(in_features=5120, out_features=16, bias=False)
)
(lora_B): ModuleDict(
(default): Linear(in_features=16, out_features=13824, bias=False)
)
(lora_embedding_A): ParameterDict()
(lora_embedding_B): ParameterDict()
(lora_magnitude_vector): ModuleDict()
)
(down_proj): lora.Linear(
(base_layer): Linear(in_features=13824, out_features=5120, bias=False)
(lora_dropout): ModuleDict(
(default): Identity()
)
(lora_A): ModuleDict(
(default): Linear(in_features=13824, out_features=16, bias=False)
)
(lora_B): ModuleDict(
(default): Linear(in_features=16, out_features=5120, bias=False)
)
(lora_embedding_A): ParameterDict()
(lora_embedding_B): ParameterDict()
(lora_magnitude_vector): ModuleDict()
)
(act_fn): SiLU()
)
(input_layernorm): LlamaRMSNorm((5120,), eps=1e-05)
(post_attention_layernorm): LlamaRMSNorm((5120,), eps=1e-05)
)
)
(norm): LlamaRMSNorm((5120,), eps=1e-05)
(rotary_emb): LlamaRotaryEmbedding()
)
(lm_head): Linear(in_features=5120, out_features=32000, bias=False)
)
)
)
In [7]:
from trl import SFTTrainer
from transformers import TrainingArguments
max_seq_length = 25536
trainer = SFTTrainer(
model = model,
tokenizer = tokenizer,
train_dataset = dataset,
dataset_text_field = "text",
max_seq_length = max_seq_length,
dataset_num_proc = 2,
packing = False, # Can make training 5x faster for short sequences.
args = TrainingArguments(
per_device_train_batch_size = 2,
gradient_accumulation_steps = 4,
warmup_steps = 4,
max_steps = 2, # Use num_train_epochs=1 instead. We use this to make training faster
learning_rate = 2e-4,
fp16 = not torch.cuda.is_bf16_supported(),
bf16 = torch.cuda.is_bf16_supported(),
logging_steps = 1,
optim = "adamw_8bit",
weight_decay = 0.01,
lr_scheduler_type = "linear",
seed = 3407,
output_dir = "outputs",
),
)
Map (num_proc=2): 0%| | 0/12000 [00:00<?, ? examples/s]
max_steps is given, it will override any value given in num_train_epochs
In [8]:
#@title Show current memory stats
gpu_stats = torch.cuda.get_device_properties(0)
start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
print(f"{start_gpu_memory} GB of memory reserved.")
GPU = NVIDIA A100-SXM4-40GB. Max memory = 39.394 GB. 25.162 GB of memory reserved.
In [9]:
trainer_stats = trainer.train()
==((====))== Unsloth - 2x faster free finetuning | Num GPUs = 1 \\ /| Num examples = 12,000 | Num Epochs = 1 O^O/ \_/ \ Batch size per device = 2 | Gradient Accumulation steps = 4 \ / Total batch size = 8 | Total steps = 20 "-____-" Number of trainable parameters = 62,586,880
torch.Size([2, 25536, 5120]) torch.Size([2, 7327, 5120]) torch.Size([2, 10390, 5120]) torch.Size([2, 7596, 5120])
[20/20 22:16, Epoch 0/1]
| Step | Training Loss |
|---|---|
| 1 | 3.766700 |
| 2 | 3.724200 |
| 3 | 3.965900 |
| 4 | 3.507000 |
| 5 | 3.210900 |
| 6 | 2.953400 |
| 7 | 2.868400 |
| 8 | 2.678200 |
| 9 | 2.723600 |
| 10 | 2.845700 |
| 11 | 2.839100 |
| 12 | 2.554300 |
| 13 | 2.861300 |
| 14 | 2.358900 |
| 15 | 2.459000 |
| 16 | 2.356900 |
| 17 | 2.416600 |
| 18 | 2.476900 |
| 19 | 2.433400 |
| 20 | 2.524200 |
torch.Size([2, 25536, 5120]) torch.Size([2, 9019, 5120]) torch.Size([2, 18487, 5120]) torch.Size([2, 6547, 5120]) torch.Size([2, 10562, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 17146, 5120]) torch.Size([2, 9230, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 7845, 5120]) torch.Size([2, 9486, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 7607, 5120]) torch.Size([2, 8356, 5120]) torch.Size([2, 10552, 5120]) torch.Size([2, 8206, 5120]) torch.Size([2, 178, 5120]) torch.Size([2, 8355, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 18572, 5120]) torch.Size([2, 15266, 5120]) torch.Size([2, 97, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 9336, 5120]) torch.Size([2, 14387, 5120]) torch.Size([2, 381, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 14744, 5120]) torch.Size([2, 16753, 5120]) torch.Size([2, 154, 5120]) torch.Size([2, 8847, 5120]) torch.Size([2, 9558, 5120]) torch.Size([2, 10207, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 10274, 5120]) torch.Size([2, 17392, 5120]) torch.Size([2, 7959, 5120]) torch.Size([2, 93, 5120]) torch.Size([2, 18240, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 11225, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 17025, 5120]) torch.Size([2, 8291, 5120]) torch.Size([2, 9287, 5120]) torch.Size([2, 18961, 5120]) torch.Size([2, 130, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 17034, 5120]) torch.Size([2, 8802, 5120]) torch.Size([2, 18578, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 12574, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 17591, 5120]) torch.Size([2, 12382, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 9138, 5120]) torch.Size([2, 9189, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 12923, 5120]) torch.Size([2, 24009, 5120]) torch.Size([2, 8834, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 16868, 5120]) torch.Size([2, 7755, 5120]) torch.Size([2, 8075, 5120]) torch.Size([2, 25536, 5120]) torch.Size([2, 8777, 5120]) torch.Size([2, 18509, 5120]) torch.Size([2, 9561, 5120])
In [10]:
#@title Show final memory and time stats
used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
used_percentage = round(used_memory /max_memory*100, 3)
lora_percentage = round(used_memory_for_lora/max_memory*100, 3)
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.")
print(f"Peak reserved memory = {used_memory} GB.")
print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
print(f"Peak reserved memory % of max memory = {used_percentage} %.")
print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")
1394.9632 seconds used for training. 23.25 minutes used for training. Peak reserved memory = 38.33 GB. Peak reserved memory for training = 13.168 GB. Peak reserved memory % of max memory = 97.299 %. Peak reserved memory for training % of max memory = 33.426 %.
In [ ]: