Continued pretraining cu replay: adaptezi un LLM la texte științifice fără să-i strici baza generală

Acest tutorial construiește un experiment complet de continued pretraining: pornim de la distilgpt2, îl adaptăm pe texte științifice din SciTLDR și păstrăm o porțiune de date generale WikiText-2 ca replay, ca să reducem riscul de catastrophic forgetting. Rezultatul final este un director runs/cpt-scitldr-replay-final/, un fișier runs/metrics.json cu perplexity înainte/după și un script de inferență care folosește modelul salvat local.

Ideea este simplă: nu antrenezi modelul doar pe domeniul nou, pentru că poate deveni mai slab pe limbaj general. Amesteci datele specializate cu o cantitate controlată de date generale, așa cum ai face într-un proiect mai mare de Data Mixture Engineering. Aici folosim o variantă mică, reproductibilă, potrivită pentru un laborator local.

Ce construim

Construim un pipeline cu patru componente: descărcarea dataseturilor, pregătirea amestecului 3:1 între domeniu și replay general, continued pretraining causal language modeling și evaluare înainte/după pe două seturi: SciTLDR validation și WikiText-2 validation.

Domeniul specializat este text științific din SciTLDR, configurația AIC. Setul general de protecție este WikiText-2 raw. Modelul de bază este distilgpt2, ales pentru că este suficient de mic pentru un tutorial local și folosește interfața standard Hugging Face Transformers. Pentru producție, aceeași logică se aplică pe modele mai mari, dar costurile și validarea cresc substanțial.

Mediul exact

Mediu testat pentru pașii de mai jos:

• Ubuntu Server 24.04 LTS
• Python 3.11
• PyTorch 2.4.1
• Transformers 4.45.2
• Datasets 3.0.1
• Accelerate 1.0.1
• GPU NVIDIA cu cel puțin 8 GB VRAM sau CPU pentru rulare lentă
• minimum 16 GB RAM
• aproximativ 30-90 de minute, în funcție de GPU și conexiune

Pe CPU funcționează, dar antrenarea va fi semnificativ mai lentă. Valorile exacte de perplexity pot varia ușor între GPU-uri, drivere și biblioteci, dar fișierele generate, structura rezultatului și comenzile de verificare rămân aceleași.

Ce îți trebuie

Creează un director curat pentru experiment:

mkdir -p cpt-scitldr-replay
cd cpt-scitldr-replay
python3.11 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip

Instalează pachetele folosite în tutorial:

pip install torch==2.4.1 --index-url https://download.pytorch.org/whl/cu121
pip install transformers==4.45.2 datasets==3.0.1 accelerate==1.0.1

Ce trebuie să vezi:

Successfully installed ...

Dacă folosești CPU și instalarea cu indexul CUDA nu este potrivită pentru sistemul tău, folosește pagina oficială PyTorch pentru comanda exactă a platformei. Restul tutorialului rămâne identic.

Pasul 1: verificăm versiunile

Creează fișierul check_env.py:

import torch
import transformers
import datasets
import accelerate

print("torch", torch.__version__)
print("transformers", transformers.__version__)
print("datasets", datasets.__version__)
print("accelerate", accelerate.__version__)
print("cuda_available", torch.cuda.is_available())
if torch.cuda.is_available():
    print("gpu", torch.cuda.get_device_name(0))

Rulează verificarea:

python check_env.py

Ce trebuie să vezi, cu numele GPU-ului diferit pe sistemul tău:

torch 2.4.1+cu121
transformers 4.45.2
datasets 3.0.1
accelerate 1.0.1
cuda_available True
gpu NVIDIA ...

Pasul 2: scriem scriptul complet de continued pretraining

Creează fișierul train_cpt_replay.py. Scriptul descarcă SciTLDR și WikiText-2, construiește amestecul de antrenare, evaluează modelul inițial, rulează continued pretraining și salvează metricele finale.

import json
import math
import os
import random
from pathlib import Path

import torch
from datasets import Dataset, load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
    set_seed,
)

SEED = 47
BASE_MODEL = "distilgpt2"
OUTPUT_DIR = Path("runs/cpt-scitldr-replay")
FINAL_DIR = Path("runs/cpt-scitldr-replay-final")
METRICS_PATH = Path("runs/metrics.json")
BLOCK_SIZE = 256
DOMAIN_TRAIN_ROWS = 384
GENERAL_TRAIN_ROWS = 128
EVAL_ROWS = 128
MAX_STEPS = 120

