Skip to content
BranchesTags

Repository files navigation

LLM Fine-Tuning + NVIDIA NIM Deployment

End-to-end pipeline for fine-tuning a domain-specific language model (Llama-3.2-3B-Instruct) on arXiv paper data and deploying it as a production inference service using NVIDIA NIM.

Overview

Training: LoRA Supervised Fine-Tuning with Unsloth
Deployment: NVIDIA NIM + OpenAI-compatible API
Goal: Complete ML lifecycle from data preparation through training, evaluation, and production deployment

Architecture

arXiv Papers (Qdrant) 
  → Dataset Generation (Evol-Instruct + CoT)
  → LoRA Fine-Tuning (Unsloth)
  → Model Evaluation
  → NIM Deployment
  → Production API

Table of Contents

Prerequisites

  • NVIDIA GPU with CUDA support (24GB+ VRAM recommended)
  • Docker installed and configured
  • NVIDIA Container Toolkit installed
  • Python 3.10+
  • NGC API Key (for NIM deployment - get from https://ngc.nvidia.com/setup/api-key)

Setup

1. Docker Installation

# Update packages
sudo apt-get update
sudo apt-get install -y \
    ca-certificates \
    curl \
    gnupg \
    lsb-release

# Add Docker's GPG key
sudo mkdir -p /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg

# Set up Docker repository
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
  $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

# Install Docker
sudo apt-get update
sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

# Add user to docker group (logout/login required)
sudo usermod -aG docker $USER

2. NVIDIA Container Toolkit

# Install prerequisites
sudo apt-get update && sudo apt-get install -y --no-install-recommends \
    ca-certificates \
    curl \
    gnupg2

# Configure NVIDIA repository
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

# Install NVIDIA Container Toolkit
sudo apt-get update
export NVIDIA_CONTAINER_TOOLKIT_VERSION=1.18.2-1
sudo apt-get install -y \
    nvidia-container-toolkit=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
    nvidia-container-toolkit-base=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
    libnvidia-container-tools=${NVIDIA_CONTAINER_TOOLKIT_VERSION} \
    libnvidia-container1=${NVIDIA_CONTAINER_TOOLKIT_VERSION}

# Configure Docker
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

3. Python Environment

# Create virtual environment
python3 -m venv .venv
source .venv/bin/activate

# Install dependencies
pip install -r requirements.txt

4. Environment Variables

Create a .env file in the project root:

# For dataset generation (if using)
QDRANT_URL=http://localhost
QDRANT_PORT=6333
OPENAI_API_KEY=your_openai_api_key
OPENAI_MODEL=gpt-4o-mini  # or gpt-4, gpt-3.5-turbo, etc.

# For NIM deployment
NGC_API_KEY=your_ngc_api_key
LOCAL_PEFT_DIRECTORY=outputs

Dataset Generation

Generate instruction-tuning dataset from arXiv paper chunks using Evol-Instruct + Chain-of-Thought.

Pipeline Overview

  1. Extract Chunks: Extract text chunks from Qdrant collection
  2. Generate Pairs: Use OpenAI LLM to generate Evol-Instruct + CoT pairs
  3. Filter Quality: Filter low-quality or malformed examples
  4. Format Alpaca: Convert to standard Alpaca format and split train/eval

Usage

python generate_dataset.py \
    --collection-name your_collection \
    --output-dir datasets/arxiv \
    --batch-size 50

Options

  • --skip-extraction: Skip chunk extraction (use existing raw_chunks.json)
  • --skip-generation: Skip generation (use existing generated_pairs.jsonl)
  • --skip-filtering: Skip quality filtering
  • --skip-formatting: Skip Alpaca formatting
  • --resume: Resume from checkpoint if available
  • --max-chunks N: Limit number of chunks (for testing)

Dataset Format

Output: datasets/arxiv/train.jsonl and datasets/arxiv/eval.jsonl

Format (Alpaca):

{
  "instruction": "Your analytical question here",
  "input": "Context or input text",
  "output": "Chain-of-thought reasoning answer"
}

Training

Fine-tune Llama-3.2-3B-Instruct using LoRA with Unsloth for efficient training.

Basic Training

python train.py \
    --train_data datasets/arxiv/train.jsonl \
    --eval_data datasets/arxiv/eval.jsonl \
    --output_dir outputs

With Custom Hyperparameters

python train.py \
    --train_data datasets/arxiv/train.jsonl \
    --eval_data datasets/arxiv/eval.jsonl \
    --output_dir outputs \
    --max_steps 1000 \
    --learning_rate 1e-4 \
    --per_device_batch_size 1 \
    --gradient_accumulation_steps 8

Resume from Checkpoint

# Resume from specific checkpoint
python train.py \
    --train_data datasets/arxiv/train.jsonl \
    --eval_data datasets/arxiv/eval.jsonl \
    --resume_from_checkpoint outputs/checkpoint-100

# Auto-resume from latest checkpoint
python train.py \
    --train_data datasets/arxiv/train.jsonl \
    --eval_data datasets/arxiv/eval.jsonl \
    --resume

Default Configuration

  • Model: unsloth/Llama-3.2-3B-Instruct
  • Max Sequence Length: 8192 tokens
  • LoRA Rank: 32
  • LoRA Alpha: 32
  • Batch Size: 2 per device, 4 gradient accumulation (effective = 8)
  • Max Steps: 600 (~2-3 epochs)
  • Learning Rate: 2e-4
  • Checkpointing: Every 50 steps, keep last 3
  • Quantization: 4-bit (bitsandbytes)

Key Features

  • Memory Efficient: 4-bit quantization + LoRA adapters
  • Long Sequences: Handles sequences up to 8192 tokens
  • Sequence Packing: Automatically packs multiple samples
  • Checkpointing: Automatic with resume support
  • Wandb Tracking: Experiment tracking (optional)

Output

After training, model saved to output_dir:

  • adapter_model.safetensors - LoRA adapter weights
  • adapter_config.json - Adapter configuration
  • tokenizer*.json - Tokenizer files
  • checkpoint-*/ - Training checkpoints

Evaluation

Compare fine-tuned model against baseline.

Basic Evaluation

python evaluate.py \
    --eval_data datasets/arxiv/eval.jsonl \
    --finetuned_model outputs

With LLM-as-Judge

python evaluate.py \
    --eval_data datasets/arxiv/eval.jsonl \
    --finetuned_model outputs \
    --use_llm_judge \
    --llm_judge_sample_size 50

Evaluation Metrics

  • Traditional Metrics: BLEU, ROUGE-L, ROUGE-1, ROUGE-2
  • LLM-as-Judge: GPT-4 evaluation of response quality
  • Side-by-Side Comparison: Manual review of examples

Output

Evaluation reports saved to evaluation_reports/:

  • evaluation_report_*.json - Full results
  • evaluation_report_*.md - Human-readable summary

Deployment

Deploy fine-tuned model using NVIDIA NIM (NVIDIA Inference Microservices) with Multi-LoRA support.

1. Setup NIM

cd deployment
chmod +x setup_nim.sh
./setup_nim.sh

This authenticates with NVIDIA Container Registry and pulls the NIM container.

2. Prepare Model Repository

python3 -m deployment.prepare_model_repo

This organizes adapter files into NIM model repository structure:

model_repository/
└── llama-3.2-3b-arxiv/
    └── 1/
        ├── adapter_config.json
        ├── adapter_model.safetensors
        └── ...

3. Start NIM Container

chmod +x run_nim.sh
./run_nim.sh

Note: First run compiles TensorRT-LLM engine (10-30 minutes).

4. Wait for Service Ready

python3 -m deployment.wait_for_nim

5. Test Inference

# Basic test
python3 -m deployment.test_endpoint

# Custom prompt
python3 -m deployment.test_endpoint --prompt "Explain attention mechanisms"

# Test with evaluation dataset
python3 -m deployment.test_endpoint --test-file ../datasets/arxiv/eval.jsonl

# Test concurrency
python3 -m deployment.test_concurrency --requests 20 --concurrency 5

API Usage

NIM exposes OpenAI-compatible API at http://localhost:8000/v1/chat/completions:

curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "llama-3.2-3b-arxiv",
    "messages": [
      {"role": "system", "content": "You are a helpful AI assistant."},
      {"role": "user", "content": "Explain attention mechanisms."}
    ],
    "temperature": 0.7,
    "max_tokens": 512
  }'

