tech3space3-0.6B / full_fine_tunning_model_SFTTraining..py
ankitkushwaha90's picture
Upload full_fine_tunning_model_SFTTraining..py
de11e8f verified
Raw
History Blame Contribute Delete
3.11 kB
# train_fft.py
import os
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig
# ============================================================
# PATHS
# ============================================================
MODEL_PATH = "./Qwen3-0.6B"
DATASET_PATH = "./dataset/train.jsonl"
OUTPUT_DIR = "./outputs/qwen3_0.6b_fft"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# ============================================================
# DATASET
# ============================================================
print("Loading dataset...")
dataset = load_dataset(
"json",
data_files=DATASET_PATH,
split="train"
)
print(f"Dataset size: {len(dataset)}")
# ============================================================
# TOKENIZER
# ============================================================
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
MODEL_PATH,
trust_remote_code=True
)
# Fix for models that don't have pad_token
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# ============================================================
# MODEL
# ============================================================
print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
MODEL_PATH,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
device_map="auto" # Optional: helps with memory on single GPU
)
model.config.use_cache = False
# ============================================================
# TRAINING CONFIG
# ============================================================
training_args = SFTConfig(
output_dir=OUTPUT_DIR,
num_train_epochs=3,
learning_rate=5e-6,
per_device_train_batch_size=1,
gradient_accumulation_steps=16,
bf16=True,
logging_steps=10,
save_strategy="steps",
save_steps=200,
save_total_limit=2,
lr_scheduler_type="cosine",
warmup_ratio=0.03,
max_length=512,
packing=False,
gradient_checkpointing=True,
report_to="none",
# Optional but recommended:
dataloader_num_workers=2,
remove_unused_columns=False,
)
# ============================================================
# TRAINER
# ============================================================
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=dataset,
processing_class=tokenizer, # Newer TRL uses processing_class
# tokenizer=tokenizer, # You can use this if processing_class doesn't work
)
# ============================================================
# TRAIN
# ============================================================
print("Starting full fine-tuning...")
trainer.train()
# ============================================================
# SAVE MODEL
# ============================================================
print("Saving model...")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
print("=" * 60)
print("✅ FULL FINE TUNING COMPLETED")
print(f"Model saved to: {OUTPUT_DIR}")
print("=" * 60)