docs/llava_lora_guide.md

LLaVA LoRA Fine-tuning Guide

Overview

This guide explains how to use LoRA (Low-Rank Adaptation) to fine-tune the LLaVA model for the robotic grasping task. LoRA enables efficient adaptation of the language model while keeping most parameters frozen.

Key Benefits

• Memory Efficient: Only ~0.1-1% of parameters need updating
• Fast Training: Smaller optimizer states and gradients
• Preserve Base Knowledge: Original LLaVA weights remain unchanged
• Task Adaptation: Model learns task-specific language-vision patterns

Implementation Details

Model Architecture

The legrconv4x_llava_lora.py model:

• Inherits from the frozen v2 model for backward compatibility
• Uses HuggingFace PEFT for LoRA implementation
• Uses TaskType.CAUSAL_LM (correct for LLaVA)
• PEFT automatically manages trainable parameters
• Properly handles train/eval modes for LoRA dropout
• Provides methods for saving/loading LoRA weights only

Configuration Options

# LoRA Configuration in YAML
use_lora: True                      # Enable LoRA
lora_rank: 16                       # Low-rank dimension (typical: 8-64)
lora_alpha: 32                      # Scaling factor (typical: 2x rank)
lora_dropout: 0.05                  # Dropout for regularization
lora_target_modules: ["q_proj", "v_proj"]  # Which layers to adapt
lora_lr_multiplier: 1.0             # LR multiplier for LoRA params

Target Module Options

1. Minimal (Recommended Start):

   lora_target_modules: ["q_proj", "v_proj"]
   

   • Adapts Query and Value projections only
   • ~0.1% trainable parameters
   • Good for initial experiments

2. Extended Attention:

   lora_target_modules: ["q_proj", "v_proj", "k_proj", "o_proj"]
   

   • All attention projection layers
   • ~0.2% trainable parameters
   • Better capacity for complex adaptations

3. Full Model:

   lora_target_modules: ["q_proj", "v_proj", "k_proj", "o_proj", 
                        "gate_proj", "up_proj", "down_proj"]
   

   • Attention + Feed-forward layers
   • ~0.5% trainable parameters
   • Maximum adaptation capacity

Training Commands

Basic Training

python train_network_unified.py \
    --config config/OCID-VLG/legrconv_llava_r50_lora.yaml

Experiment with Different Ranks

# Small rank (r=8)
python train_network_unified.py \
    --config config/OCID-VLG/legrconv_llava_r50_lora.yaml \
    --opts lora_rank 8 lora_alpha 16

# Large rank (r=32)  
python train_network_unified.py \
    --config config/OCID-VLG/legrconv_llava_r50_lora.yaml \
    --opts lora_rank 32 lora_alpha 64

Recommended Training Strategy

1. Start Small: Begin with rank=8 or 16 and minimal modules
2. Monitor Performance: Track validation metrics and training efficiency
3. Scale Up If Needed: Increase rank or add modules if performance plateaus
4. Learning Rate: LoRA often benefits from slightly higher LR than full fine-tuning

Memory and Performance

Configuration	Trainable Params	GPU Memory	Training Speed
Frozen (v2)	~10M (heads only)	~16GB	Fastest
LoRA r=8 minimal	~20M	~18GB	Fast
LoRA r=16 minimal	~30M	~20GB	Fast
LoRA r=16 extended	~50M	~22GB	Moderate
LoRA r=32 full	~150M	~28GB	Moderate

Saving and Loading

Save LoRA Weights Only

# In your training script
model.save_lora_weights("path/to/lora_checkpoint")

This saves only the LoRA adapter weights (~10-100MB).

Load LoRA Weights

# Load base model + LoRA
model = build_net(args)  # This creates the PEFT model
model.load_lora_weights("path/to/lora_checkpoint")  # Uses load_adapter internally

Merge and Export

# Merge LoRA into base model for deployment
model.merge_and_unload()
torch.save(model.state_dict(), "merged_model.pth")

Troubleshooting

Out of Memory

• Reduce batch size
• Use smaller LoRA rank
• Use gradient accumulation
• Reduce number of target modules

Poor Performance

• Increase LoRA rank
• Add more target modules
• Increase learning rate for LoRA
• Train for more epochs

Unstable Training

• Reduce learning rate
• Increase LoRA dropout
• Use warmup scheduler
• Check gradient clipping

Best Practices

1. Always start with frozen baseline to establish performance floor
2. Use validation set to select optimal rank and modules
3. Save LoRA weights separately for easy experimentation
4. Monitor trainable parameters - PEFT will print this at initialization
5. Use AdamW optimizer with weight decay for better regularization