Тюнинг LLM
и энкодеров
достаточно подробно

# Установка (важен порядок — unsloth зависит от flash-attn версии)
# pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
# pip install --no-deps trl peft accelerate bitsandbytes

from unsloth import FastLanguageModel
import torch

# ── 1. Загрузка модели через Unsloth ─────────────────────────
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="Qwen/Qwen2.5-7B-Instruct",
    max_seq_length=2048,           # RoPE масштабируется автоматически
    dtype=None,                    # None = auto (BF16 на Ampere+, FP16 на V100)
    load_in_4bit=True,            # QLoRA: NF4 квантизация
    # token="hf_..." если приватная модель
)

# ── 2. LoRA конфигурация через Unsloth ───────────────────────
model = FastLanguageModel.get_peft_model(
    model,
    r=16,                          # LoRA rank (4, 8, 16, 32, 64)
    target_modules=[               # Qwen2.5 attention + FFN matrices
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_alpha=32,                 # scaling = alpha/r = 2.0
    lora_dropout=0,               # Unsloth рекомендует 0 для скорости
    bias="none",
    use_gradient_checkpointing="unsloth",  # Unsloth custom GC
    random_state=42,
    use_rslora=False,             # RSLoRA: масштаб = alpha/sqrt(r)
    loftq_config=None,            # LoftQ инициализация (опционально)
)

print(model.print_trainable_parameters())
# trainable params: 40,108,032 || all: 7,241,732,096 || trainable%: 0.5539

from unsloth import FastLanguageModel, is_bfloat16_supported
from datasets import load_dataset
from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
from transformers import TrainingArguments
import torch

# ── 1. Модель и токенизатор ───────────────────────────────────
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="Qwen/Qwen2.5-7B-Instruct",
    max_seq_length=2048,
    dtype=None,
    load_in_4bit=True,
)
model = FastLanguageModel.get_peft_model(
    model, r=16, lora_alpha=32,
    target_modules=["q_proj","k_proj","v_proj","o_proj",
                    "gate_proj","up_proj","down_proj"],
    use_gradient_checkpointing="unsloth", random_state=42,
)

# ── 2. Датасет и форматирование ──────────────────────────────
dataset = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft[:5000]")

def format_chat(example):
    """Конвертируем conversation в ChatML строку для Qwen2.5."""
    messages = example["messages"]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False,
    )
    return {"text": text}

dataset = dataset.map(format_chat, num_proc=4)

# Пример форматированного текста:
# <|im_start|>system
# You are a helpful assistant.<|im_end|>
# <|im_start|>user
# What is photosynthesis?<|im_end|>
# <|im_start|>assistant
# Photosynthesis is the process by which...<|im_end|>

# ── 3. Label masking: loss только по assistant ───────────────
# Qwen2.5 response template:
response_template_ids = tokenizer.encode(
    "<|im_start|>assistant\n",
    add_special_tokens=False
)
collator = DataCollatorForCompletionOnlyLM(
    response_template=response_template_ids,
    tokenizer=tokenizer,
)

# ── 4. TrainingArguments ─────────────────────────────────────
training_args = TrainingArguments(
    output_dir="./qwen2.5-7b-sft",
    num_train_epochs=3,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,     # eff_batch=16
    warmup_ratio=0.03,
    learning_rate=2e-4,               # QLoRA: выше чем full FT
    lr_scheduler_type="cosine",
    optim="adamw_bnb_8bit",           # 8-bit Adam из bitsandbytes
    bf16=is_bfloat16_supported(),
    fp16=not is_bfloat16_supported(),
    logging_steps=10,
    save_strategy="steps",
    save_steps=200,
    eval_strategy="steps",
    eval_steps=200,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    report_to="wandb",               # W&B трекинг
    run_name="qwen2.5-7b-qlora-sft",
)

# ── 5. SFTTrainer ─────────────────────────────────────────────
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=dataset,
    dataset_text_field="text",
    data_collator=collator,
    max_seq_length=2048,
    dataset_num_proc=4,
    packing=False,                   # False при label masking
)

trainer.train()

from unsloth import FastLanguageModel

# ── Вариант 1: Merge и сохранение в HuggingFace формат ───────
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="./qwen2.5-7b-sft/checkpoint-best",  # наш чекпоинт
    max_seq_length=2048, dtype=None, load_in_4bit=True,
)

