Comment affiner les modèles Baichuan sur des données propriétaires avec LoRA et Hugging Face PEFT

Dans les secteurs régulés tels que la finance, la santé, le juridique et le support d'entreprise, l'utilisation directe d'API de modèles propriétaires hébergés sur le cloud public soulève des contraintes majeures de souveraineté des données, de conformité et de confidentialité. L'adaptation de modèles ouverts (open-weights) sur une infrastructure interne dédiée est devenue la norme pour les équipes d'ingénierie TAL (NLP). La gamme de modèles développée par Baichuan Intelligent (notamment Baichuan-7B/13B et Baichuan2-7B/13B en versions Base et Chat) s'impose comme une référence pour les contextes bilingues et techniques grâce à son pré-entraînement sur 2,6 billions de tokens, son vocabulaire étendu de 125 696 tokens et ses innovations architecturales (RoPE sur le 7B, ALiBi sur le 13B, normalisation de tête NormHead).

En exploitant la méthode Low-Rank Adaptation (LoRA) via la bibliothèque Parameter-Efficient Fine-Tuning (peft) de Hugging Face, les ingénieurs peuvent adapter un modèle Baichuan en ne modifiant que moins de 1 % de ses paramètres. Cette approche permet de réaliser l'apprentissage sur un seul GPU de 24 Go grâce à QLoRA en 4 bits, ou sur une carte A100/H100 de 80 Go en précision bfloat16 native.

1. Spécificités architecturales des modèles Baichuan

Avant de configurer le pipeline d'entraînement, l'ingénieur TAL doit intégrer les particularités qui distinguent Baichuan des architectures LLaMA classiques :

1. Vocabulaire BPE de 125 696 tokens : Ce dictionnaire volumineux compresse très efficacement les termes techniques et les caractères non latins. Par rapport aux tokeniseurs basés sur LLaMA, la longueur des séquences est réduite de 15 % à 20 %, ce qui diminue la pression sur la mémoire VRAM lors de la constitution des lots.
2. Mécanismes d'encodage positionnel :
   • Modèles 7B : Utilisent les Rotary Position Embeddings (RoPE) avec une fenêtre de contexte standard de 4 096 tokens.
   • Modèles 13B : Implémentent l'Attention with Linear Biases (ALiBi). Au lieu d'additionner des vecteurs positionnels aux projections Query et Key, ALiBi applique une pénalité linéaire dépendante de la distance sur les scores d'attention. Cette conception offre une capacité d'extrapolation supérieure sur les documents longs (contrats, dossiers médicaux).
3. Normalisation de la tête de sortie (NormHead) : Baichuan 2 normalise les poids de la couche linéaire de projection finale (lm_head). Cela atténue l'influence de la norme des tokens rares sur le calcul des logits et stabilise le pré-entraînement. Dans le cadre d'un ajustement LoRA, il est fortement recommandé de ne pas inclure lm_head dans les modules cibles, sous peine de perturber la distribution des probabilités calibrée en amont.
4. Exécution obligatoire du code personnalisé : Les modèles Baichuan s'appuient sur des scripts personnalisés intégrés à leur dépôt. L'argument trust_remote_code=True est indispensable lors du chargement du modèle et du tokeniseur.

2. Préparation des données d'entreprise et tokens de contrôle

Les jeux de données propriétaires sont généralement constitués de paires question-réponse, de journaux de tickets ou de résumés de procédures internes. Pour le fine-tuning supervisé (SFT), Baichuan s'appuie sur des tokens de contrôle dédiés à la gestion des tours de parole :

• <reserved_106> (Token ID : 195) : Indique le début de la consigne ou question de l'utilisateur (Human).
• <reserved_107> (Token ID : 196) : Indique le début de la réponse de l'assistant (Assistant).

Format des données (Schéma ShareGPT / Conversations)

Structurez vos données sous forme de fichiers JSON ou JSON Lines (.jsonl) :

[
  {
    "conversations": [
      {
        "from": "human",
        "value": "Quelles sont les obligations de conformité pour les virements transfrontaliers de fonds supérieurs à 50 000 dollars selon notre politique interne ?"
      },
      {
        "from": "assistant",
        "value": "Conformément à la section 4.2 du manuel de conformité financière, les virements internationaux excédant 50 000 USD requièrent une double validation managériale ainsi qu'une vérification AML automatisée dans le système sous 24 heures ouvrées."
      }
    ]
  }
]

