Fine-Tuning von Baichuan-Modellen auf proprietären Domänendaten mit LoRA und Hugging Face PEFT

In regulierten Branchen wie dem Finanzsektor, dem Gesundheitswesen, der Rechtsberatung und dem unternehmensweiten Wissensmanagement stoßen öffentliche Cloud-APIs schnell an Grenzen bezüglich Datensouveränität, Vertraulichkeit und fachspezifischer Terminologie. Das lokale Fine-Tuning von Open-Weights-Modellen auf eigener Infrastruktur hat sich daher als Standard für NLP-Engineering-Teams etabliert. Die Baichuan-Modellfamilie von Baichuan Intelligent (darunter Baichuan-7B/13B und Baichuan2-7B/13B in den Varianten Base und Chat) hat sich dank ihres Pre-Trainings auf 2,6 Billionen Tokens, ihres 125.696 Tokens umfassenden Vokabulars sowie optimierter Architekturmerkmale (RoPE auf 7B, ALiBi auf 13B, normalisierter Ausgabekopf NormHead) zu einer gefragten Basis für spezialisierte Fachanwendungen entwickelt.

Durch den Einsatz von Low-Rank Adaptation (LoRA) über die Hugging Face-Bibliothek peft (Parameter-Efficient Fine-Tuning) können Entwickler ein Baichuan-Modell anpassen, indem sie weniger als 1 % der Modellparameter trainieren. Dies ermöglicht die Ausführung auf einer einzelnen 24-GB-GPU mittels 4-Bit-QLoRA oder auf einer einzelnen A100/H100 mit 80 GB VRAM in nativer bfloat16-Präzision.

1. Architektonische Besonderheiten der Baichuan-Modelle

Vor der Konfiguration der Trainingspipeline müssen NLP-Ingenieure mehrere spezifische Konstruktionsmerkmale berücksichtigen, die Baichuan von Standard-LLaMA-Architekturen unterscheiden:

1. Erweitertes BPE-Vokabular (125.696 Tokens): Das BPE-Vokabular komprimiert technische Fachtermini und mehrsprachige Texte mit hoher Effizienz. Verglichen mit LLaMA-Tokenizern verkürzt sich die Sequenzlänge um 15 % bis 20 %, was den VRAM-Bedarf bei der Sequenzbündelung (Sequence Packing) merklich senkt.
2. Positionskodierung (7B vs. 13B):
   • 7B-Modelle: Verwenden Rotary Position Embeddings (RoPE) mit einem Standard-Kontextfenster von 4.096 Tokens.
   • 13B-Modelle: Nutzen Attention with Linear Biases (ALiBi). ALiBi addiert keine Positionsvektoren zu Query- und Key-Projektionen, sondern wendet distanzbasierte Strafsteigungen direkt auf die Attention-Score-Matrix an. Dies ermöglicht eine überlegene Kontext-Extrapolation bei langen Dokumenten wie Verträgen oder Krankenakten.
3. Normalisierter Projektionskopf (NormHead): Baichuan 2 führt eine Gewichtsnormierung auf der linearen Ausgabeschicht (lm_head) ein. Dies verringert den störenden Einfluss seltener Token-Normen auf die Logit-Berechnung und stabilisiert die Pre-Training-Dynamik. Beim LoRA-Fine-Tuning sollte die Schicht lm_head eingefroren bleiben, um die etablierte Ausgabeverteilung nicht zu deregulieren.
4. Erforderliche Ausführung von benutzerdefiniertem Code: Da Baichuan eigene Modellierungs- und Tokenizer-Definitionen verwendet, muss beim Laden in Hugging Face Transformers stets trust_remote_code=True angegeben werden.

2. Vorbereitung proprietärer Domänendaten und Steuerungs-Tokens

Proprietäre Unternehmensdaten liegen typischerweise als strukturierte Dokumentationsauszüge, regulatorische Richtlinien oder Dialogprotokolle vor. Für das überwachte Fine-Tuning (Supervised Fine-Tuning, SFT) setzt Baichuan dedizierte Steuerungs-Tokens zur Rollentrennung ein:

