Fine-tuning Zephyr-7B for Finite Volumes: Methodology and Results

Context and Goal

I wanted to test fine-tuning on a domain I know well: finite volume methods in numerical simulation. The goal was to make Zephyr-7B able to explain the concepts and schemes clearly, so it could answer technical questions on the topic.

Data Preparation: The Heart of the Project

The quality of a fine-tune depends on the training data. I built a dataset of a few question/answer pairs covering:

  1. Theoretical concepts:
    • “What is the CFL condition and how do you apply it to an explicit scheme?”
    • “What is the difference between an upwind scheme and a centered scheme for discretizing a convective term?”
  2. Practical cases:
    • “How do you handle a Neumann boundary condition in a finite volume solver?”
    • “What are the stability criteria for an implicit scheme applied to the heat equation?”
  3. Method comparisons:
    • “What are the advantages and disadvantages of explicit vs implicit schemes?”
    • “When should an upwind scheme be preferred over a centered scheme?”

Data format: Each example follows a standard format for SFTTrainer:

{
  "text": "### Instruction: [Question about finite volumes]\n### Response: [Detailed technical answer]"
}

Example:

{
  "text": "### Instruction: Explain the CFL condition for an explicit finite volume scheme.
### Response: The CFL (Courant-Friedrichs-Lewy) condition is a stability criterion for explicit numerical schemes. For a convection problem (∂u/∂t + c∂u/∂x = 0), it is written as Δt ≤ Δx/|c|, where Δt is the time step, Δx the spatial step, and c the convection speed. For diffusion (∂u/∂t = ν∂²u/∂x²), it becomes Δt ≤ Δx²/(2ν). These conditions ensure that numerical information does not propagate faster than the physics of the problem."
}

Why this format?

  • Clear structure: The model learns to distinguish the question (Instruction) from the answer (Response).
  • Compatibility: trl’s SFTTrainer is optimized for this format, which simplifies training.

Fine-tuning Methodology: Zephyr-7B + LoRA in 4-bit

Model and Technique Choice

I selected Zephyr-7B for its instruction-following ability and its size, which is suitable for a consumer GPU. To reduce memory usage, I used:

  • LoRA (Low-Rank Adaptation): A technique that trains only a few additional weights (the “adapters”), without modifying the base model.
  • 4-bit quantization: Reduces weight precision to save memory (via bitsandbytes).

Technical Setup

  1. Model loading: 
    from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,          # 4-bit quantization
    bnb_4bit_quant_type="nf4",  # Optimal quantization type
    bnb_4bit_compute_dtype=torch.float16,  # FP16 computation
)

model = AutoModelForCausalLM.from_pretrained(
    "HuggingFaceH4/zephyr-7b-beta",
    quantization_config=bnb_config,
    device_map="auto",  # Automatically uses the GPU
)
  2. LoRA setup: 
    from peft import LoraConfig

lora_config = LoraConfig(
    r=8,                     # LoRA matrix rank (smaller = fewer parameters)
    lora_alpha=32,           # Scaling factor
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # Target layers
    lora_dropout=0.05,       # Dropout to reduce overfitting
    bias="none",             # No bias in the adapters
    task_type="CAUSAL_LM",   # Task type (language model)
)
  3. Preparing the model for LoRA: 
    model = prepare_model_for_kbit_training(model)  # Prepares the model for LoRA
model = get_peft_model(model, lora_config)      # Applies LoRA

Memory Optimizations

To avoid CUDA out of memory errors on a T4 GPU (15 GB of VRAM), I:

  • Reduced the batch size: per_device_train_batch_size=2.
  • Used gradient accumulation: gradient_accumulation_steps=8 (equivalent to a batch size of 16).
  • Limited the sequence length: max_seq_length=256.
  • Cleared memory before training: 
    import torch, gc
torch.cuda.empty_cache()
gc.collect()

Training: Parameters and Results

Training Parameters

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="zephyr-volumes-finis",      # Output directory
    per_device_train_batch_size=2,          # Batch size
    gradient_accumulation_steps=8,          # Gradient accumulation
    learning_rate=2e-4,                     # Learning rate
    num_train_epochs=3,                     # Number of epochs
    fp16=True,                              # Mixed precision (FP16)
    save_steps=50,                          # Periodic checkpoints
    logging_steps=10,                       # Periodic logs
    optim="paged_adamw_8bit",                # Memory-efficient optimizer
    report_to="none",                       # Disable Weights & Biases
)

Running Training

from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset,          # Prepared dataset
    peft_config=lora_config,        # LoRA configuration
    max_seq_length=256,             # Maximum sequence length
)

trainer.train()  # Start training

Results and Observations

  • Duration: ~ minutes for 3 epochs on a few examples (T4 GPU).
  • Loss progression:
    • Epoch 1: Loss ~0.5 (the model is learning the answer structure).
    • Epoch 2: Loss ~0.25 (the answers become more precise).
    • Epoch 3: Loss ~0.15 (stabilization, correct generalization).
  • Answer quality:
    • The model generates precise technical explanations, even for questions not present in the dataset.
    • The responses are structured and context-aware (for example, convection/diffusion distinctions).

Example generated answer for “What are the advantages of an implicit scheme?”:

*“An implicit scheme has three major advantages:

  1. Unconditional stability: Unlike explicit schemes, it is not subject to the CFL condition, which allows larger time steps (Δt ≫ Δx²/(2ν) for diffusion).
  2. Time accuracy: By avoiding restrictions on Δt, it better captures slow dynamics (for example, diffusion over long periods).
  3. Robustness: Less sensitive to numerical instabilities, even for stiff problems (for example, equations with strong source terms). Drawback: It requires solving a linear system at each time step, which increases the computational cost (tridiagonal matrix for 1D problems).”*

Validation and Limits

To validate the model, I tested it on questions not present in the dataset:

  1. “How do you discretize a source term in the convection-diffusion equation?” → A coherent answer explaining how the source term is integrated into the flux balance.
  2. “What is the influence of the mesh on the accuracy of a finite volume scheme?” → A detailed answer about truncation error and convergence.

Identified limits:

  • Limited precision for very specific problems (for example, 3D unstructured meshes).
  • Dataset dependence: The model generalizes well within its training domain, but may lack detail for uncovered cases.

Future Improvements

  1. Expand the dataset:
    • Add examples on non-uniform meshes, 2D/3D problems, and high-order schemes.
    • Include cases with source terms or chemical reactions (for example, a reaction-diffusion equation).
  2. Optimize training:
    • Test with more epochs (5-10) to improve accuracy.
    • Use a validation set to evaluate generalization.
  3. Practical integration:
    • Develop an API to interact with the model through a web interface.
    • Create a Jupyter notebook with interactive PDE-solving examples.

Conclusion

This project shows that a large model like Zephyr-7B can be adapted to a specific technical domain with little data and training time. The model already answers finite volume questions correctly and explains the concepts clearly.

Why is this useful?

  • Saves time when looking for explanations.
  • Gives pedagogical answers adapted to the audience.
  • Can be extended to other numerical methods.

Next steps:

  • …

Project link: GitHub Gist