Update train_llama4.py

train_llama4.py

@@ -1,8 +1,7 @@
 # train_llama4.py
-# Script to fine-tune Llama 4 Maverick for healthcare fraud detection (text-only with CPU offloading)
+# Script to fine-tune Llama 4 Maverick for healthcare fraud detection
 
-from transformers import AutoTokenizer, Llama4ForConditionalGeneration, Trainer, TrainingArguments
-from transformers import BitsAndBytesConfig
+from transformers import AutoTokenizer, Llama4ForConditionalGeneration
 import datasets
 import torch
 from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
@@ -10,11 +9,8 @@ from accelerate import Accelerator
 import huggingface_hub
 import os
 
-# Version and CUDA check
-print(f"PyTorch version: {torch.__version__}")
-print(f"CUDA version: {torch.version.cuda}")
-print(f"Is CUDA available: {torch.cuda.is_available()}")
-print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+# Debug: Confirm file version
+print("Running train_llama4.py with CPU offloading (version: 2025-04-21 v2)")
 
 # Authenticate with Hugging Face
 LLama = os.getenv("LLama")
@@ -22,111 +18,85 @@ if not LLama:
     raise ValueError("LLama token not found. Set it in Hugging Face Space secrets as 'LLama'.")
 huggingface_hub.login(token=LLama)
 
-# Load Llama 4 model and tokenizer
+# Model setup
 MODEL_ID = "meta-llama/Llama-4-Maverick-17B-128E-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-# Add padding token if it doesn't exist
 if tokenizer.pad_token is None:
     tokenizer.add_special_tokens({'pad_token': '[PAD]'})
 
-# Custom device map to offload some layers to CPU
+# Custom device map for CPU offloading
 device_map = {
     "model.embed_tokens": 0,
-    "model.layers.0-15": 0,  # Keep first 16 layers on GPU
-    "model.layers.16-31": "cpu",  # Offload remaining layers to CPU
+    "model.layers.0-15": 0,
+    "model.layers.16-31": "cpu",
     "model.norm": 0,
     "lm_head": 0
 }
 
-# Quantization config for A100 80 GB VRAM
-quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+# Debug: Confirm offloading settings
+print("Loading model with CPU offloading: llm_int8_enable_fp32_cpu_offload=True, device_map=", device_map)
 
+# Load model with 8-bit quantization and CPU offloading
 model = Llama4ForConditionalGeneration.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.bfloat16,
     device_map=device_map,
-    quantization_config=quantization_config,
+    quantization_config={"load_in_8bit": True},
     llm_int8_enable_fp32_cpu_offload=True,
     attn_implementation="flex_attention"
 )
 
-# Prepare for LoRA
-model = prepare_model_for_kbit_training(model)
-peft_config = LoraConfig(
+# Resize token embeddings
+model.resize_token_embeddings(len(tokenizer))
+
+# Initialize Accelerator
+accelerator = Accelerator()
+model = accelerator.prepare(model)
+
+# Load dataset
+dataset = datasets.load_dataset('json', data_files="Bingaman_training_data.json")['train']
+
+# LoRA configuration
+lora_config = LoraConfig(
     r=16,
     lora_alpha=32,
+    target_modules=["q_proj", "v_proj"],
     lora_dropout=0.05,
     bias="none",
-    task_type="CAUSAL_LM",
-    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]
+    task_type="CAUSAL_LM"
 )
-model = get_peft_model(model, peft_config)
-model.print_trainable_parameters()
 
-# Load dataset
-dataset = datasets.load_dataset("json", data_files="Bingaman_training_data.json", field="training_pairs")
-print("First example from dataset:", dataset["train"][0])
-
-# Tokenization
-def tokenize_data(example):
-    formatted_text = f"<s>[INST] {example['input']} [/INST] {example['output']}</s>"
-    inputs = tokenizer(formatted_text, padding="max_length", truncation=True, max_length=4096, return_tensors="pt")
-    input_ids = inputs["input_ids"].squeeze(0).tolist()
-    attention_mask = inputs["attention_mask"].squeeze(0).tolist()
-    labels = input_ids.copy()
-    return {
-        "input_ids": input_ids,
-        "labels": labels,
-        "attention_mask": attention_mask
-    }
-
-tokenized_dataset = dataset["train"].map(tokenize_data, batched=False, remove_columns=dataset["train"].column_names)
-print("First tokenized example:", {k: (type(v), len(v)) for k, v in tokenized_dataset[0].items()})
-
-# Data collator
-def custom_data_collator(features):
-    input_ids = [torch.tensor(f["input_ids"]) for f in features]
-    attention_mask = [torch.tensor(f["attention_mask"]) for f in features]
-    labels = [torch.tensor(f["labels"]) for f in features]
-    return {
-        "input_ids": torch.stack(input_ids),
-        "attention_mask": torch.stack(attention_mask),
-        "labels": torch.stack(labels)
-    }
-
-# Training setup
-accelerator = Accelerator()
-training_args = TrainingArguments(
-    output_dir="./fine_tuned_llama4_healthcare",
-    per_device_train_batch_size=2,
-    gradient_accumulation_steps=8,
-    eval_strategy="steps",
-    eval_steps=10,
-    save_strategy="steps",
-    save_steps=20,
-    save_total_limit=3,
-    num_train_epochs=5,
-    learning_rate=2e-5,
-    weight_decay=0.01,
-    logging_dir="./logs",
-    logging_steps=5,
-    bf16=True,
-    gradient_checkpointing=True,
-    optim="adamw_torch",
-    warmup_steps=50
-)
+# Prepare model for fine-tuning
+model = prepare_model_for_kbit_training(model)
+model = get_peft_model(model, lora_config)
+
+# Training arguments
+training_args = {
+    "output_dir": "./results",
+    "num_train_epochs": 1,
+    "per_device_train_batch_size": 2,
+    "gradient_accumulation_steps": 8,
+    "optim": "adamw_torch",
+    "save_steps": 500,
+    "logging_steps": 100,
+    "learning_rate": 2e-4,
+    "fp16": True,
+    "max_grad_norm": 0.3,
+    "warmup_ratio": 0.03,
+    "lr_scheduler_type": "cosine"
+}
 
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized_dataset,
-    eval_dataset=tokenized_dataset.select(range(min(5, len(tokenized_dataset)))),
-    data_collator=custom_data_collator
+# Initialize trainer
+trainer = accelerator.prepare(
+    datasets.Trainer(
+        model=model,
+        args=datasets.TrainingArguments(**training_args),
+        train_dataset=dataset,
+    )
 )
 
-# Start training
+# Train
 trainer.train()
-model.save_pretrained("./fine_tuned_llama4_healthcare")
-tokenizer.save_pretrained("./fine_tuned_llama4_healthcare")
-print("Training complete. Model and tokenizer saved to ./fine_tuned_llama4_healthcare")
+model.save_pretrained("./fine_tuned_model")
+print("Training completed!")