• <reserved_106> (Token-ID: 195): Markiert den Beginn der Benutzereingabe (Human).
• <reserved_107> (Token-ID: 196): Markiert den Beginn der Modellantwort (Assistant).

Datenformat (ShareGPT-Schema)

Die Daten sollten als JSON- oder JSON Lines-Dateien (.jsonl) mit klar abgegrenzten Dialogabschnitten strukturiert sein:

[
  {
    "conversations": [
      {
        "from": "human",
        "value": "Welche Freigabeprozesse sind gemäß unserer internen Richtlinie für grenzüberschreitende Überweisungen über 50.000 USD erforderlich?"
      },
      {
        "from": "assistant",
        "value": "Gemäß Abschnitt 4.2 der internen Finanz-Compliance-Richtlinie erfordern Auslandsüberweisungen von mehr als 50.000 USD eine doppelte Zeichnungsberechtigung sowie eine automatisierte AML-Prüfung im System innerhalb von 24 Arbeitsstunden."
      }
    ]
  }
]

Maskierung des Prompt-Verlusts (Prompt Loss Masking)

Um zu verhindern, dass das Modell Trainingskapazität auf das Auswendiglernen der Eingabefragen verschwendet, müssen die Prompt-Tokens maskiert werden. Alle Tokens von <reserved_106> bis einschließlich <reserved_107> erhalten das Label -100. Der Kreuzentropieverlust (Cross-Entropy Loss) wird ausschließlich über die vom Assistenten erzeugten Tokens und das Sequenzende-Token (</s> oder <reserved_107>) berechnet.

3. LoRA-Konfigurationsstrategie und Hyperparameter

LoRA zerlegt die Gewichtsaktualisierungsmatrix $\Delta W$ in das Produkt zweier Matrizen mit niedrigem Rang: $W = W_0 + \frac{\alpha}{r}(B \cdot A)$.

Auswahl der Zielmodule (target_modules)

Während frühe LoRA-Ansätze lediglich die Schichten q_proj und v_proj modifizierten, verlangt die Anpassung an komplexe Fachdomänen die Einbindung aller linearen Schichten:

• Attention-Schichten: q_proj, k_proj, v_proj, o_proj
• MLP-Schichten: gate_proj, up_proj, down_proj

In aktuellen Versionen der Bibliothek peft kann dafür einfach target_modules="all-linear" deklariert werden.

Empfohlene Hyperparameter

4. Vollständiges Implementierungsskript (Transformers & PEFT)

Das folgende Python-Skript demonstriert den vollständigen Trainingsablauf mit 4-Bit-Quantisierung (QLoRA) auf einer einzelnen GPU.

import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq
)
from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    TaskType
)

model_id = "baichuan-inc/Baichuan2-7B-Base"

# 1. Tokenizer initialisieren
tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    trust_remote_code=True,
    use_fast=False
)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# 2. 4-Bit-Quantisierung konfigurieren (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
)

# 3. Basismodell laden
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
)

# 4. Modell auf k-Bit-Training vorbereiten
model = prepare_model_for_kbit_training(model)

# 5. LoRA-Konfiguration definieren
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ]
)

model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

# 6. Datenaufbereitung mit Verlustmaskierung
USER_TOKEN_ID = 195  # <reserved_106>
ASSISTANT_TOKEN_ID = 196  # <reserved_107>

def preprocess_function(example):
    input_ids = []
    labels = []
    
    for message in example["conversations"]:
        if message["from"] == "human":
            prompt_tokens = [USER_TOKEN_ID] + tokenizer.encode(message["value"], add_special_tokens=False)
            input_ids.extend(prompt_tokens)
            labels.extend([-100] * len(prompt_tokens))
        elif message["from"] == "assistant":
            resp_tokens = [ASSISTANT_TOKEN_ID] + tokenizer.encode(message["value"], add_special_tokens=False) + [tokenizer.eos_token_id]
            input_ids.extend(resp_tokens)
            labels.extend(resp_tokens)
            
    return {"input_ids": input_ids[:2048], "labels": labels[:2048]}

