146 lines
4.3 KiB
Python
146 lines
4.3 KiB
Python
"""
|
|
LoRA (Low-Rank Adaptation) utilities for MiniMindVLM
|
|
Enables parameter-efficient fine-tuning with <2% trainable parameters
|
|
"""
|
|
import torch
|
|
from peft import LoraConfig, get_peft_model, TaskType, PeftModel
|
|
from typing import Optional
|
|
|
|
|
|
def get_lora_config(
|
|
r: int = 16,
|
|
lora_alpha: int = 32,
|
|
lora_dropout: float = 0.05,
|
|
target_modules: Optional[list] = None,
|
|
bias: str = "none",
|
|
) -> LoraConfig:
|
|
"""
|
|
Get LoRA configuration for MiniMindVLM
|
|
|
|
Args:
|
|
r: Rank of LoRA update matrices (higher = more parameters)
|
|
lora_alpha: Scaling factor for LoRA updates
|
|
lora_dropout: Dropout rate for LoRA layers
|
|
target_modules: Which modules to apply LoRA to. If None, uses default for MiniMind
|
|
bias: How to handle bias parameters ("none", "all", "lora_only")
|
|
|
|
Returns:
|
|
LoraConfig object ready to apply to model
|
|
|
|
Typical memory savings:
|
|
- r=8: ~0.5% trainable params, 99.5% frozen
|
|
- r=16: ~1.0% trainable params, 99.0% frozen (recommended)
|
|
- r=32: ~2.0% trainable params, 98.0% frozen
|
|
"""
|
|
|
|
# Default target modules for MiniMind architecture
|
|
# These are the attention projection layers
|
|
if target_modules is None:
|
|
target_modules = [
|
|
"q_proj", # Query projection in attention
|
|
"k_proj", # Key projection in attention
|
|
"v_proj", # Value projection in attention
|
|
"o_proj", # Output projection in attention
|
|
]
|
|
|
|
config = LoraConfig(
|
|
task_type=TaskType.CAUSAL_LM,
|
|
inference_mode=False,
|
|
r=r,
|
|
lora_alpha=lora_alpha,
|
|
lora_dropout=lora_dropout,
|
|
target_modules=target_modules,
|
|
bias=bias,
|
|
)
|
|
|
|
return config
|
|
|
|
|
|
def apply_lora(model, lora_config: LoraConfig, verbose: bool = True):
|
|
"""
|
|
Apply LoRA adapters to a MiniMindVLM model
|
|
|
|
Args:
|
|
model: MiniMindVLM model instance
|
|
lora_config: LoRA configuration from get_lora_config()
|
|
verbose: Print trainable parameter statistics
|
|
|
|
Returns:
|
|
Model with LoRA adapters applied
|
|
"""
|
|
# Freeze vision encoder (it's already frozen in MiniMindVLM, but be explicit)
|
|
if hasattr(model, 'vision_encoder') and model.vision_encoder is not None:
|
|
for param in model.vision_encoder.parameters():
|
|
param.requires_grad = False
|
|
|
|
# Apply LoRA
|
|
model = get_peft_model(model, lora_config)
|
|
|
|
if verbose:
|
|
model.print_trainable_parameters()
|
|
|
|
# Calculate memory savings
|
|
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
|
total = sum(p.numel() for p in model.parameters())
|
|
|
|
print(f"\n📊 LoRA Statistics:")
|
|
print(f" Total parameters: {total / 1e6:.2f}M")
|
|
print(f" Trainable parameters: {trainable / 1e6:.2f}M ({trainable / total * 100:.2f}%)")
|
|
print(f" Frozen parameters: {(total - trainable) / 1e6:.2f}M ({(total - trainable) / total * 100:.2f}%)")
|
|
print(f" Memory savings: ~{100 - (trainable / total * 100):.1f}%")
|
|
|
|
return model
|
|
|
|
|
|
def load_lora_model(base_model, lora_adapter_path: str):
|
|
"""
|
|
Load a model with LoRA adapters
|
|
|
|
Args:
|
|
base_model: Base MiniMindVLM model (without LoRA)
|
|
lora_adapter_path: Path to saved LoRA adapters
|
|
|
|
Returns:
|
|
Model with LoRA adapters loaded
|
|
"""
|
|
model = PeftModel.from_pretrained(base_model, lora_adapter_path)
|
|
return model
|
|
|
|
|
|
def merge_lora_weights(model, save_path: Optional[str] = None):
|
|
"""
|
|
Merge LoRA weights into base model for deployment
|
|
After merging, the model is a standard model without adapters
|
|
|
|
Args:
|
|
model: Model with LoRA adapters
|
|
save_path: If provided, save merged model to this path
|
|
|
|
Returns:
|
|
Merged model (standard model, no adapters)
|
|
"""
|
|
# Merge LoRA weights into base model
|
|
model = model.merge_and_unload()
|
|
|
|
if save_path:
|
|
model.save_pretrained(save_path)
|
|
print(f"✅ Merged model saved to {save_path}")
|
|
|
|
return model
|
|
|
|
|
|
# Example usage:
|
|
if __name__ == "__main__":
|
|
from model.model_vlm import MiniMindVLM, VLMConfig
|
|
|
|
# Create model
|
|
config = VLMConfig()
|
|
model = MiniMindVLM(config)
|
|
|
|
# Apply LoRA
|
|
lora_config = get_lora_config(r=16)
|
|
model = apply_lora(model, lora_config)
|
|
|
|
# Now model is ready for training with LoRA
|
|
# Only ~1% of parameters will be updated during training
|