""" Actor and Critic models for offline RL with QLoRA. This file contains the Actor and Critic model implementations using QLoRA for efficient finetuning of LLMs. """ import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig from peft import LoraConfig, get_peft_model, PeftModel, TaskType import platform import re def get_target_modules_for_model(model_id, model): """ Get the appropriate target modules for LoRA based on the model architecture. Args: model_id: The model identifier string model: The loaded model Returns: List of target module names """ # Check model architecture if "llama" in model_id.lower() or "mistral" in model_id.lower(): # Llama/Mistral models return ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] elif "gpt-j" in model_id.lower(): # GPT-J models return ["q_proj", "k_proj", "v_proj", "out_proj", "fc_in", "fc_out"] elif "gpt-neox" in model_id.lower(): # GPT-NeoX models return ["query_key_value", "dense", "dense_h_to_4h", "dense_4h_to_h"] elif "gpt2" in model_id.lower(): # GPT-2 models return ["c_attn", "c_proj", "c_fc", "c_proj"] elif hasattr(model, "config") and hasattr(model.config, "architectures"): # Try to infer from model architecture arch = model.config.architectures[0] if model.config.architectures else "" if "GPT2" in arch: return ["c_attn", "c_proj", "c_fc", "c_proj"] elif "Llama" in arch or "Mistral" in arch: return ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] elif "GPTNeoX" in arch: return ["query_key_value", "dense", "dense_h_to_4h", "dense_4h_to_h"] # Examine module names in the model module_names = [] for name, _ in model.named_modules(): if any(substr in name for substr in ["attn", "mlp", "attention"]): parts = name.split(".") if len(parts) > 1: module_names.append(parts[-1]) # Extract common module names if module_names: # Try to find attention and MLP layers attn_patterns = ["attn", "attention", "self", "q", "k", "v", "query", "key", "value"] mlp_patterns = ["mlp", "feed_forward", "fc", "dense", "linear", "ffn"] attn_modules = [name for name in module_names if any(p in name.lower() for p in attn_patterns)] mlp_modules = [name for name in module_names if any(p in name.lower() for p in mlp_patterns)] if attn_modules or mlp_modules: return list(set(attn_modules + mlp_modules)) # Default to commonly used module names as a fallback print(f"Warning: Could not determine target modules for {model_id}. Using default modules.") return ["query", "key", "value", "dense"] class LLMActorLora: """Actor model with QLoRA for LLMs.""" def __init__(self, device, model_id="meta-llama/Llama-3-8B-Instruct", lora_r=8, disable_quantization=False): """ Initialize the actor model with QLoRA. Args: device: Device to run the model on model_id: HuggingFace model ID lora_r: LoRA rank parameter disable_quantization: If True, disable 4-bit quantization (useful for Mac/CPU) """ self.device = device self.model_id = model_id # Check if we're on a Mac - often has issues with 4-bit quantization is_mac = platform.system() == 'Darwin' running_on_cpu = 'cpu' in str(device).lower() # Set up quantization config if not disabled if disable_quantization or (is_mac and running_on_cpu): if disable_quantization: print(f"4-bit quantization disabled by user request") else: print(f"4-bit quantization automatically disabled for Mac/CPU") # Without quantization self.model = AutoModelForCausalLM.from_pretrained( model_id, torch_dtype=torch.float32 if running_on_cpu else torch.bfloat16, ).to(device) else: # Load model in 4-bit mode with QLoRA self.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 base model in 4-bit print(f"Loading {model_id} with QLoRA (4-bit quantization + LoRA)") self.model = AutoModelForCausalLM.from_pretrained( model_id, quantization_config=self.bnb_config, device_map="auto" ) # Get the right target modules for this model architecture target_modules = get_target_modules_for_model(model_id, self.model) print(f"Using target modules for LoRA: {target_modules}") # Apply LoRA config self.lora_config = LoraConfig( r=lora_r, lora_alpha=2 * lora_r, target_modules=target_modules, bias="none", task_type=TaskType.CAUSAL_LM, lora_dropout=0.05, ) # Create LoRA model self.model = get_peft_model(self.model, self.lora_config) self.model.print_trainable_parameters() def parameters(self): """Return the model parameters.""" return self.model.parameters() def forward(self, input_ids, attention_mask=None): """ Forward pass through the model. Args: input_ids: Tokenized input IDs attention_mask: Attention mask Returns: Model outputs """ return self.model(input_ids, attention_mask=attention_mask) def get_log_probs(self, input_ids, action_ids, attention_mask=None): """ Calculate log probabilities for given actions. Args: input_ids: Tokenized input IDs [batch_size, seq_len] action_ids: Tokenized action IDs [batch_size, act_len] attention_mask: Attention mask [batch_size, seq_len] Returns: log_probs: Log probabilities of actions [batch_size] entropy: Entropy of the policy [batch_size] """ outputs = self.model(input_ids, attention_mask=attention_mask) logits = outputs.logits # Get the logits for the last token in each sequence last_token_logits = logits[:, -1, :] # [batch_size, vocab_size] # Get log probabilities for the first token of each action log_probs = F.log_softmax(last_token_logits, dim=-1) first_action_tokens = action_ids[:, 0] # [batch_size] selected_log_probs = log_probs.gather(1, first_action_tokens.unsqueeze(-1)).squeeze(-1) # Calculate entropy probs = torch.exp(log_probs) entropy = -(probs * log_probs).sum(dim=-1) return selected_log_probs, entropy def generate(self, input_ids, attention_mask=None, **kwargs): """ Generate text from the model. Args: input_ids: Tokenized input IDs attention_mask: Attention mask kwargs: Additional generation arguments Returns: Generated token IDs """ return self.model.generate(input_ids, attention_mask=attention_mask, **kwargs) def save_pretrained(self, path): """ Save the model to the given path. Args: path: Path to save the model to """ self.model.save_pretrained(path) def load_pretrained(self, path): """ Load the model from the given path. Args: path: Path to load the model from """ self.model = PeftModel.from_pretrained(self.model, path) class LLMCriticLora: """Critic (value function) model with QLoRA for LLMs.""" def __init__(self, device, model_id="meta-llama/Llama-3-8B-Instruct", lora_r=8, disable_quantization=False): """ Initialize the critic model with QLoRA. Args: device: Device to run the model on model_id: HuggingFace model ID lora_r: LoRA rank parameter disable_quantization: If True, disable 4-bit quantization (useful for Mac/CPU) """ self.device = device self.model_id = model_id # Check if we're on a Mac - often has issues with 4-bit quantization is_mac = platform.system() == 'Darwin' running_on_cpu = 'cpu' in str(device).lower() # Set up quantization config if not disabled if disable_quantization or (is_mac and running_on_cpu): if disable_quantization: print(f"Critic: 4-bit quantization disabled by user request") else: print(f"Critic: 4-bit quantization automatically disabled for Mac/CPU") # Without quantization self.model = AutoModelForCausalLM.from_pretrained( model_id, torch_dtype=torch.float32 if running_on_cpu else torch.bfloat16, ).to(device) else: # Load model in 4-bit mode with QLoRA self.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 base model in 4-bit print(f"Loading critic {model_id} with QLoRA (4-bit quantization + LoRA)") self.model = AutoModelForCausalLM.from_pretrained( model_id, quantization_config=self.bnb_config, device_map="auto" ) # Get the right target modules for this model architecture target_modules = get_target_modules_for_model(model_id, self.model) print(f"Critic: Using target modules for LoRA: {target_modules}") # Apply LoRA config self.lora_config = LoraConfig( r=lora_r, lora_alpha=2 * lora_r, target_modules=target_modules, bias="none", task_type=TaskType.CAUSAL_LM, lora_dropout=0.05, ) # Create LoRA model self.model = get_peft_model(self.model, self.lora_config) # Add a value head on top of the LLM hidden_size = self.model.config.hidden_size self.value_head = nn.Sequential( nn.Linear(hidden_size, hidden_size // 2), nn.ReLU(), nn.Linear(hidden_size // 2, 1) ).to(device) self.model.print_trainable_parameters() def parameters(self): """Return all trainable parameters.""" return list(self.model.parameters()) + list(self.value_head.parameters()) def forward(self, input_ids, attention_mask=None): """ Forward pass to compute the value function. Args: input_ids: Tokenized input IDs attention_mask: Attention mask Returns: Value predictions [batch_size] """ # Get the hidden states from the LLM outputs = self.model(input_ids, attention_mask=attention_mask, output_hidden_states=True) hidden_states = outputs.hidden_states[-1] # Use the last layer's hidden states # Extract the last token's hidden state for each sequence batch_size = hidden_states.shape[0] if attention_mask is not None: # Find the position of the last non-padding token last_token_positions = attention_mask.sum(dim=1) - 1 # [batch_size] last_token_hidden = hidden_states[torch.arange(batch_size), last_token_positions] else: # Just use the last token last_token_hidden = hidden_states[:, -1] # Pass through the value head to get values # Ensure dtype matches value_head if last_token_hidden.dtype != next(self.value_head.parameters()).dtype: last_token_hidden = last_token_hidden.to(next(self.value_head.parameters()).dtype) values = self.value_head(last_token_hidden) return values def save_pretrained(self, path): """ Save the model to the given path. Args: path: Path to save the model to """ self.model.save_pretrained(f"{path}/lora") torch.save(self.value_head.state_dict(), f"{path}/value_head.pt") def load_pretrained(self, path): """ Load the model from the given path. Args: path: Path to load the model from """ self.model = PeftModel.from_pretrained(self.model, f"{path}/lora") self.value_head.load_state_dict(torch.load(f"{path}/value_head.pt"))