# Merge LoRA → full model в BF16 (14 GB)
model.save_pretrained_merged(
    "./qwen2.5-7b-merged",         # куда сохранить
    tokenizer,
    save_method="merged_16bit",   # merged_16bit | merged_4bit | lora
)
# Варианты save_method:
# "lora"          → только адаптер (~80 MB)
# "merged_16bit"  → merged в BF16 (~14 GB)
# "merged_4bit"   → merged + quantize обратно в 4bit (~3.5 GB)

# ── Вариант 2: Export в GGUF для llama.cpp ────────────────────
model.save_pretrained_gguf(
    "./qwen2.5-7b-gguf",
    tokenizer,
    quantization_method="q4_k_m",   # q4_k_m = 4-bit K-Quant Mixed
    # Другие варианты: q2_k, q3_k_m, q5_k_m, q8_0, f16
)
# Результат: qwen2.5-7b-Q4_K_M.gguf (~4.2 GB)
# Запуск: ./llama.cpp/main -m qwen2.5-7b-Q4_K_M.gguf

# ── Вариант 3: Upload на HuggingFace Hub ─────────────────────
# model.push_to_hub("username/qwen2.5-7b-sft", token="hf_...")
# tokenizer.push_to_hub("username/qwen2.5-7b-sft")

# ── Ручной merge через PEFT (если не Unsloth) ─────────────────
from peft import PeftModel
from transformers import AutoModelForCausalLM
import torch

base = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct", torch_dtype=torch.bfloat16
)
peft_model = PeftModel.from_pretrained(base, "./lora-adapter")
merged = peft_model.merge_and_unload()   # W = W_frozen + A·B·scale
merged.save_pretrained("./qwen-merged-peft")

# Установка
pip install lm-eval

# ── Базовый запуск: оценка после SFT ─────────────────────────
lm_eval \
    --model hf \
    --model_args pretrained=./qwen2.5-7b-merged,dtype=bfloat16 \
    --tasks mmlu,hellaswag,arc_easy,arc_challenge,gsm8k \
    --num_fewshot 5 \
    --batch_size 8 \
    --output_path ./eval_results/qwen_sft \
    --log_samples                   # сохранить отдельные примеры

# ── Сравнение base vs finetuned ────────────────────────────── 
lm_eval \
    --model hf \
    --model_args pretrained=Qwen/Qwen2.5-7B-Instruct,dtype=bfloat16 \
    --tasks mmlu \
    --num_fewshot 5 \
    --output_path ./eval_results/qwen_base

# ── Быстрая проверка на 1 GPU (4-bit для скорости) ──────────
lm_eval \
    --model hf \
    --model_args "pretrained=./qwen2.5-7b-merged,load_in_4bit=True" \
    --tasks arc_easy,hellaswag \
    --num_fewshot 0 \
    --batch_size auto

import lm_eval
import wandb, json

# Запуск из Python (удобно для пайплайна)
results = lm_eval.simple_evaluate(
    model="hf",
    model_args="pretrained=./qwen2.5-7b-merged,dtype=bfloat16",
    tasks=["mmlu", "hellaswag", "gsm8k"],
    num_fewshot=5,
    batch_size=8,
)

# Парсим результаты
for task, metrics in results["results"].items():
    acc = metrics.get("acc,none", metrics.get("exact_match,none", None))
    print(f"{task:30s}: {acc:.4f}")

# mmlu                          : 0.7234
# hellaswag                     : 0.8156
# gsm8k                         : 0.6543

# Логируем в W&B
wandb.init(project="qwen-eval", name="sft-checkpoint-best")
flat_metrics = {}
for task, metrics in results["results"].items():
    for metric, val in metrics.items():
        if isinstance(val, (int, float)):
            flat_metrics[f"eval/{task}/{metric}"] = val
wandb.log(flat_metrics)
wandb.finish()

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import torch.nn as nn

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct", torch_dtype=torch.bfloat16
)

old_vocab_size = len(tokenizer)   # 151,936 в Qwen2.5
print(f"Исходный словарь: {old_vocab_size} токенов")

# ── 1. Добавляем специальные/доменные токены ──────────────────
new_special = ["<|tool_start|>", "<|tool_end|>", "<|reasoning|>"]
new_regular = ["гиперглипсемия", "нейроинтерфейс", "квантовычисление"]

# Special tokens: добавляем через add_special_tokens
tokenizer.add_special_tokens({"additional_special_tokens": new_special})
# Regular tokens: add_tokens
tokenizer.add_tokens(new_regular)

new_vocab_size = len(tokenizer)
num_new = new_vocab_size - old_vocab_size
print(f"Новый словарь: {new_vocab_size} токенов (+{num_new})")

# ── 2. Resize embedding matrix и lm_head ─────────────────────
# ВАЖНО: resize_token_embeddings обновляет ОБА: embed_tokens и lm_head
model.resize_token_embeddings(new_vocab_size, pad_to_multiple_of=64)
# pad_to_multiple_of=64: выравнивание для эффективных CUDA kernel

# ── 3. Инициализация новых эмбеддингов ───────────────────────
with torch.no_grad():
    embed_weight = model.model.embed_tokens.weight
    lm_head_weight = model.lm_head.weight

    for i, tok in enumerate(new_special + new_regular):
        new_id = old_vocab_size + i

        # Вариант A: среднее существующих токенов → хорошо для регулярных
        sub_ids = tokenizer.encode(tok, add_special_tokens=False)
        sub_ids = [x for x in sub_ids if x < old_vocab_size]

        if sub_ids:
            mean_emb = embed_weight[sub_ids].mean(0)
            embed_weight[new_id] = mean_emb
            lm_head_weight[new_id] = mean_emb
        else:
            # Вариант B: случайная инициализация со std базовых эмбеддингов
            std = embed_weight[:old_vocab_size].std()
            embed_weight[new_id] = torch.randn_like(embed_weight[0]) * std

# ── 4. Заморозка старых эмбеддингов: учим только новые ───────
# Вместо полного разморозки embeddings — используем hook
def freeze_old_embeddings_hook(grad: torch.Tensor) -> torch.Tensor:
    """Обнуляем градиенты для старых токенов."""
    grad[:old_vocab_size] = 0
    return grad

model.model.embed_tokens.weight.register_hook(freeze_old_embeddings_hook)
model.lm_head.weight.register_hook(freeze_old_embeddings_hook)

# Теперь обучаем модель через LoRA + только новые эмбеддинги
# Остальные параметры заморожены LoRA как обычно
tokenizer.save_pretrained("./qwen-extended-tok")
print(f"Проверка: {tokenizer.encode('<|tool_start|>')}")
# [151936]  ← один токен вместо разбиения на части

from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType
import torch

# ── QLoRA для энкодера ────────────────────────────────────────
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# Используем multilingual-e5-large для русского SberQuAD
MODEL_NAME = "intfloat/multilingual-e5-large"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
encoder = AutoModel.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
)

# LoRA для энкодера: target_modules зависят от архитектуры
# У RoBERTa-based (E5): query, key, value, output в attention
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["query", "key", "value", "dense"],
    lora_dropout=0.1,
    bias="none",
    task_type=TaskType.FEATURE_EXTRACTION,  # для энкодеров
)

encoder = get_peft_model(encoder, lora_config)
encoder.print_trainable_parameters()
# trainable params: 1,179,648 || all: 559,890,432 || trainable%: 0.2107

# ── Pooling: mean pooling для dense retrieval ─────────────────
def mean_pooling(model_output, attention_mask):
    token_embs = model_output.last_hidden_state
    mask = attention_mask.unsqueeze(-1).float()
    return (token_embs * mask).sum(1) / mask.sum(1).clamp(min=1e-9)

def encode(texts, prefix=""):
    """E5 требует prefix: 'query: ' или 'passage: '"""
    texts = [prefix + t for t in texts]
    enc = tokenizer(texts, padding=True, truncation=True,
                    max_length=512, return_tensors="pt")
    with torch.no_grad():
        out = encoder(**enc)
    embs = mean_pooling(out, enc["attention_mask"])
    return torch.nn.functional.normalize(embs, p=2, dim=1)

from datasets import load_dataset
from torch.utils.data import Dataset, DataLoader
import torch
import torch.nn.functional as F
import faiss, numpy as np
from tqdm import tqdm

# ── 1. Загрузка SberQuAD ──────────────────────────────────────
# SberQuAD: ~45k вопросов из русской Википедии
raw = load_dataset("sberquad")
# Структура: {'id', 'context', 'question', 'answers': {'text', 'answer_start'}}

print(f"Train: {len(raw['train'])} | Val: {len(raw['validation'])}")
# Train: 43561 | Val: 5036
print(raw['train'][0]['question'])
# Какова была цель эксперимента Гейгера-Марсдена?
print(raw['train'][0]['context'][:100])
# В 1909 году, через Рутерфорда, Гейгер и Марсден провели эксперимент...

