Skip to content

Instantly share code, notes, and snippets.

@abhishekkrthakur
Created July 16, 2023 09:09
Show Gist options
  • Save abhishekkrthakur/82a09c9428d48a343c5911ffb701c1d4 to your computer and use it in GitHub Desktop.
Save abhishekkrthakur/82a09c9428d48a343c5911ffb701c1d4 to your computer and use it in GitHub Desktop.
Train LLMs in 50 lines of code. This is a reference code for YouTube tutorial: https://www.youtube.com/watch?v=JNMVulH7fCo&ab_channel=AbhishekThakur
import torch
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer
def train():
train_dataset = load_dataset("tatsu-lab/alpaca", split="train")
tokenizer = AutoTokenizer.from_pretrained("Salesforce/xgen-7b-8k-base", trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
"Salesforce/xgen-7b-8k-base", load_in_4bit=True, torch_dtype=torch.float16, device_map="auto"
)
model.resize_token_embeddings(len(tokenizer))
model = prepare_model_for_int8_training(model)
peft_config = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM")
model = get_peft_model(model, peft_config)
training_args = TrainingArguments(
output_dir="xgen-7b-tuned-alpaca-l1",
per_device_train_batch_size=4,
optim="adamw_torch",
logging_steps=100,
learning_rate=2e-4,
fp16=True,
warmup_ratio=0.1,
lr_scheduler_type="linear",
num_train_epochs=1,
save_strategy="epoch",
push_to_hub=True,
)
trainer = SFTTrainer(
model=model,
train_dataset=train_dataset,
dataset_text_field="text",
max_seq_length=1024,
tokenizer=tokenizer,
args=training_args,
packing=True,
peft_config=peft_config,
)
trainer.train()
trainer.push_to_hub()
if __name__ == "__main__":
train()
@srikant86panda
Copy link

I also have this same issue using: transformers: 4.30.2
bitsandbytes: 0.40.0. Have created bitsandbytes-foundation/bitsandbytes#600

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment