Fine-Tuning Baichuan Models on Proprietary Domain Datasets with LoRA and Hugging Face PEFT

Fine-tuning open-weights foundation models on proprietary domain data—such as internal financial records, clinical guidelines, compliance handbooks, or enterprise customer support transcripts—requires balancing adaptation quality, data privacy, and compute budgets. Baichuan models (including Baichuan-7B/13B and Baichuan2-7B/13B Base and Chat) have emerged as preferred architectures for bilingual Chinese-English enterprise applications due to their custom 125,696-token vocabulary, specialized positional encodings (RoPE on 7B, ALiBi on 13B), and stabilized output projection heads (NormHead).

Using Low-Rank Adaptation (LoRA) through Hugging Face's Parameter-Efficient Fine-Tuning (peft) library allows NLP engineers to adapt a 7B or 13B Baichuan model on modest hardware—often a single 24GB GPU using 4-bit QLoRA, or a single 80GB A100/H100 in bfloat16—while updating less than 1% of the model parameters.

1. Architectural Foundations: Why Baichuan Differs

Before configuring fine-tuning pipelines, NLP engineers must account for Baichuan’s architectural characteristics:

1. Vocabulary Density & Tokenization: Baichuan models feature an expansive 125,696-token vocabulary based on byte-pair encoding (BPE). This reduces sequence lengths for Chinese domain text by 15% to 20% compared to standard LLaMA-derived tokenizers, saving GPU memory during sequence packing.
2. Positional Encodings:
   • 7B Models: Employ Rotary Position Embeddings (RoPE), standard for 4,096-token context windows.
   • 13B Models: Implement Attention with Linear Biases (ALiBi). ALiBi alters attention score calculations directly rather than injecting positional embeddings into queries and keys, offering superior zero-shot context length extrapolation for long documents.
3. Normalized LM Head (NormHead): Baichuan 2 introduces weight normalization on the linear output head (lm_head). This reduces the influence of rare token embeddings on logit calculations and stabilizes pre-training dynamics. During parameter-efficient tuning, this layer should remain frozen unless expanding vocabulary.
4. Mandatory Custom Code Execution: Baichuan architectures require custom modeling files hosted on the model repository. Any loading call in Hugging Face Transformers must specify trust_remote_code=True.

2. Dataset Preparation and Special Token Formatting

Proprietary enterprise datasets typically exist as structured QA pairs, customer-agent dialogues, or document-summary extractions. For instruction fine-tuning, Baichuan relies on dedicated control tokens for chat turns:

• <reserved_106>: Signifies the start of the user/human prompt.
• <reserved_107>: Signifies the start of the assistant/model response.

Data Format (ShareGPT / Conversations Schema)

Organize domain data into JSON or JSON Lines (.jsonl) files where conversations are explicitly delineated:

[
  {
    "conversations": [
      {
        "from": "human",
        "value": "What are the compliance disclosure obligations for cross-border asset transfers under our proprietary policy?"
      },
      {
        "from": "assistant",
        "value": "Under Policy Section 4.2, cross-border asset transfers exceeding $50,000 require dual-officer sign-off and an automated AML audit clearance within 24 hours."
      }
    ]
  }
]

Loss Masking on Prompts

A critical engineering practice during supervised fine-tuning (SFT) is prompt masking. To prevent the model from expending capacity memorizing prompts, assign a target label of -100 to all prompt tokens (including <reserved_106>). Cross-entropy loss is computed strictly on tokens generated by the assistant following <reserved_107> up to the end-of-sequence token (</s> or <reserved_107>).

3. LoRA Configuration Strategy for Baichuan

LoRA decomposes weight updates $\Delta W$ for a weight matrix $W_0 \in \mathbb{R}^{d \times k}$ into two low-rank matrices $B \in \mathbb{R}^{d \times r}$ and $A \in \mathbb{R}^{r \times k}$, where $r \ll \min(d, k)$:

$$W = W_0 + \frac{\alpha}{r} (B \cdot A)$$

Selecting target_modules

While early LoRA recipes targeted only the query (q_proj) and value (v_proj) projections, empirical results show that adapting all linear layers across the self-attention and MLP blocks is necessary for complex domain knowledge injection:

• Attention Layers: q_proj, k_proj, v_proj, o_proj
• Feedforward (MLP) Layers: gate_proj, up_proj, down_proj

Modern versions of peft accept target_modules="all-linear", which automatically resolves to all linear projections across Baichuan’s transformer decoder blocks.

Hyperparameter Recommendations

4. End-to-End Implementation Script

Below is a complete, production-ready PyTorch script using Hugging Face transformers and peft with 4-bit quantization (QLoRA) for resource-constrained training.

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. Load Tokenizer
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. Configure 4-bit Quantization (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. Load Model
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
)

# 4. Prepare Model for k-bit Training
model = prepare_model_for_kbit_training(model)

# 5. Define LoRA Configuration
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. Data Formatting Function with Loss Masking
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. Training Arguments
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. Initialize Trainer and Train
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. Save LoRA Adapter
model.save_pretrained("./baichuan2-domain-lora-adapter")
tokenizer.save_pretrained("./baichuan2-domain-lora-adapter")

5. Weight Merging and Production Serving

In a microservices architecture, serving raw LoRA adapters through Hugging Face introduces inference latency and limits serving framework optimizations. To deploy with high-throughput engines like vLLM, Triton Inference Server, or LMDeploy, engineers merge the adapter into the base weights.

Offline Weight Fusion Script

Important: Do not merge adapters directly into 4-bit quantized base models. Load the unquantized base model in torch.bfloat16 or torch.float16 into host memory or an unquantized GPU instance:

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)

base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    torch_dtype=torch.bfloat16,
    device_map="cpu",  # CPU or GPU with sufficient VRAM
    trust_remote_code=True
)

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

merged_model.save_pretrained(export_path, safe_serialization=True)
tokenizer.save_pretrained(export_path)

The resulting standalone model in export_path contains standard Hugging Face weights that can be immediately ingested by vLLM via vllm serve ./baichuan2-7b-domain-merged --trust-remote-code.

6. Engineering Checkpoints and Trade-offs

1. Base vs. Chat as the Starting Point:
   • Start with Baichuan2-Base when building strict structured-output agents (e.g., deterministic JSON extraction, SQL generation, or translation) where pre-existing conversational conditioning may cause verbose or conversational divergence.
   • Start with Baichuan2-Chat when adapting conversational enterprise assistants, role-playing customer agents, or triage bots that require out-of-the-box conversational etiquette and safety alignment.
2. 7B (RoPE) vs. 13B (ALiBi):
   • For document analysis exceeding 4,000 tokens, the 13B model provides superior extrapolation due to ALiBi, reducing attention drift over multi-page contracts and medical histories.
   • The 7B model offers higher throughput and lower memory footprint, fitting comfortably onto entry-level enterprise GPUs (e.g., NVIDIA L4 or RTX 4090).
3. Evaluating Domain Retention:
   • Track both task-specific validation metrics (e.g., domain accuracy, BLEU/ROUGE for summarization) and general language perplexity on a held-out set to monitor catastrophic forgetting of foundational reasoning.

References

• 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