# ── 2. Dataset для MNRL ───────────────────────────────────────
class SberQuADRetrieval(Dataset):
    def __init__(self, data):
        self.pairs = []
        seen_contexts = set()
        for ex in data:
            q = ex['question']
            p = ex['context']
            # Один контекст может встречаться для нескольких вопросов
            self.pairs.append({'query': q, 'passage': p})

    def __len__(self): return len(self.pairs)
    def __getitem__(self, i): return self.pairs[i]

train_ds = SberQuADRetrieval(raw["train"])

# ── 3. Multiple Negatives Ranking Loss ────────────────────────
class MNRLoss(torch.nn.Module):
    def __init__(self, temperature: float = 0.02):
        super().__init__()
        self.temp = temperature
        self.ce = torch.nn.CrossEntropyLoss()

    def forward(self, q_emb, p_emb):
        """
        q_emb: (B, D) — query embeddings
        p_emb: (B, D) — passage embeddings (i-й passage = позитив к i-му query)
        In-batch negatives: passage j≠i автоматически негатив для query i
        """
        # Нормализуем для cosine similarity
        q = F.normalize(q_emb, p=2, dim=-1)
        p = F.normalize(p_emb, p=2, dim=-1)

        # Similarity matrix: (B, B)
        scores = torch.mm(q, p.T) / self.temp

        # Метки: диагональ — правильные пары
        labels = torch.arange(scores.size(0), device=scores.device)

        # CrossEntropy: q_i → p_i, а не p_j (j≠i)
        return self.ce(scores, labels)

# ── 4. Тренировочный цикл ─────────────────────────────────────
optimizer = torch.optim.AdamW(encoder.parameters(), lr=2e-5)
loss_fn = MNRLoss(temperature=0.02)

def collate_fn(batch):
    queries   = [b['query']   for b in batch]
    passages  = [b['passage'] for b in batch]
    q_enc = tokenizer(queries,  padding=True, truncation=True,
                       max_length=64,  return_tensors="pt")
    p_enc = tokenizer(passages, padding=True, truncation=True,
                       max_length=512, return_tensors="pt")
    return q_enc, p_enc

loader = DataLoader(train_ds, batch_size=64,  # больший батч = больше негативов
                     shuffle=True, collate_fn=collate_fn, num_workers=4)

encoder.train()
for epoch in range(3):
    total_loss = 0
    for q_enc, p_enc in tqdm(loader):
        q_out = encoder(**q_enc).last_hidden_state[:, 0]  # CLS token
        p_out = encoder(**p_enc).last_hidden_state[:, 0]
        loss = loss_fn(q_out, p_out)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, loss={total_loss/len(loader):.4f}")

# ── Hit@K: доля запросов где правильный passage в топ-K ──────
def evaluate_hit_at_k(encoder, tokenizer, val_data, k: int = 5,
                      batch_size: int = 64) -> float:
    """
    1. Собираем корпус уникальных passage
    2. Эмбеддируем все passages → FAISS index
    3. Для каждого query: ищем топ-K passages
    4. Hit = 1 если правильный passage в топ-K
    """
    encoder.eval()

    # Уникальные passages (один контекст = одна запись в индексе)
    passages = list({ex['context']: None for ex in val_data}.keys())
    passage2idx = {p: i for i, p in enumerate(passages)}

    # Эмбеддируем passages батчами
    all_p_embs = []
    for i in range(0, len(passages), batch_size):
        batch_p = passages[i:i+batch_size]
        enc = tokenizer(batch_p, padding=True, truncation=True,
                         max_length=512, return_tensors="pt")
        with torch.no_grad():
            out = encoder(**enc).last_hidden_state[:, 0]
            out = F.normalize(out, p=2, dim=-1)
        all_p_embs.append(out.cpu().numpy())
    p_matrix = np.vstack(all_p_embs)

    # FAISS IndexFlatIP для cosine (векторы нормированы)
    index = faiss.IndexFlatIP(p_matrix.shape[1])
    index.add(p_matrix)

    # Оцениваем Hit@K
    hits = 0
    total = 0
    for i in range(0, len(val_data), batch_size):
        batch = val_data[i:i+batch_size]
        queries = [ex['question'] for ex in batch]
        correct_ids = [passage2idx[ex['context']] for ex in batch]

        enc = tokenizer(queries, padding=True, truncation=True,
                         max_length=64, return_tensors="pt")
        with torch.no_grad():
            q_emb = encoder(**enc).last_hidden_state[:, 0]
            q_emb = F.normalize(q_emb, p=2, dim=-1).cpu().numpy()

        _, I = index.search(q_emb, k)   # I: (batch, K) — indices топ-K
        for j, correct in enumerate(correct_ids):
            if correct in I[j]:
                hits += 1
        total += len(batch)

    return hits / total