os.environ["TOKENIZERS_PARALLELISM"] = "false"
set_seed(SEED)
random.seed(SEED)
torch.manual_seed(SEED)

if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

def non_empty_wikitext(example):
    return example["text"].strip() != ""

def scitldr_to_text(example):
    source = " ".join(example["source"])
    target = " ".join(example["target"])
    return {
        "text": "Scientific paper context:n"
        + source
        + "nShort technical summary:n"
        + target
    }

def load_text_datasets():
    domain = load_dataset("allenai/scitldr", "AIC")
    general = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1")

    domain_train = domain["train"].select(range(DOMAIN_TRAIN_ROWS)).map(
        scitldr_to_text,
        remove_columns=domain["train"].column_names,
    )

    domain_eval = domain["validation"].select(range(EVAL_ROWS)).map(
        scitldr_to_text,
        remove_columns=domain["validation"].column_names,
    )

    general_train_source = general["train"].filter(non_empty_wikitext)
    general_eval_source = general["validation"].filter(non_empty_wikitext)

    general_train = general_train_source.select(range(GENERAL_TRAIN_ROWS))
    general_eval = general_eval_source.select(range(EVAL_ROWS))

    train_texts = []
    for i in range(GENERAL_TRAIN_ROWS):
        start = i * 3
        train_texts.extend(domain_train[start:start + 3]["text"])
        train_texts.append(general_train[i]["text"])

    train_dataset = Dataset.from_dict({"text": train_texts})
    return train_dataset, domain_eval, general_eval

def tokenize_dataset(dataset, tokenizer):
    def tokenize_batch(batch):
        return tokenizer(
            batch["text"],
            truncation=True,
            max_length=BLOCK_SIZE,
            padding=False,
        )

    tokenized = dataset.map(
        tokenize_batch,
        batched=True,
        remove_columns=dataset.column_names,
    )

    return tokenized.filter(lambda example: len(example["input_ids"]) > 8)

def perplexity_from_loss(loss):
    if loss > 20:
        return float("inf")
    return math.exp(loss)

def evaluate_perplexity(trainer, eval_dataset, name):
    result = trainer.evaluate(eval_dataset=eval_dataset, metric_key_prefix=name)
    loss = result[f"{name}_loss"]
    return {
        "loss": loss,
        "perplexity": perplexity_from_loss(loss),
    }

def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    FINAL_DIR.mkdir(parents=True, exist_ok=True)
    METRICS_PATH.parent.mkdir(parents=True, exist_ok=True)

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
    model.config.pad_token_id = tokenizer.eos_token_id

    train_dataset, domain_eval_text, general_eval_text = load_text_datasets()

    train_tokens = tokenize_dataset(train_dataset, tokenizer)
    domain_eval_tokens = tokenize_dataset(domain_eval_text, tokenizer)
    general_eval_tokens = tokenize_dataset(general_eval_text, tokenizer)

    collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,
    )

    args = TrainingArguments(
        output_dir=str(OUTPUT_DIR),
        overwrite_output_dir=True,
        max_steps=MAX_STEPS,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        learning_rate=5e-5,
        warmup_steps=10,
        weight_decay=0.01,
        logging_steps=20,
        save_steps=120,
        save_total_limit=1,
        report_to=[],
        seed=SEED,
        data_seed=SEED,
        fp16=torch.cuda.is_available(),
        dataloader_num_workers=0,
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_tokens,
        data_collator=collator,
        tokenizer=tokenizer,
    )

    before_domain = evaluate_perplexity(trainer, domain_eval_tokens, "before_domain")
    before_general = evaluate_perplexity(trainer, general_eval_tokens, "before_general")

    trainer.train()

    after_domain = evaluate_perplexity(trainer, domain_eval_tokens, "after_domain")
    after_general = evaluate_perplexity(trainer, general_eval_tokens, "after_general")

    trainer.save_model(str(FINAL_DIR))
    tokenizer.save_pretrained(str(FINAL_DIR))

    metrics = {
        "base_model": BASE_MODEL,
        "seed": SEED,
        "block_size": BLOCK_SIZE,
        "domain_dataset": "allenai/scitldr/AIC",
        "general_replay_dataset": "Salesforce/wikitext/wikitext-2-raw-v1",
        "mixture": "3 domain examples : 1 general replay example",
        "max_steps": MAX_STEPS,
        "before_domain": before_domain,
        "before_general": before_general,
        "after_domain": after_domain,
        "after_general": after_general,
    }

    METRICS_PATH.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    print(json.dumps(metrics, indent=2))

