import os import torch import glob import gc from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments, Trainer, DataCollatorForLanguageModeling, AutoTokenizer, LlamaConfig, AutoConfig ) from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training from datasets import Dataset from huggingface_hub import snapshot_download from tqdm import tqdm import gradio as gr import math from accelerate import Accelerator import subprocess import sys import json import shutil # --- Configuration --- YOUR_HF_USERNAME = "Twelve2five" MODEL_REPO_NAME = "llama-3-8b-rvq-resized" DATASET_REPO_NAME = "podcast-dialogue-rvq-pairs-3items" hf_model_repo_id = f"{YOUR_HF_USERNAME}/{MODEL_REPO_NAME}" hf_dataset_repo_id = f"{YOUR_HF_USERNAME}/{DATASET_REPO_NAME}" # Output directories OUTPUT_TRAINING_DIR = "./llama3-8b-rvq-qlora-finetuned-run" LOGGING_DIR = "./llama3-8b-rvq-qlora-logs-run" local_download_path = "./downloaded_dataset_files" # Training parameters NUM_EPOCHS = 1 BATCH_SIZE_PER_DEVICE = 1 GRAD_ACCUMULATION_STEPS = 64 LEARNING_RATE = 1e-4 WEIGHT_DECAY = 0.01 WARMUP_RATIO = 0.03 LR_SCHEDULER = "cosine" OPTIMIZER = "paged_adamw_8bit" MAX_SEQ_LENGTH = 256 MICRO_BATCH_SIZE = 1 # Multi-GPU configuration accelerator = Accelerator() # Configure environment for multi-GPU os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:32" # Print GPU information print(f"Available GPUs: {torch.cuda.device_count()}") for i in range(torch.cuda.device_count()): print(f"GPU {i}: {torch.cuda.get_device_name(i)} with {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB") def seq2seq_causal_collator(features): """ Collator that concatenates context (input_ids) and target (labels) for Causal LM sequence-to-sequence training. Masks the loss for the context part of the sequence. Pads sequences to the maximum length in the batch. """ batch = {} concatenated_input_ids = [] concatenated_labels = [] max_len = 0 # --- First pass: Concatenate, create masked labels, find max length --- for feature in features: # Dataset transform should provide tensors here input_ids = feature['input_ids'] labels = feature['labels'] # Ensure tensors are 1D (handle potential extra dims if any) if input_ids.dim() > 1: input_ids = input_ids.squeeze() if labels.dim() > 1: labels = labels.squeeze() context_len = input_ids.shape[0] target_len = labels.shape[0] # Concatenate context and target for input combined_ids = torch.cat([input_ids, labels], dim=0) concatenated_input_ids.append(combined_ids) # Create labels: -100 for context, actual labels for target masked_labels = torch.cat([ torch.full((context_len,), -100, dtype=torch.long, device=input_ids.device), labels ], dim=0) concatenated_labels.append(masked_labels) # Track max length for padding if combined_ids.shape[0] > max_len: max_len = combined_ids.shape[0] # --- Second pass: Pad to max length --- padded_input_ids = [] padded_labels = [] input_pad_token_id = 0 label_pad_token_id = -100 for i in range(len(features)): ids = concatenated_input_ids[i] lbls = concatenated_labels[i] padding_len = max_len - ids.shape[0] # Pad on the right side padded_input_ids.append(torch.nn.functional.pad( ids, (0, padding_len), value=input_pad_token_id )) padded_labels.append(torch.nn.functional.pad( lbls, (0, padding_len), value=label_pad_token_id )) # --- Stack and create final batch --- batch['input_ids'] = torch.stack(padded_input_ids) batch['labels'] = torch.stack(padded_labels) # Create attention mask (1 for real tokens, 0 for padding) batch['attention_mask'] = batch['input_ids'].ne(input_pad_token_id).long() return batch def prepare_for_dataset(batch): output = {'input_ids': [], 'labels': []} for item in batch: output['input_ids'].append(item['input_ids'].cpu().tolist()) output['labels'].append(item['labels'].cpu().tolist()) return output def load_model(): print(f"Loading base model architecture from: {hf_model_repo_id}") # Get information about GPU with most free memory gpu_id = 0 # Default to first GPU max_free_memory = 0 for i in range(torch.cuda.device_count()): free_memory = torch.cuda.get_device_properties(i).total_memory - torch.cuda.memory_allocated(i) if free_memory > max_free_memory: max_free_memory = free_memory gpu_id = i print(f"Loading model on GPU {gpu_id} with {max_free_memory / 1e9:.2f}GB free memory") # Configure quantization bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16 ) # Load the model try: # First update transformers to make sure we have latest version subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "transformers"]) # Now try loading with explicit config class to avoid auto-detection issues from transformers import LlamaConfig # Load config first config = LlamaConfig.from_pretrained( hf_model_repo_id, trust_remote_code=True ) # Then load model with explicit config model = AutoModelForCausalLM.from_pretrained( hf_model_repo_id, config=config, quantization_config=bnb_config, device_map="auto", trust_remote_code=True ) log.append(f"Loaded model vocab size: {model.config.vocab_size}") log.append(f"Input embedding shape: {model.get_input_embeddings().weight.shape}") except Exception as e: error_msg = f"Error loading model from Hub: {e}" log.append(error_msg) # Try with a fallback method try: log.append("Attempting alternative loading method...") # Try loading without auto detection model = AutoModelForCausalLM.from_pretrained( hf_model_repo_id, quantization_config=bnb_config, device_map="auto", trust_remote_code=True, torch_dtype=torch.bfloat16, # Add these to help with the loading revision="main", low_cpu_mem_usage=True, ) log.append("Alternative loading successful!") log.append(f"Loaded model vocab size: {model.config.vocab_size}") except Exception as e2: log.append(f"Alternative loading also failed: {e2}") return "\n".join(log) # --- Load Tokenizer --- progress(0.3, desc="Loading tokenizer...") try: log.append("Loading a compatible tokenizer...") # Use the tokenizer from Meta's official Llama models - should be compatible with Llama 3.2 tokenizer_id = "meta-llama/Llama-3-1B" # This is a reliable source for a Llama tokenizer # Try with specified tokenizer first try: tokenizer = AutoTokenizer.from_pretrained( tokenizer_id, use_fast=True, padding_side="right", trust_remote_code=True ) log.append(f"Successfully loaded tokenizer from {tokenizer_id}") except Exception as e: log.append(f"Could not load from {tokenizer_id}: {e}") # Fallback to Llama-2 tokenizer try: tokenizer = AutoTokenizer.from_pretrained( "meta-llama/Llama-2-7b-hf", use_fast=True, padding_side="right" ) log.append("Loaded Llama-2 tokenizer as fallback") except Exception as e2: # If that fails too, try the most basic option from transformers import LlamaTokenizer tokenizer = LlamaTokenizer.from_pretrained( "hf-internal-testing/llama-tokenizer", use_fast=False, padding_side="right" ) log.append("Loaded basic Llama tokenizer from testing repo") # Set pad token if not already set if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token log.append("Set pad_token to eos_token") # Make sure we have necessary special tokens if tokenizer.bos_token is None: tokenizer.bos_token = "" log.append("Set bos_token to ") if tokenizer.eos_token is None: tokenizer.eos_token = "" log.append("Set eos_token to ") log.append(f"Loaded model vocab size: {len(tokenizer)}") except Exception as e: error_msg = f"All attempts to load a tokenizer failed: {e}" log.append(error_msg) return "\n".join(log) # Print information about input embeddings print(f"Input embedding shape: {model.get_input_embeddings().weight.shape}") # Prepare model for k-bit training model = prepare_model_for_kbit_training(model) # Define LoRA configuration lora_config = LoraConfig( r=16, lora_alpha=32, target_modules=[ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], lora_dropout=0.05, bias="none", task_type=TaskType.CAUSAL_LM ) # Apply LoRA to model model = get_peft_model(model, lora_config) model.print_trainable_parameters() return model, tokenizer # Return both model and tokenizer def load_dataset(): # --- Download the dataset repository files --- try: os.makedirs(local_download_path, exist_ok=True) downloaded_repo_root = snapshot_download( repo_id=hf_dataset_repo_id, repo_type="dataset", local_dir=local_download_path, local_dir_use_symlinks=False ) print(f"Dataset repository content downloaded to: {downloaded_repo_root}") except Exception as e: print(f"Error downloading dataset: {e}") return None # --- Load .pt files into a Hugging Face Dataset object --- pairs_dir = os.path.join(downloaded_repo_root, "final_rvq_pairs") all_pair_files = glob.glob(os.path.join(pairs_dir, "*_rvq_pairs.pt")) if not all_pair_files: all_pair_files = glob.glob(os.path.join(downloaded_repo_root, "*_rvq_pairs.pt")) if not all_pair_files: print("No RVQ pair files found!") return None print(f"Found {len(all_pair_files)} RVQ pair files.") # Load data from .pt files into memory all_data_pairs = [] for file_path in tqdm(all_pair_files, desc="Loading pair files"): try: episode_pairs = torch.load(file_path, map_location='cpu') all_data_pairs.extend(episode_pairs) except Exception as e: print(f"Warning: Could not load file {file_path}: {e}") if not all_data_pairs: return None print(f"Loaded {len(all_data_pairs)} training pairs.") # Convert to Hugging Face Dataset chunk_size = 1000 processed_data = {'input_ids': [], 'labels': []} for i in tqdm(range(0, len(all_data_pairs), chunk_size), desc="Preparing data"): batch = all_data_pairs[i:i + chunk_size] prepared_batch = prepare_for_dataset(batch) processed_data['input_ids'].extend(prepared_batch['input_ids']) processed_data['labels'].extend(prepared_batch['labels']) hf_dataset = Dataset.from_dict(processed_data) # Transform to get tensors back hf_dataset.set_transform(lambda batch: { 'input_ids': [torch.tensor(ids, dtype=torch.long) for ids in batch['input_ids']], 'labels': [torch.tensor(lbls, dtype=torch.long) for lbls in batch['labels']] }) # Cleanup del all_data_pairs del processed_data gc.collect() return hf_dataset # Memory cleaning function def clean_memory(): gc.collect() if torch.cuda.is_available(): for i in range(torch.cuda.device_count()): with torch.cuda.device(f'cuda:{i}'): torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() def train_model( hf_username, model_repo_name, dataset_repo_name, epochs=1, batch_size=4, # Increased for A100 grad_accum_steps=4, learning_rate=2e-4, progress=gr.Progress() ): progress(0, desc="Setting up environment...") log = [] # Clean up any existing model files to save space if os.path.exists("./model_files"): try: shutil.rmtree("./model_files") except Exception as e: log.append(f"Warning: Could not remove existing model files: {e}") if os.path.exists("./downloaded_dataset_files"): try: shutil.rmtree("./downloaded_dataset_files") except Exception as e: log.append(f"Warning: Could not remove existing dataset files: {e}") # Print GPU info - using imported torch, not a local variable if torch.cuda.is_available(): log.append(f"Available GPUs: {torch.cuda.device_count()}") for i in range(torch.cuda.device_count()): gpu_name = torch.cuda.get_device_name(i) gpu_memory = torch.cuda.get_device_properties(i).total_memory / (1024**3) log.append(f"GPU {i}: {gpu_name} with {gpu_memory:.2f} GB") # Import required libraries try: from datasets import Dataset from huggingface_hub import snapshot_download # Don't import torch again, since it's already imported import transformers from transformers import AutoModelForCausalLM, AutoTokenizer from transformers import BitsAndBytesConfig, TrainingArguments, Trainer from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training log.append(f"Transformers version: {transformers.__version__}") log.append(f"PyTorch version: {torch.__version__}") except ImportError as e: log.append(f"Error importing libraries: {e}") return "\n".join(log) # --- Configuration --- progress(0.05, desc="Setting up configuration...") hf_model_repo_id = f"{hf_username}/{model_repo_name}" hf_dataset_repo_id = f"{hf_username}/{dataset_repo_name}" log.append(f"Model repo: {hf_model_repo_id}") log.append(f"Dataset repo: {hf_dataset_repo_id}") # Check if running on multiple GPUs n_gpus = torch.cuda.device_count() log.append(f"Number of GPUs available: {n_gpus}") # --- DeepSpeed Configuration --- # Create DeepSpeed config file progress(0.1, desc="Setting up DeepSpeed configuration...") # Create a simpler config since we have plenty of memory on A100 ds_config = { "bf16": { "enabled": "auto" }, "zero_optimization": { "stage": 1, # Lower stage is fine for A100-80GB "contiguous_gradients": True, "overlap_comm": True }, "gradient_accumulation_steps": grad_accum_steps, "gradient_clipping": 1.0, "train_batch_size": batch_size * grad_accum_steps * max(1, n_gpus) } ds_config_path = "ds_config.json" with open(ds_config_path, "w") as f: json.dump(ds_config, f, indent=4) log.append("DeepSpeed configuration created successfully") # --- Download and Load Model --- progress(0.15, desc="Downloading model...") try: # Download model files local_model_path = "./model_files" snapshot_download( repo_id=hf_model_repo_id, local_dir=local_model_path, use_auth_token=False, resume_download=True ) log.append(f"Model files downloaded to {local_model_path}") # Check and fix the model config if needed config_path = os.path.join(local_model_path, "config.json") if os.path.exists(config_path): with open(config_path, 'r') as f: config_data = json.load(f) # Fix the rope_scaling configuration if 'rope_scaling' in config_data: if not isinstance(config_data['rope_scaling'], dict): config_data['rope_scaling'] = {"type": "linear", "factor": 2.0} elif 'rope_type' in config_data['rope_scaling']: # Convert complex rope_scaling to the simple format expected rope_factor = config_data['rope_scaling'].get('factor', 2.0) config_data['rope_scaling'] = {"type": "linear", "factor": rope_factor} # Write the updated config back with open(config_path, 'w') as f: json.dump(config_data, f, indent=2) log.append("Updated model configuration for rope_scaling") # Create a bnb configuration for loading the model in 4-bit progress(0.25, desc="Loading model...") bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=False ) # Load the model with fixed configuration model = AutoModelForCausalLM.from_pretrained( local_model_path, quantization_config=bnb_config, device_map="auto", use_cache=False, # Needed for gradient checkpointing trust_remote_code=True, # Following reference code torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16, ) # --- Load Tokenizer (from a compatible model) --- # Following the pattern from reference code progress(0.3, desc="Loading tokenizer...") # Try to load a compatible tokenizer try: # First try loading from standard Llama 3 model tokenizer = AutoTokenizer.from_pretrained( "meta-llama/Llama-3-8B", # Using standard Llama 3 tokenizer padding_side="right", use_fast=True, trust_remote_code=True ) log.append("Loaded tokenizer from meta-llama/Llama-3-8B") except Exception as e1: log.append(f"Couldn't load Llama-3 tokenizer: {e1}") try: # Fallback to Llama 2 tokenizer = AutoTokenizer.from_pretrained( "meta-llama/Llama-2-7b-hf", padding_side="right", use_fast=True ) log.append("Loaded Llama-2 tokenizer as fallback") except Exception as e2: log.append(f"Couldn't load Llama-2 tokenizer: {e2}") # Final fallback from transformers import LlamaTokenizer tokenizer = LlamaTokenizer.from_pretrained( "hf-internal-testing/llama-tokenizer", padding_side="right" ) log.append("Loaded testing Llama tokenizer as final fallback") # Set pad token and ensure it's usable if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token log.append(f"Loaded model vocab size: {model.config.vocab_size}") log.append(f"Input embedding shape: {model.get_input_embeddings().weight.shape}") # --- QLoRA Preparation --- progress(0.35, desc="Preparing model for k-bit training...") model = prepare_model_for_kbit_training(model) log.append("Model prepared for k-bit training") # Define LoRA configuration # Based on your reference code lora_config = LoraConfig( task_type=TaskType.CAUSAL_LM, r=16, # Rank lora_alpha=32, # Alpha parameter lora_dropout=0.05, # Dropout probability bias="none", # Bias type target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] ) # Apply LoRA to model progress(0.4, desc="Applying LoRA to model...") model_to_train = get_peft_model(model, lora_config) log.append("LoRA applied to model") # Cleanup to free up memory gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e: error_msg = f"Error preparing model for training: {str(e)}" log.append(error_msg) return "\n".join(log) # --- Download and Process Dataset --- progress(0.4, desc="Downloading dataset...") try: dataset_path = "./downloaded_dataset_files" snapshot_download( repo_id=hf_dataset_repo_id, local_dir=dataset_path, use_auth_token=False, resume_download=True ) log.append(f"Dataset repository content downloaded to: {dataset_path}") # Load dataset from PT files progress(0.5, desc="Processing dataset...") # Load RVQ pairs pair_files = glob.glob(f"{dataset_path}/*_rvq_pairs.pt") log.append(f"Found {len(pair_files)} RVQ pair files.") all_pairs = [] for file in pair_files: pairs = torch.load(file) all_pairs.extend(pairs) log.append(f"Loaded a total of {len(all_pairs)} training pairs into memory.") # Process pairs into a format suitable for training all_texts = [] for pair in all_pairs: # Create instruction format if isinstance(pair, dict): instruction = pair.get("instruction", "") input_text = pair.get("input", "") output = pair.get("output", "") # ALPACA format if instruction and input_text: text = f"### Instruction: {instruction}\n### Input: {input_text}\n### Response: {output}" elif instruction: text = f"### Instruction: {instruction}\n### Response: {output}" else: text = output else: # Simple prompt-completion format if isinstance(pair, tuple) and len(pair) == 2: prompt, completion = pair text = f"{prompt}{completion}" else: text = str(pair) all_texts.append({"text": text}) # Create HF dataset train_dataset = Dataset.from_list(all_texts) # Function to tokenize the dataset def tokenize_function(examples): return tokenizer( examples["text"], padding=False, truncation=True, max_length=2048, return_tensors=None, ) # Tokenize the dataset tokenized_dataset = train_dataset.map( tokenize_function, batched=True, remove_columns=["text"], desc="Tokenizing dataset", ) train_dataset = tokenized_dataset # Data collator from transformers import DataCollatorForLanguageModeling data_collator = DataCollatorForLanguageModeling( tokenizer=tokenizer, mlm=False ) except Exception as e: error_msg = f"Error loading dataset: {str(e)}" log.append(error_msg) return "\n".join(log) # --- Training Arguments --- progress(0.75, desc="Setting up training arguments...") output_dir = f"./results_{model_repo_name}" os.makedirs(output_dir, exist_ok=True) # Optimize settings for A100 training_args = TrainingArguments( output_dir=output_dir, num_train_epochs=float(epochs), per_device_train_batch_size=batch_size, gradient_accumulation_steps=grad_accum_steps, learning_rate=learning_rate, weight_decay=0.01, logging_dir=f"{output_dir}/logs", logging_steps=10, save_steps=100, save_total_limit=3, remove_unused_columns=False, push_to_hub=False, disable_tqdm=False, warmup_ratio=0.03, lr_scheduler_type="cosine", report_to="tensorboard", bf16=True if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else False, gradient_checkpointing=True, # Still useful for efficiency gradient_checkpointing_kwargs={'use_reentrant': False}, ddp_find_unused_parameters=False, deepspeed=ds_config_path if n_gpus > 1 else None, # Only use DeepSpeed for multi-GPU ) # --- Initialize Trainer --- progress(0.8, desc="Initializing trainer...") trainer = Trainer( model=model_to_train, args=training_args, train_dataset=train_dataset, data_collator=data_collator, ) log.append("Trainer initialized for training.") # --- Start Training --- # Clear cache before starting gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() try: progress(0.85, desc="Starting training...") log.append("Starting training...") train_result = trainer.train() progress(0.95, desc="Saving model...") # Save final model (adapter weights) and training state final_save_path = os.path.join(training_args.output_dir, "final_checkpoint") log.append(f"Saving final model checkpoint to {final_save_path}...") trainer.save_model(final_save_path) trainer.save_state() # Log metrics metrics = train_result.metrics trainer.log_metrics("train", metrics) trainer.save_metrics("train", metrics) for key, value in metrics.items(): log.append(f"{key}: {value}") except Exception as e: error_msg = f"An error occurred during training: {e}" log.append(error_msg) return "\n".join(log) progress(1.0, desc="Training complete!") log.append("Training process complete.") return "\n".join(log) # Define the Gradio interface def create_interface(): with gr.Blocks(title="Llama 3.2 1B RVQ Fine-tuning") as demo: gr.Markdown("# Llama 3.2 1B RVQ LoRA Fine-tuning") gr.Markdown("Fine-tune a Llama 3.2 1B model with RVQ token embeddings using LoRA") with gr.Row(): with gr.Column(): hf_username = gr.Textbox(label="HuggingFace Username", value="Twelve2five") model_repo = gr.Textbox(label="Model Repository Name", value="llama-3.2-1b-rvq") dataset_repo = gr.Textbox(label="Dataset Repository Name", value="podcast-dialogue-rvq-pairs-3items") with gr.Column(): epochs = gr.Number(label="Number of Epochs", value=3, minimum=1, maximum=10) batch_size = gr.Number(label="Batch Size per Device", value=4, minimum=1, maximum=16) grad_accum = gr.Number(label="Gradient Accumulation Steps", value=2, minimum=1, maximum=16) lr = gr.Number(label="Learning Rate", value=2e-4) start_btn = gr.Button("Start Training") output = gr.Textbox(label="Training Log", lines=20) start_btn.click( fn=train_model, inputs=[hf_username, model_repo, dataset_repo, epochs, batch_size, grad_accum, lr], outputs=output ) return demo # Create and launch the interface demo = create_interface() if __name__ == "__main__": demo.launch()