Important: The model parameter must match your adapter name (default: llama-3.2-3b-arxiv).

Monitor GPU

python3 -m deployment.monitor_gpu --duration 300

Project Structure

finetune_nim/
├── dataset_generation/          # Dataset generation pipeline
│   ├── extract_chunks.py        # Qdrant extraction
│   ├── generate_pairs.py        # OpenAI LLM generation
│   ├── filter_quality.py        # Quality filtering
│   ├── format_alpaca.py         # Alpaca formatting
│   ├── openai_client.py         # OpenAI client
│   └── config.py                # Configuration
├── training/                     # Training pipeline
│   ├── trainer.py               # ArxivFineTuner class
│   ├── data_formatter.py        # Dataset formatting
│   └── config.py               # Training config
├── evaluation/                  # Evaluation system
│   ├── evaluator.py            # ModelEvaluator class
│   ├── metrics.py              # Traditional metrics
│   ├── llm_judge.py            # LLM-as-judge
│   ├── report.py               # Report generation
│   └── config.py               # Evaluation config
├── deployment/                  # NIM deployment
│   ├── prepare_model_repo.py   # Repository preparation
│   ├── run_nim.sh              # Start NIM container
│   ├── setup_nim.sh            # Setup script
│   ├── wait_for_nim.py         # Wait for readiness
│   ├── test_endpoint.py        # Test inference
│   ├── test_concurrency.py     # Concurrency testing
│   ├── monitor_gpu.py          # GPU monitoring
│   └── config.py               # Deployment config
├── generate_dataset.py          # Dataset generation CLI
├── train.py                     # Training CLI
├── evaluate.py                  # Evaluation CLI
└── requirements.txt             # Dependencies

