File size: 2,910 Bytes
38b4eff | 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 |
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, TrainingArguments, Trainer, DataCollatorForLanguageModeling, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset
def train():
model_id = 'HuggingFaceTB/SmolLM2-360M-Instruct'
# 1. Load Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token
# 2. Model & Quantization
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type='nf4',
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True
)
config = AutoConfig.from_pretrained(model_id)
config.use_cache = False
config._attn_implementation = 'sdpa'
model = AutoModelForCausalLM.from_pretrained(
model_id,
config=config,
quantization_config=bnb_config,
device_map='auto',
trust_remote_code=True
)
model = prepare_model_for_kbit_training(model)
# 3. LoRA Setup (Targeting more modules for better reasoning)
peft_config = LoraConfig(
r=32,
lora_alpha=64,
target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],
lora_dropout=0.05,
bias='none',
task_type='CAUSAL_LM'
)
model = get_peft_model(model, peft_config)
# 4. Load Authentic Ground-Truth Data (UltraChat 200k subset)
# This dataset covers 100+ topics and complex instructions
dataset = load_dataset('HuggingFaceH4/ultrachat_200k', split='train_sft[:2000]')
def tokenize(example):
# UltraChat uses the 'messages' format which apply_chat_template handles perfectly
text = tokenizer.apply_chat_template(example['messages'], tokenize=False)
return tokenizer(text, truncation=True, max_length=1024, padding='max_length')
tokenized_dataset = dataset.map(tokenize, remove_columns=dataset.column_names)
# 5. Scaled Training Arguments
args = TrainingArguments(
output_dir='./checkpoints_advanced',
num_train_epochs=1, # 1 epoch on 2000 high-quality samples is substantial for 360M model
per_device_train_batch_size=2,
gradient_accumulation_steps=8,
learning_rate=1e-4,
fp16=True,
logging_steps=20,
save_strategy='no',
optim='paged_adamw_32bit',
lr_scheduler_type='cosine'
)
trainer = Trainer(
model=model,
train_dataset=tokenized_dataset,
args=args,
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
)
print('Starting advanced training on 2000 ground-truth samples...')
trainer.train()
model.save_pretrained('./checkpoints/final_model_advanced')
print('Advanced Training complete!')
if __name__ == "__main__":
train()
|