Maaac/CodeLLaMA-Linux-BugFix

1# QLoRA fine-tuning for CodeLLaMA-7B-Instruct on 1x H200
2# Requirements: transformers, peft, accelerate, bitsandbytes, datasets
3from transformers import (
4    AutoTokenizer,
5    AutoModelForCausalLM,
6    TrainingArguments,
7    Trainer,
8    BitsAndBytesConfig,
9    DataCollatorForSeq2Seq
10)
11from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
12from datasets import load_dataset
13import torch
14import os
15import wandb
16
17os.environ["WANDB_PROJECT"] = "codellama-7b-instruct-qlora-linux-bugfix"
18os.environ["WANDB_NAME"] = "run-v1"
19# Paths and model
20BASE_MODEL = "codellama/CodeLLaMA-7b-Instruct-hf"
21DATA_PATH = "../dataset/training_data_100k.jsonl"
22OUTPUT_DIR = "./output/qlora-codellama-bugfix"
23
24# Load dataset (prompt-completion format)
25dataset = load_dataset("json", data_files=DATA_PATH, split="train")
26
27# BitsandBytes config for QLoRA
28bnb_config = BitsAndBytesConfig(
29    load_in_4bit=True,
30    bnb_4bit_use_double_quant=True,
31    bnb_4bit_quant_type="nf4",
32    bnb_4bit_compute_dtype=torch.bfloat16  # optimized for H100/H200
33)
34
35# Load tokenizer and model
36tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
37tokenizer.pad_token = tokenizer.eos_token
38tokenizer.padding_side = "right"
39
40model = AutoModelForCausalLM.from_pretrained(
41    BASE_MODEL,
42    quantization_config=bnb_config,
43    device_map="auto"
44)
45model = prepare_model_for_kbit_training(model)
46model.gradient_checkpointing_enable()
47torch.backends.cuda.matmul.allow_tf32 = True
48
49# Apply QLoRA (LoRA config)
50lora_config = LoraConfig(
51    r=64,
52    lora_alpha=16,
53    lora_dropout=0.1,
54    bias="none",
55    task_type="CAUSAL_LM"
56)
57model = get_peft_model(model, lora_config)
58model.config.use_cache = False
59model.config.return_dict = True
60model.config.pad_token_id = tokenizer.pad_token_id
61model.print_trainable_parameters()
62
63# Format and tokenize the dataset
64model_max_len = tokenizer.model_max_length
65
66def format(example):
67    prompt_ids = tokenizer(example["prompt"], truncation=True, max_length=1024)["input_ids"]
68    completion_ids = tokenizer(example["completion"], truncation=True, max_length=512)["input_ids"]
69
70    input_ids = prompt_ids + completion_ids
71    labels = [-100] * len(prompt_ids) + completion_ids
72
73    # pad both input_ids and labels to the same length
74    max_len = min(len(input_ids), tokenizer.model_max_length)
75    input_ids = input_ids[:max_len]
76    labels = labels[:max_len]
77
78    return {
79        "input_ids": input_ids,
80        "labels": labels,
81    }
82
83
84# Sanity check
85print("__ Sanity checking one example...")
86sample = format(dataset[0])
87test_input = torch.tensor(sample["input_ids"]).unsqueeze(0).to(model.device)
88test_labels = torch.tensor(sample["labels"]).unsqueeze(0).to(model.device)
89model.train()
90out = model(input_ids=test_input, labels=test_labels)
91assert out.loss.requires_grad, "Sanity check failed: Loss does not require grad."
92print("__ Sanity check passed. Proceeding to map()...")
93
94# Apply formatting to entire dataset
95dataset = dataset.map(format, remove_columns=["prompt", "completion"])
96collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model, return_tensors="pt", pad_to_multiple_of=8)
97
98# Training arguments
99training_args = TrainingArguments(
100    report_to="wandb",
101    run_name="codellama-7b-instruct-qlora-linux-bugfix",
102    logging_dir=f"{OUTPUT_DIR}/logs",
103    
104    output_dir=OUTPUT_DIR,
105    num_train_epochs=3,
106    per_device_train_batch_size=64,
107    gradient_accumulation_steps=4,
108    learning_rate=2e-4,
109    lr_scheduler_type="cosine",
110    warmup_ratio=0.03,
111    gradient_checkpointing=True,
112    bf16=True,                     # Important for H200
113    fp16=False,
114    max_grad_norm=1.0,
115    save_strategy="steps",
116    save_steps=500,
117    save_total_limit=2,
118    logging_steps=50,
119    push_to_hub=False,
120    label_names=["labels"],
121    remove_unused_columns=False,  # Critical to prevent data loss
122)
123
124# Trainer setup
125trainer = Trainer(
126    model=model,
127    args=training_args,
128    train_dataset=dataset,
129    tokenizer=tokenizer,
130    data_collator=collator
131)
132
133
134# Begin training
135model.train()
136print(f"Track this run in Weights & Biases: https://wandb.ai/{os.environ['WANDB_PROJECT']}/{os.environ['WANDB_NAME']}")
137trainer.train(resume_from_checkpoint=True)
138
139
140# Save final model
141model.save_pretrained(OUTPUT_DIR, safe_serialization=True)
142tokenizer.save_pretrained(OUTPUT_DIR)
143print(f"[DONE] Model saved to {OUTPUT_DIR}")
144