Masquage de la perte sur le prompt (Prompt Loss Masking)

Pour éviter que le modèle ne gaspille sa capacité à mémoriser les formulations des questions, appliquez un masquage strict. Tous les tokens du prompt (de <reserved_106> à <reserved_107>) reçoivent le label -100. La perte d'entropie croisée n'est calculée que sur les tokens générés par l'assistant et le token de fin de séquence (</s> ou <reserved_107>).

3. Stratégie de configuration LoRA et hyperparamètres

LoRA décompose la mise à jour des poids $\Delta W$ en un produit de deux matrices de rang faible : $W = W_0 + \frac{\alpha}{r}(B \cdot A)$.

Choix des modules cibles (target_modules)

Alors que les premières implémentations se limitaient à q_proj et v_proj, l'injection efficace de connaissances dans des domaines techniques exige d'adapter l'ensemble des couches linéaires :

• Attention : q_proj, k_proj, v_proj, o_proj
• Réseau feed-forward (MLP) : gate_proj, up_proj, down_proj

Avec les versions récentes de la bibliothèque peft, définir target_modules="all-linear" permet de cibler automatiquement toutes les projections linéaires du décodeur Baichuan.

Recommandations d'hyperparamètres

4. Script d'implémentation complet (Transformers et PEFT)

Le script suivant présente le flux d'apprentissage complet en 4 bits (QLoRA) sur un seul 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. Chargement du tokeniseur
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. Configuration de la quantification 4 bits (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. Chargement du modèle de base
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
)

# 4. Préparation pour l'entraînement quantifié
model = prepare_model_for_kbit_training(model)

# 5. Configuration de l'adaptateur LoRA
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. Prétraitement avec masquage de la perte sur le prompt
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. Paramètres d'entraînement
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. Initialisation du Trainer et exécution
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. Sauvegarde de l'adaptateur LoRA
model.save_pretrained("./baichuan2-domain-lora-adapter")
tokenizer.save_pretrained("./baichuan2-domain-lora-adapter")

5. Fusion des poids et déploiement haute performance

En production, charger dynamiquement des adaptateurs LoRA engendre une surcharge de calcul et empêche les moteurs d'inférence modernes tels que vLLM, Triton ou LMDeploy d'appliquer la fusion d'opérateurs. La démarche recommandée consiste à fusionner hors ligne l'adaptateur avec les poids de base.

Script de fusion des poids (Merge and Unload)

Attention : Ne jamais exécuter la fusion directement sur un modèle quantifié en 4 bits. Chargez le modèle d'origine non quantifié en torch.bfloat16 ou torch.float16 en mémoire centrale ou sur un GPU non saturé :

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)

# Charger le modèle de base non quantifié
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    torch_dtype=torch.bfloat16,
    device_map="cpu",  # Utiliser "cuda" si la VRAM est suffisante
    trust_remote_code=True
)

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

# Sauvegarde au format standard Hugging Face
merged_model.save_pretrained(export_path, safe_serialization=True)
tokenizer.save_pretrained(export_path)

Le répertoire généré peut ensuite être déployé directement sous vLLM : vllm serve ./baichuan2-7b-domain-merged --trust-remote-code

6. Critères de choix et bonnes pratiques

1. Modèle de départ (Base vs Chat) :
   • Choisissez Baichuan2-Base pour des tâches à format strict (extraction JSON, requêtes SQL, classification), évitant ainsi les tournures conversationnelles superflues.
   • Choisissez Baichuan2-Chat pour les agents conversationnels ou les robots de support client qui requièrent un comportement poli et une modération préalablement intégrée.
2. Taille du modèle (7B vs 13B) :
   • Pour les documents techniques excédant 4 000 tokens, le 13B avec ALiBi offre une stabilité supérieure dans la conservation du contexte long.
   • Le 7B assure une latence plus faible et un encombrement mémoire réduit, idéal pour des applications d'entreprise à fort trafic.
3. Prévention de l'oubli catastrophique :
   • Intégrez 5 % à 10 % de données d'instructions généralistes dans le jeu d'entraînement pour préserver les capacités de raisonnement global du modèle.

Références

• 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