Qwen3-1.7BSFTLoRA/Finetunning_lora_SFTTrainer.py

#!/usr/bin/env python3

import os
import torch

from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
)

from peft import (
    LoraConfig,
    get_peft_model,
)

from trl import (
    SFTTrainer,
    SFTConfig,
)


# ============================================================
# CONFIG
# ============================================================

MODEL_NAME = "./Qwen3-1.7B"
DATASET_PATH = "./dataset/train.jsonl"

OUTPUT_DIR = "./qwen3_lora_sft_pro_try"
MERGED_DIR = OUTPUT_DIR + "_merged"

MAX_LENGTH = 512

os.makedirs(OUTPUT_DIR, exist_ok=True)


# ============================================================
# TOKENIZER
# ============================================================

print("Loading tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

tokenizer.padding_side = "right"


# ============================================================
# MODEL
# ============================================================

print("Loading model...")

dtype = (
    torch.bfloat16
    if torch.cuda.is_available()
    and torch.cuda.is_bf16_supported()
    else torch.float16
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    torch_dtype=dtype,
    device_map="auto",
)

model.config.use_cache = False
model.gradient_checkpointing_enable()


# ============================================================
# LORA CONFIG
# ============================================================

print("Applying LoRA...")

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
)

model = get_peft_model(
    model,
    lora_config,
)

model.print_trainable_parameters()


# ============================================================
# DATASET
# ============================================================

print("Loading dataset...")

dataset = load_dataset(
    "json",
    data_files=DATASET_PATH,
    split="train",
)


def preprocess(example):
    text = tokenizer.apply_chat_template(
        example["messages"],
        tokenize=False,
        add_generation_prompt=False,
    )

    return {
        "text": text
    }


dataset = dataset.map(
    preprocess,
    remove_columns=dataset.column_names,
)

print(dataset)


# ============================================================
# TRAINING CONFIG
# ============================================================

training_args = SFTConfig(
    output_dir=OUTPUT_DIR,

    max_length=MAX_LENGTH,
    packing=True,

    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,

    num_train_epochs=3,

    learning_rate=2e-4,
    warmup_ratio=0.03,

    lr_scheduler_type="cosine",
    weight_decay=0.05,

    logging_steps=10,

    save_steps=500,
    save_total_limit=3,

    bf16=torch.cuda.is_available()
    and torch.cuda.is_bf16_supported(),

    fp16=not (
        torch.cuda.is_available()
        and torch.cuda.is_bf16_supported()
    ),

    report_to="none",

    dataset_text_field="text",
)


# ============================================================
# TRAINER
# ============================================================

print("Initializing trainer...")

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
)


# ============================================================
# TRAIN
# ============================================================

print("Starting training...")

trainer.train()

print("Training completed.")


# ============================================================
# SAVE LORA ADAPTER
# ============================================================

print("Saving LoRA adapter...")

trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)


# ============================================================
# MERGE LORA
# ============================================================

print("Merging LoRA into base model...")

merged_model = model.merge_and_unload()

merged_model.save_pretrained(
    MERGED_DIR,
    safe_serialization=True,
)

tokenizer.save_pretrained(
    MERGED_DIR,
)

print(f"Merged model saved to: {MERGED_DIR}")

print("Done.")