Troubleshooting

Training Issues

Out of Memory (OOM)

# Reduce batch size and increase gradient accumulation
--per_device_batch_size 1 --gradient_accumulation_steps 8

CUDA Not Available

  • Ensure NVIDIA GPU with CUDA support
  • Verify CUDA toolkit installed
  • Check PyTorch has CUDA support

Deployment Issues

Container Won't Start

  • Check GPU: nvidia-smi
  • Verify NVIDIA Container Toolkit: docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi
  • Check logs: docker logs nim-llama-3.2-3b

Authentication Errors

  • Verify NGC_API_KEY in .env
  • Test login: echo $NGC_API_KEY | docker login nvcr.io -u $oauthtoken --password-stdin

Model Not Found

  • Verify repository structure: ls -la model_repository/{adapter_name}/1/
  • Check required files: adapter_config.json, adapter_model.safetensors

Port Already in Use

  • Change port in .env: NIM_PORT=8001
  • Or stop existing container: docker stop nim-llama-3.2-3b

Dataset Generation Issues

API Rate Limits

  • Reduce batch size: --batch-size 25
  • Use --resume to continue from checkpoint
  • Consider using a cheaper model like gpt-4o-mini for cost efficiency

Malformed JSON

  • Quality filtering should remove bad examples
  • Check logs for parsing errors

Quick Reference

Training

python train.py \
    --train_data datasets/arxiv/train.jsonl \
    --eval_data datasets/arxiv/eval.jsonl \
    --output_dir outputs

Evaluation

python evaluate.py \
    --eval_data datasets/arxiv/eval.jsonl \
    --finetuned_model outputs

Deploy NIM

cd deployment
./setup_nim.sh
python3 -m deployment.prepare_model_repo
./run_nim.sh

Test Inference

python3 -m deployment.test_endpoint --prompt "Your prompt here"

Key Technologies

  • Training: Unsloth, LoRA, 4-bit quantization
  • Base Model: Llama-3.2-3B-Instruct
  • Deployment: NVIDIA NIM (TensorRT-LLM)
  • Dataset Format: Alpaca (instruction/input/output)
  • Tracking: Weights & Biases (optional)

License

See individual component licenses.

About

No description, website, or topics provided.

Resources

Readme
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages