train.py

"""
Phi-3-mini QLoRA Training Script
GTX 1650 4GB + 16GB RAM uchun optimallashtirilgan
"""

import os
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig

# ============== KONFIGURATSIYA ==============
MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"
OUTPUT_DIR = "./procurement-lora"
DATA_FILE = "train_data.jsonl"

# LoRA parametrlari (3000 samples uchun optimallashtirilgan)
LORA_R = 16  # Kattaroq rank = ko'proq o'rganish qobiliyati
LORA_ALPHA = 32
LORA_DROPOUT = 0.05

# Training parametrlari
BATCH_SIZE = 1
GRADIENT_ACCUMULATION = 8  # Kattaroq effective batch
LEARNING_RATE = 1e-4  # Kichikroq LR ko'p data uchun
NUM_EPOCHS = 5  # Ko'proq epochs
MAX_SEQ_LENGTH = 1024

# ============================================

def create_prompt(example):
    """Training prompt formatlash."""
    instruction = example.get("instruction", "")
    input_text = example.get("input", "")
    output = example.get("output", "")

    if input_text:
        text = f"""<|user|>
{instruction}

Ma'lumot:
{input_text}<|end|>
<|assistant|>
{output}<|end|>"""
    else:
        text = f"""<|user|>
{instruction}<|end|>
<|assistant|>
{output}<|end|>"""

    return {"text": text}


def main():
    print("=" * 50)
    print("Phi-3 QLoRA Training")
    print("=" * 50)

    # GPU tekshirish
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    else:
        print("WARNING: GPU topilmadi! CPU da training juda sekin bo'ladi.")

    # 4-bit quantization konfiguratsiyasi
    print("\n[1/5] 4-bit quantization sozlanmoqda...")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,  # GTX 1650 uchun float16
    )

    # Tokenizer yuklash
    print("\n[2/5] Tokenizer yuklanmoqda...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    # Model yuklash
    print("\n[3/5] Model yuklanmoqda (4-bit)...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        attn_implementation="eager",  # Flash attention o'rniga oddiy attention
    )
    model.config.use_cache = False

    # Gradient checkpointing
    model.gradient_checkpointing_enable()
    model = prepare_model_for_kbit_training(model)

    # LoRA konfiguratsiyasi
    print("\n[4/5] LoRA adapter sozlanmoqda...")
    lora_config = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=[
            "qkv_proj",
            "o_proj",
            "gate_up_proj",
            "down_proj",
        ],
    )

    model = get_peft_model(model, lora_config)

    # Trainable parametrlarni ko'rsatish
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Trainable: {trainable_params:,} / {total_params:,} ({100 * trainable_params / total_params:.2f}%)")

    # Dataset yuklash
    print("\n[5/5] Dataset yuklanmoqda...")
    dataset = load_dataset("json", data_files=DATA_FILE, split="train")
    dataset = dataset.map(create_prompt)
    print(f"Training samples: {len(dataset)}")

    # SFT konfiguratsiyasi (yangi TRL versiya uchun)
    sft_config = SFTConfig(
        output_dir=OUTPUT_DIR,
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRADIENT_ACCUMULATION,
        learning_rate=LEARNING_RATE,
        weight_decay=0.01,
        logging_steps=1,
        save_strategy="epoch",
        save_total_limit=2,
        fp16=False,  # GTX 1650 uchun FP32 ishlatamiz
        optim="paged_adamw_8bit",
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        report_to="none",
        gradient_checkpointing=True,
        max_grad_norm=0.3,
        max_length=MAX_SEQ_LENGTH,
        dataset_text_field="text",
        packing=False,
    )

    # Trainer
    trainer = SFTTrainer(
        model=model,
        args=sft_config,
        train_dataset=dataset,
        processing_class=tokenizer,
    )

    # Training
    print("\n" + "=" * 50)
    print("TRAINING BOSHLANDI")
    print("=" * 50)

    trainer.train()

    # Adapter saqlash
    print("\n" + "=" * 50)
    print("Model saqlanmoqda...")
    print("=" * 50)

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

    print(f"\nAdapter saqlandi: {OUTPUT_DIR}/")
    print("Training yakunlandi!")

    # VRAM tozalash
    del model
    del trainer
    torch.cuda.empty_cache()


if __name__ == "__main__":
    main()