val_list = list(raw["validation"])
for K in [1, 5, 10]:
    score = evaluate_hit_at_k(encoder, tokenizer, val_list, k=K)
    print(f"Hit@{K}: {score:.4f}")
# До обучения (базовый multilingual-e5-large):
# Hit@1:  0.6123 | Hit@5:  0.8234 | Hit@10: 0.8891
# После QLoRA MNRL (3 эпохи, bs=64):
# Hit@1:  0.7456 | Hit@5:  0.9012 | Hit@10: 0.9423
# Прирост Hit@5: +7.8% — значимое улучшение!

import wandb
import os, json, hashlib
from pathlib import Path

# ── 1. Инициализация W&B run ──────────────────────────────────
# WANDB_API_KEY в переменных окружения или wandb login
run = wandb.init(
    project="qwen-sft-experiments",
    name="qwen2.5-7b-qlora-r16-ultrachat",
    config={
        # Все гиперпараметры — будут доступны для сравнения
        "model":          "Qwen/Qwen2.5-7B-Instruct",
        "dataset":        "HuggingFaceH4/ultrachat_200k",
        "lora_r":         16,
        "lora_alpha":     32,
        "quantization":   "nf4",
        "learning_rate":  2e-4,
        "batch_size":     2,
        "grad_accum":     8,
        "epochs":         3,
        "max_seq_len":    2048,
        "optimizer":      "adamw_bnb_8bit",
        "scheduler":      "cosine",
        "warmup_ratio":   0.03,
    },
    tags=["sft", "qlora", "qwen2.5", "instruction-tuning"],
    notes="Baseline SFT на ultrachat. Следующий шаг: добавить hard negatives.",
)

# ── 2. Трекинг через HuggingFace Trainer (автоматически) ─────
# В TrainingArguments: report_to="wandb" → метрики логируются автоматически
# Логируется: train_loss, eval_loss, learning_rate, grad_norm, VRAM

# ── 3. Ручное логирование дополнительных метрик ───────────────
class WandbCallback:
    """Кастомный callback для дополнительного логирования."""
    def on_evaluate(self, args, state, control, metrics, **kwargs):
        # Логируем GPU utilization и VRAM
        try:
            import pynvml
            pynvml.nvmlInit()
            handle = pynvml.nvmlDeviceGetHandleByIndex(0)
            info = pynvml.nvmlDeviceGetMemoryInfo(handle)
            wandb.log({
                "gpu/vram_used_gb":  info.used / 1e9,
                "gpu/vram_total_gb": info.total / 1e9,
                "gpu/util_pct": pynvml.nvmlDeviceGetUtilizationRates(handle).gpu,
            })
        except: pass

# ── 4. Версионирование артефактов (чекпоинты, адаптеры) ──────
def save_adapter_artifact(checkpoint_dir: str, step: int):
    """Сохраняет LoRA адаптер как W&B артефакт с версией."""
    artifact = wandb.Artifact(
        name=f"qwen2.5-lora-adapter",
        type="model",
        description=f"LoRA adapter step {step}",
        metadata={"step": step, "framework": "peft"},
    )
    artifact.add_dir(checkpoint_dir)  # загружает всю папку чекпоинта
    wandb.run.log_artifact(artifact)
    print(f"Artifact logged: step={step}")

# ── 5. Хеш конфигурации для воспроизводимости ────────────────
def config_hash(config: dict) -> str:
    """Детерминированный хеш конфигурации."""
    config_str = json.dumps(config, sort_keys=True)
    return hashlib.md5(config_str.encode()).hexdigest()[:8]

run_hash = config_hash(dict(wandb.config))
print(f"Config hash: {run_hash}")  # e.g. "a3f7c891"
# Сохраняем в артефакт для точного воспроизведения

# ── 6. Воспроизводимость: seed и окружение ───────────────────
def set_seed(seed: int = 42):
    import random
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # Детерминизм cuDNN (медленнее, но воспроизводимо)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)

# Логируем версии зависимостей
import pkg_resources
deps = {'transformers', 'peft', 'trl', 'bitsandbytes', 'torch', 'unsloth'}
versions = {
    pkg: pkg_resources.get_distribution(pkg).version
    for pkg in deps
}
wandb.config.update({"deps": versions})
print(json.dumps(versions, indent=2))
# {"transformers": "4.45.2", "peft": "0.12.0",
#  "trl": "0.11.4", "bitsandbytes": "0.44.1",
#  "torch": "2.4.1", "unsloth": "2024.10.7"}