if __name__ == "__main__":
    main()

Ce trebuie observat în cod: nu facem doar fine-tuning pe SciTLDR. La fiecare trei exemple de domeniu adăugăm un exemplu WikiText-2. Acesta este mecanismul de replay care limitează degradarea pe text general. Dacă vrei să optimizezi ulterior eficiența pe secvențe mai lungi, merită citit și ghidul despre Sequence Packing avansat, dar aici păstrăm pipeline-ul simplu și verificabil.

Pasul 3: rulează antrenarea

Pornește continued pretraining:

python train_cpt_replay.py

Ce trebuie să vezi în timpul rulării:

{'loss': ..., 'grad_norm': ..., 'learning_rate': ..., 'epoch': ...}

La final trebuie să apară un JSON cu aceste chei:

{
  "base_model": "distilgpt2",
  "seed": 47,
  "block_size": 256,
  "domain_dataset": "allenai/scitldr/AIC",
  "general_replay_dataset": "Salesforce/wikitext/wikitext-2-raw-v1",
  "mixture": "3 domain examples : 1 general replay example",
  "max_steps": 120,
  "before_domain": {
    "loss": ...,
    "perplexity": ...
  },
  "before_general": {
    "loss": ...,
    "perplexity": ...
  },
  "after_domain": {
    "loss": ...,
    "perplexity": ...
  },
  "after_general": {
    "loss": ...,
    "perplexity": ...
  }
}

Interpretarea corectă este defensivă, nu triumfalistă: vrem ca perplexity pe domeniu să se îmbunătățească sau să nu se degradeze, iar perplexity pe WikiText-2 să nu sară brusc. Dacă generalul se degradează puternic, amestecul are prea puțin replay sau rata de învățare este prea agresivă.

Pasul 4: verifică fișierele produse

Rulează:

find runs -maxdepth 2 -type f | sort

Ce trebuie să vezi:

runs/cpt-scitldr-replay-final/config.json
runs/cpt-scitldr-replay-final/generation_config.json
runs/cpt-scitldr-replay-final/merges.txt
runs/cpt-scitldr-replay-final/model.safetensors
runs/cpt-scitldr-replay-final/special_tokens_map.json
runs/cpt-scitldr-replay-final/tokenizer.json
runs/cpt-scitldr-replay-final/tokenizer_config.json
runs/cpt-scitldr-replay-final/vocab.json
runs/metrics.json

Verifică metricele:

python -m json.tool runs/metrics.json

Ce trebuie să vezi: același JSON formatat, cu patru grupuri de metrice: before_domain, before_general, after_domain, after_general.

Pasul 5: inferență cu modelul salvat

Creează infer.py:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_DIR = "runs/cpt-scitldr-replay-final"

tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModelForCausalLM.from_pretrained(MODEL_DIR)
model.eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

prompt = (
    "Scientific paper context:n"
    "A neural language model is adapted to a narrow corpus while retaining "
    "general language ability through replay data.n"
    "Short technical summary:n"
)

inputs = tokenizer(prompt, return_tensors="pt").to(device)

with torch.no_grad():
    output = model.generate(
        **inputs,
        max_new_tokens=80,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
    )

print(tokenizer.decode(output[0], skip_special_tokens=True))

Rulează:

python infer.py

Ce trebuie să vezi:

Scientific paper context:
A neural language model is adapted ...
Short technical summary:
...

Textul generat de distilgpt2 nu trebuie tratat ca informație științifică verificată. În acest tutorial, inferența confirmă că modelul local se încarcă și generează; calitatea reală cere evaluare separată, inclusiv verificări de contaminare ca în Benchmark Contamination Detection.

Verificarea finală

Ai un rezultat valid dacă toate comenzile de mai jos trec:

test -f runs/metrics.json
test -f runs/cpt-scitldr-replay-final/model.safetensors
python -m json.tool runs/metrics.json >/dev/null
python infer.py >/tmp/cpt_infer.txt
test -s /tmp/cpt_infer.txt
echo "VERIFICARE_OK"

Output așteptat:

VERIFICARE_OK

Deschide apoi runs/metrics.json și compară cele două axe: domeniu și general. Continued pretraining-ul este util doar dacă adaptarea pe domeniu nu vine cu o pierdere necontrolată pe general. În proiecte mari, etapa următoare ar fi să testezi mai multe amestecuri de date și să păstrezi modelul care are cel mai bun compromis, nu doar cel mai mic loss pe domeniul nou.

Dacă vrei să servești modelul mai târziu, nu să-l antrenezi, separă clar etapa de training de etapa de runtime. Pentru inferență scalată, principiile din articolul despre deploy cu vLLM, continuous batching și KV cache sunt relevante după ce ai stabilit că modelul chiar merită publicat intern.

Probleme frecvente

Eroare:

ModuleNotFoundError: No module named 'transformers'

Cauză: mediul virtual nu este activat sau instalarea nu a rulat în același Python. Rezolvare:

source .venv/bin/activate
pip install transformers==4.45.2 datasets==3.0.1 accelerate==1.0.1

Eroare:

CUDA out of memory

Cauză: GPU-ul nu are memorie suficientă pentru batch, activări și optimizer. Rezolvare pentru acest tutorial: rulează cu batch-ul deja setat la 1 și oprește alte procese GPU. Verifică:

nvidia-smi

Dacă problema persistă, rulează pe CPU prin ascunderea GPU-ului:

CUDA_VISIBLE_DEVICES="" python train_cpt_replay.py

Eroare:

ValueError: Asking to pad but the tokenizer does not have a padding token

Cauză: GPT-2 nu are token de padding implicit. În script, linia tokenizer.pad_token = tokenizer.eos_token rezolvă problema. Dacă ai modificat scriptul, pune linia înapoi și rulează din nou.

Eroare:

FileNotFoundError: runs/cpt-scitldr-replay-final

Cauză: inferența a fost pornită înainte ca antrenarea să salveze modelul final. Rezolvare:

python train_cpt_replay.py
python infer.py

Revenire și curățenie

Pentru a șterge complet experimentul local:

cd ..
rm -rf cpt-scitldr-replay

Dacă vrei să păstrezi doar rezultatele, arhivează directorul modelului și metricele:

tar -czf cpt-scitldr-replay-final.tar.gz cpt-scitldr-replay/runs

Nu publica modelul pe Hugging Face Hub sau într-un serviciu intern fără o verificare separată a licențelor dataseturilor, a calității și a riscului de memorare. Dataset card-ul SciTLDR listează licența ca necunoscută, deci experimentul de aici trebuie tratat ca laborator local, nu ca bază automată pentru distribuție comercială.

Mai departe

Primul pas util este să rulezi același script cu două rapoarte diferite, de exemplu 1:1 și 7:1 între domeniu și general, apoi să compari runs/metrics.json. Al doilea este să adaugi un set de evaluare propriu, curat, pe domeniul real al organizației. Al treilea este să separi adaptarea prin continued pretraining de specializări mai mici, unde LoRA adapters în producție pot fi mai potrivite decât modificarea completă a modelului.

Întrebări frecvente

De ce folosim WikiText-2 ca replay general?

Pentru că tutorialul are nevoie de un set general mic, public și ușor de descărcat. El nu garantează păstrarea tuturor capabilităților generale, dar oferă un semnal clar pentru degradarea pe limbaj general.

De ce nu antrenăm doar pe SciTLDR?

Antrenarea doar pe domeniul nou crește riscul ca modelul să se specializeze prea mult și să piardă performanță pe text general. Amestecul 3:1 introduce replay general în fiecare ciclu de training.

Este modelul rezultat potrivit pentru producție?

Nu direct. Tutorialul produce un laborator reproductibil cu distilgpt2, metrice și inferență locală. Pentru producție trebuie verificate licențele, calitatea, memorarea și comportamentul pe evaluări mai largi.

Ce metrică urmărim aici?

Urmărim loss și perplexity înainte/după pe două seturi: SciTLDR validation pentru domeniu și WikiText-2 validation pentru general. Un rezultat bun nu înseamnă doar domeniu mai bun, ci și degradare controlată pe general.

Surse

• Hugging Face Transformers — Causal language modeling
• Hugging Face Transformers — Trainer
• Hugging Face Transformers — Data collator
• Hugging Face Datasets — Loading methods
• Salesforce/wikitext dataset card
• allenai/scitldr dataset card