A Step-by-Step Tutorial for Fine-Tuning CodeLlama 70B with LoRA and 4x A100s

import subprocess
import sys
import torch
import warnings
from packaging import version
from typing import List, Dict

def run_cmd(cmd: str, check: bool = True) -> subprocess.CompletedProcess:
    """Run a shell command and handle errors with informative messages."""
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            check=check,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        if result.stderr:
            warnings.warn(f"Command {cmd} produced stderr: {result.stderr}")
        return result
    except subprocess.CalledProcessError as e:
        print(f"❌ Command failed: {cmd}")
        print(f"Stdout: {e.stdout}")
        print(f"Stderr: {e.stderr}")
        sys.exit(1)

def validate_gpu_setup() -> None:
    """Check that 4x A100 80GB GPUs are available and correctly configured."""
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available. Check NVIDIA driver installation.")
    gpu_count = torch.cuda.device_count()
    if gpu_count != 4:
        raise RuntimeError(f"Expected 4 GPUs, found {gpu_count}. This tutorial requires 4x A100 80GB GPUs.")
    for i in range(gpu_count):
        gpu_name = torch.cuda.get_device_name(i)
        if "A100" not in gpu_name or "80GB" not in gpu_name:
            raise RuntimeError(f"GPU {i} is {gpu_name}, expected A100 80GB.")
        mem = torch.cuda.get_device_properties(i).total_mem
        if mem < 80 * 1024 * 1024 * 1024:  # 80GB in bytes
            raise RuntimeError(f"GPU {i} has {mem//1024**3}GB memory, expected 80GB.")
    print(f"✅ Validated {gpu_count}x A100 80GB GPUs")

def install_dependencies() -> None:
    """Install exact dependency versions validated for this tutorial."""
    deps = [
        "torch==2.1.0 --index-url https://download.pytorch.org/whl/cu121",
        "transformers==4.36.2",
        "peft==0.7.1",
        "deepspeed==0.12.0",
        "datasets==2.16.1",
        "accelerate==0.25.0",
        "bitsandbytes==0.41.1",
        "evaluate==0.4.1",
        "rouge-score==0.1.2"
    ]
    for dep in deps:
        print(f"Installing {dep.split('==')[0]}...")
        run_cmd(f"pip install {dep}")

if __name__ == "__main__":
    print("Starting environment setup for CodeLlama 70B LoRA fine-tuning...")
    validate_gpu_setup()
    install_dependencies()
    # Validate installed versions
    import transformers
    import peft
    import deepspeed
    assert version.parse(transformers.__version__) >= version.parse("4.36.2"), "Transformers version too old"
    assert version.parse(peft.__version__) >= version.parse("0.7.1"), "PEFT version too old"
    assert version.parse(deepspeed.__version__) >= version.parse("0.12.0"), "DeepSpeed version too old"
    print("✅ All dependencies installed and validated")

import os
import json
import glob
import ast
import logging
from typing import List, Dict, Optional
from datasets import Dataset, DatasetDict
from transformers import AutoTokenizer

# Configure logging for error tracking
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

class CodeDatasetProcessor:
    def __init__(self, tokenizer_name: str = "codellama/CodeLlama-70b-hf", max_length: int = 2048):
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.max_length = max_length
        logger.info(f"Initialized processor with tokenizer {tokenizer_name}, max length {max_length}")

    def extract_python_functions(self, file_path: str) -> List[Dict]:
        """Extract function definitions from a Python file, handling syntax errors."""
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                source = f.read()
        except UnicodeDecodeError:
            logger.warning(f"Skipping {file_path}: non-UTF-8 encoding")
            return []
        try:
            tree = ast.parse(source)
        except SyntaxError as e:
            logger.warning(f"Skipping {file_path}: syntax error {e}")
            return []
        functions = []
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                func_source = ast.get_source_segment(source, node)
                if func_source is None:
                    continue
                # Create instruction-response pair: complete the function given docstring
                docstring = ast.get_docstring(node)
                instruction = f"Complete the following Python function:\n{def node.name} {ast.unparse(node.args)}"
                if docstring:
                    instruction += f"\nDocstring: {docstring}"
                response = func_source
                functions.append({
                    "instruction": instruction,
                    "response": response,
                    "file_path": file_path
                })
        return functions

    def process_codebase(self, codebase_dir: str) -> List[Dict]:
        """Recursively process all Python files in a codebase directory."""
        python_files = glob.glob(os.path.join(codebase_dir, "**/*.py"), recursive=True)
        logger.info(f"Found {len(python_files)} Python files in {codebase_dir}")
        all_samples = []
        for file_path in python_files:
            try:
                samples = self.extract_python_functions(file_path)
                all_samples.extend(samples)
            except Exception as e:
                logger.error(f"Failed to process {file_path}: {e}")
        logger.info(f"Extracted {len(all_samples)} total function samples")
        return all_samples

    def format_for_training(self, samples: List[Dict]) -> Dataset:
        """Format samples into CodeLlama's training format with prompt templating."""
        def tokenize_fn(examples: Dict) -> Dict:
            prompts = []
            for instr, resp in zip(examples["instruction"], examples["response"]):
                # CodeLlama instruction format
                prompt = f"[INST] {instr} [/INST] {resp}"
                prompts.append(prompt)
            tokenized = self.tokenizer(
                prompts,
                max_length=self.max_length,
                padding="max_length",
                truncation=True,
                return_tensors="pt"
            )
            tokenized["labels"] = tokenized["input_ids"].clone()
            # Mask instruction tokens to not compute loss on them
            for i, (instr, resp) in enumerate(zip(examples["instruction"], examples["response"])):
                instr_len = len(self.tokenizer(f"[INST] {instr} [/INST]", return_tensors="pt")["input_ids"][0])
                tokenized["labels"][i, :instr_len] = -100
            return tokenized

        dataset = Dataset.from_list(samples)
        tokenized_dataset = dataset.map(
            tokenize_fn,
            batched=True,
            remove_columns=["instruction", "response", "file_path"]
        )
        return tokenized_dataset

if __name__ == "__main__":
    processor = CodeDatasetProcessor()
    # Process internal codebase (replace with your own path)
    samples = processor.process_codebase("./internal_python_codebase")
    if len(samples) < 1000:
        logger.warning(f"Only {len(samples)} samples found. Recommended minimum 10k for 70B fine-tuning.")
    tokenized_train = processor.format_for_training(samples)
    # Split into train/validation (90/10)
    dataset_dict = DatasetDict({
        "train": tokenized_train.shuffle(seed=42).select(range(int(0.9 * len(tokenized_train)))),
        "validation": tokenized_train.shuffle(seed=42).select(range(int(0.9 * len(tokenized_train)), len(tokenized_train)))
    })
    dataset_dict.save_to_disk("./processed_codellama_dataset")
    logger.info(f"Saved processed dataset to ./processed_codellama_dataset with {len(dataset_dict['train'])} train, {len(dataset_dict['validation'])} validation samples")

import os
import sys
import json
import logging
import argparse
from typing import Optional
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_from_disk
import deepspeed

# Enable DeepSpeed distributed training
os.environ["TOKENIZERS_PARALLELISM"] = "false"
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def parse_args():
    parser = argparse.ArgumentParser(description="Fine-tune CodeLlama 70B with LoRA on 4x A100s")
    parser.add_argument("--model_name", type=str, default="codellama/CodeLlama-70b-hf")
    parser.add_argument("--dataset_path", type=str, default="./processed_codellama_dataset")
    parser.add_argument("--output_dir", type=str, default="./codellama-70b-lora-finetuned")
    parser.add_argument("--lora_r", type=int, default=64, help="LoRA rank")
    parser.add_argument("--lora_alpha", type=int, default=128, help="LoRA alpha")
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--batch_size", type=int, default=1, help="Per-device batch size")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=16)
    return parser.parse_args()

def main():
    args = parse_args()
    logger.info(f"Starting training with args: {args}")

    # Load tokenizer
    try:
        tokenizer = AutoTokenizer.from_pretrained(args.model_name)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        logger.info(f"Loaded tokenizer {args.model_name}")
    except Exception as e:
        logger.error(f"Failed to load tokenizer: {e}")
        sys.exit(1)

    # Load model in 4-bit precision to fit 4x A100 80GB
    try:
        model = AutoModelForCausalLM.from_pretrained(
            args.model_name,
            load_in_4bit=True,
            device_map="auto",
            torch_dtype=torch.bfloat16,
            use_flash_attention_2=True  # Requires A100 CUDA 12.1+
        )
        model = prepare_model_for_kbit_training(model)
        logger.info(f"Loaded model {args.model_name} in 4-bit precision")
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        sys.exit(1)

    # Configure LoRA
    lora_config = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],  # All linear layers in CodeLlama
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM"
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()  # Should show ~0.1% trainable params
    logger.info(f"Applied LoRA config: r={args.lora_r}, alpha={args.lora_alpha}")

    # Load dataset
    try:
        dataset = load_from_disk(args.dataset_path)
        logger.info(f"Loaded dataset: {dataset}")
    except Exception as e:
        logger.error(f"Failed to load dataset: {e}")
        sys.exit(1)

    # Training arguments with DeepSpeed
    training_args = TrainingArguments(
        output_dir=args.output_dir,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        num_train_epochs=args.epochs,
        learning_rate=2e-4,
        bf16=True,
        save_steps=500,
        save_total_limit=3,
        evaluation_strategy="steps",
        eval_steps=500,
        logging_steps=10,
        report_to="none",  # Disable wandb/tensorboard unless configured
        deepspeed="ds_config.json",  # DeepSpeed config file
        remove_unused_columns=False
    )

    # Data collator
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # Initialize Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        data_collator=data_collator,
        tokenizer=tokenizer
    )

    # Start training
    try:
        trainer.train()
        trainer.save_model(args.output_dir)
        logger.info(f"Training complete. Model saved to {args.output_dir}")
    except Exception as e:
        logger.error(f"Training failed: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()

{
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "bf16": {"enabled": true},
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "none"},
        "offload_param": {"device": "none"}
    },
    "steps_per_print": 10
}

# Load model with Flash Attention 2 and 4-bit quantization
model = AutoModelForCausalLM.from_pretrained(
    "codellama/CodeLlama-70b-hf",
    load_in_4bit=True,
    use_flash_attention_2=True,  # Requires CUDA 12.1+ and A100
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# LoRA config tuned for 70B code models
lora_config = LoraConfig(
    r=64,
    lora_alpha=128,  # 2*r as per best practice
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    task_type="CAUSAL_LM"
)

{
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "bf16": {"enabled": true},
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "none"},
        "offload_param": {"device": "none"}
    },
    "steps_per_print": 10
}

codellama-70b-lora-finetuning/
├── setup/
│   └── 01_setup_environment.py  # Environment validation and dependency install
├── data/
│   ├── 02_process_codebase.py    # Codebase processing and dataset creation
│   └── processed_dataset/        # Saved tokenized dataset
├── training/
│   ├── 03_train_lora.py          # Main training script
│   ├── ds_config.json            # DeepSpeed ZeRO-3 config
│   └── lora_config.json          # LoRA hyperparameters
├── inference/
│   └── 04_deploy_vllm.py         # vLLM deployment script for fine-tuned model
├── benchmarks/
│   └── accuracy_eval.py          # Human eval and ROUGE score calculation
└── README.md                     # Full tutorial and setup instructions