dataset = load_dataset("json", data_files="domain_data.jsonl")
tokenized_dataset = dataset["train"].map(preprocess_function, remove_columns=["conversations"])

# 7. Trainingsparameter festlegen
training_args = TrainingArguments(
    output_dir="./baichuan2-lora-output",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.05,
    logging_steps=10,
    num_train_epochs=3,
    gradient_checkpointing=True,
    bf16=True,
    save_strategy="epoch",
    report_to="none"
)

# 8. Trainer instanziieren und Training starten
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8, return_tensors="pt")
)

trainer.train()

# 9. LoRA-Adapter speichern
model.save_pretrained("./baichuan2-domain-lora-adapter")
tokenizer.save_pretrained("./baichuan2-domain-lora-adapter")

5. Gewichtsverschmelzung und Bereitstellung in der Produktion

Im Produktivbetrieb führt das separate dynamische Laden von LoRA-Adaptern zu Latenzeinbußen und verhindert Kernel-Optimierungen in Hochdurchsatz-Engines wie vLLM, Triton oder LMDeploy. Der etablierte Standardweg ist daher die Offline-Verschmelzung der Adaptergewichte mit dem Basismodell.

Skript zur Gewichtsverschmelzung (Merge and Unload)

Wichtiger Hinweis: Führen Sie die Verschmelzung niemals direkt auf einem 4-Bit-quantisierten Modell aus. Laden Sie das unquantisierte Basismodell in torch.bfloat16 oder torch.float16 im Arbeitsspeicher oder auf einer unquantisierten GPU:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

base_model_path = "baichuan-inc/Baichuan2-7B-Base"
adapter_path = "./baichuan2-domain-lora-adapter"
export_path = "./baichuan2-7b-domain-merged"

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

# Unquantisiertes Basismodell laden
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    torch_dtype=torch.bfloat16,
    device_map="cpu",  # Bei ausreichend VRAM kann "cuda" verwendet werden
    trust_remote_code=True
)

peft_model = PeftModel.from_pretrained(base_model, adapter_path)
merged_model = peft_model.merge_and_unload()

# Im Standard-Hugging-Face-Format exportieren
merged_model.save_pretrained(export_path, safe_serialization=True)
tokenizer.save_pretrained(export_path)

Das exportierte Modell kann direkt über vLLM bereitgestellt werden: vllm serve ./baichuan2-7b-domain-merged --trust-remote-code

6. Entscheidungskriterien und Empfehlungen

1. Wahl des Ausgangsmodells (Base vs. Chat):
   • Baichuan2-Base eignet sich optimal für strukturierte Informationsextraktion, SQL-Generierung oder standardisierte Klassifikation, da keine vorangehenden Dialogfloskeln stören.
   • Baichuan2-Chat ist vorzuziehen, wenn interne Dialogassistenten oder Support-Bots entwickelt werden, die bereits grundlegende Gesprächsetikette und Sicherheitsabstimmungen mitbringen sollen.
2. Skalierung (7B vs. 13B):
   • Bei umfangreichen Dokumenten über 4.000 Tokens bietet das 13B-Modell mit ALiBi signifikant stabilere Aufmerksamkeit über weite Distanzen.
   • Das 7B-Modell punktet mit minimalen Latenzen und niedrigem Speicherbedarf, ideal für hochfrequente Abfragen.
3. Schutz vor katastrophalem Vergessen:
   • Mischen Sie dem proprietären Datensatz 5 % bis 10 % allgemeine Instruktionsdaten bei, um die logische Urteilskraft des Basismodells zu bewahren.

Referenzen

• GitHub - baichuan-inc/Baichuan2: A series of large language models developed by Baichuan Intelligent Technology
• GitHub - huggingface/peft: 🤗 PEFT: State-of-the-art Parameter-Efficient Fine-Tuning.
• [2309.10305] Baichuan 2: Open Large-scale Language Models
• [2106.09685] LoRA: Low-Rank Adaptation of Large Language Models