File size: 13,425 Bytes
2a347f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 |
"""
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")) |