# 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

```yaml
# 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)**:
   ```yaml
   lora_target_modules: ["q_proj", "v_proj"]
   ```
   - Adapts Query and Value projections only
   - ~0.1% trainable parameters
   - Good for initial experiments

2. **Extended Attention**:
   ```yaml
   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**:
   ```yaml
   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
```bash
python train_network_unified.py \
    --config config/OCID-VLG/legrconv_llava_r50_lora.yaml
```

### Experiment with Different Ranks
```bash
# 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
```python
# In your training script
model.save_lora_weights("path/to/lora_checkpoint")
```
This saves only the LoRA adapter weights (~10-100MB).

### Load LoRA Weights
```python
# 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
```python
# 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