← Back to Blog

Fine-Tuning Open-Source LLMs on Private Enterprise Codebases



Fine-Tuning Open-Source LLMs on Private Enterprise Codebases

Privacy vs. Performance

Off-the-shelf LLMs generate acceptable generic code, but they often struggle with proprietary frameworks, internal design patterns, and custom microservices. RAG helps introduce context into prompts, but it does not modify how a model naturally generates syntax. Fine-tuning an open-source model such as Llama 3 or Mistral on your codebase resolves this issue, though traditional full-parameter fine-tuning requires substantial VRAM. Quantized Low-Rank Adaptation (QLoRA) makes this practical using Parameter-Efficient Fine-Tuning (PEFT).

Training Setup with QLoRA

Here is a Python script utilizing Hugging Face's transformers, trl, and peft libraries. It quantizes the base model to 4-bit precision, enabling training on a single GPU (such as an NVIDIA A10G or L4) while fine-tuning adapter weights on your internal codebase.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer

# Line 1: Configure 4-bit quantization to drastically reduce memory footprint bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16 )

# Line 2: Load the base model with the quantization configuration model_id = "meta-llama/Meta-Llama-3-8B-Instruct" model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map="auto") model = prepare_model_for_kbit_training(model)

# Line 3: Define the LoRA configuration targeting attention modules peft_config = LoraConfig( r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM" ) model = get_peft_model(model, peft_config)

# Line 4: Initialize the Supervised Fine-Tuning (SFT) Trainer trainer = SFTTrainer( model=model, train_dataset=internal_codebase_dataset, # Pre-formatted Dataset object dataset_text_field="text", max_seq_length=2048, args=TrainingArguments( per_device_train_batch_size=4, gradient_accumulation_steps=4, warmup_steps=100, learning_rate=2e-4, fp16=True, output_dir="outputs/enterprise-coder-v1" ), peft_config=peft_config, )

# Line 5: Execute training and save the adapter weights trainer.train() trainer.model.save_pretrained("enterprise-coder-v1-adapter")

How the Configuration Works

  • Line 1: bnb_4bit_quant_type="nf4" (Normal Float 4) performs effectively for normally distributed LLM weights. Combined with a bfloat16 compute dtype, weights are stored in 4-bit format while matrix multiplications in the forward and backward passes execute in 16-bit to maintain gradient accuracy.
  • Line 2: prepare_model_for_kbit_training converts layer norms and the final LM head to fp32. This avoids loss spikes and numerical instability when fine-tuning quantized models.
  • Line 3: The LoRA setup targets the query, key, value, and output projection matrices (q_proj, k_proj, v_proj, o_proj). Focusing on attention modules instead of MLP layers balances training speed and syntax adaptation without causing catastrophic forgetting.
  • Lines 4 & 5: Gradient accumulation (per_device_train_batch_size=4 and gradient_accumulation_steps=4) creates an effective batch size of 16 to reduce VRAM consumption. Saving only the adapter weights produces a compact output (~100MB) rather than saving the full base model again.

Utilizing QLoRA allows engineering teams to adapt open-source models to private codebases entirely behind a firewall, eliminating privacy risks while gaining domain-specific coding assistance.