File size: 2,953 Bytes
53aeb2a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from peft import LoraConfig, TaskType
from trl import SFTTrainer, SFTConfig
import trackio
model_name = "./SmolLM3-3B-Base/"
dataset_path = "./MathInstruct/MathInstruct.json"
output_dir = "./SmolLMathematician-3B"
project_name = "SmolLMathematician-3B"
MAX_SEQ_LENGTH = 4096
trackio.init(project=project_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map="auto",
dtype=torch.bfloat16,
low_cpu_mem_usage=True,
trust_remote_code=True,
attn_implementation="flash_attention_2",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = model.config.eos_token_id
with open("chat_template.jinja", "r") as f:
chat_template = f.read()
tokenizer.chat_template = chat_template
model.gradient_checkpointing_enable()
dataset = load_dataset("json", data_files=dataset_path, split="train")
def formatInstructionWithTemplate(example: dict) -> str:
messages = [
{"role": "user", "content": example["instruction"]},
{"role": "assistant", "content": example["output"]},
]
return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
def checkSequenceLength(example: dict) -> bool:
formatted_text = formatInstructionWithTemplate(example)
tokens = tokenizer(formatted_text)
return len(tokens['input_ids']) <= MAX_SEQ_LENGTH
original_size = len(dataset)
train_dataset = dataset.filter(checkSequenceLength)
new_size = len(train_dataset)
print(f"Dataset: {original_size} → {new_size} samples (removed: {original_size - new_size})")
torch.cuda.empty_cache()
peft_config = LoraConfig(
r=16,
lora_alpha=32,
lora_dropout=0.1,
target_modules=['q_proj', 'v_proj'],
bias="none",
task_type=TaskType.CAUSAL_LM,
)
training_args = SFTConfig(
output_dir=output_dir,
num_train_epochs=1,
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
optim="paged_adamw_8bit",
learning_rate=2e-5,
weight_decay=0.01,
adam_epsilon=1e-6,
max_grad_norm=1.0,
lr_scheduler_type="cosine",
warmup_ratio=0.1,
logging_steps=8,
eval_strategy="no",
save_strategy="steps",
save_steps=32,
save_total_limit=4,
resume_from_checkpoint=True,
report_to="trackio",
bf16=True,
packing=True,
max_length=MAX_SEQ_LENGTH,
dataloader_pin_memory=False,
gradient_checkpointing_kwargs={"use_reentrant": False},
)
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=train_dataset,
peft_config=peft_config,
formatting_func=formatInstructionWithTemplate,
)
torch.cuda.empty_cache()
trainer.train()
torch.cuda.empty_cache()
trainer.save_model(output_dir)
print(f"LoRA adapter saved to {output_dir}")
trackio.finish()
|