[tab] test_lora_ckpt.py (python)
import os
import shutil
import tempfile
import json
import torch
import torch.nn as nn
from pathlib import Path
from unittest.mock import MagicMock, patch
from argparse import Namespace

# Mock dependencies that might be hard to import or run in isolation
import sys
sys.modules["miles.utils.tracking_utils"] = MagicMock()
sys.modules["miles.utils.train_dump_utils"] = MagicMock()
sys.modules["ring_flash_attn"] = MagicMock()

# Import the modified modules
from miles.models.peft import LoRAConfig, apply_lora
from miles.backends.fsdp_utils import checkpoint

def test_hf_lora_checkpointing():
    # Setup
    temp_dir = tempfile.mkdtemp()
    save_path = Path(temp_dir) / "checkpoints"
    
    try:
        # Create a simple model with LoRA
        model = nn.Sequential(
            nn.Linear(10, 10),
            nn.Linear(10, 10)
        )
        # Mock named_modules to return q_proj/v_proj names for LoRA application
        # But apply_lora iterates named_modules. Let's just manually wrap or use a model structure that matches.
        # Simpler: just create a dummy actor with args and a model that has "lora_" keys in state_dict
        
        class DummyModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.layer = nn.Linear(10, 10)
                self.lora_A = nn.Parameter(torch.randn(4, 10))
                self.lora_B = nn.Parameter(torch.randn(10, 4))
                
        model = DummyModel()
        
        # Mock Actor
        actor = MagicMock()
        actor.model = model
        actor.global_step = 0
        actor.micro_step = 0
        actor.args = Namespace(
            save=str(save_path),
            use_lora=True,
            lora_rank=4,
            lora_alpha=8,
            lora_dropout=0.1,
            lora_target_modules=["layer"]
        )
        
        # Mock Distributed
        with patch("torch.distributed.get_rank", return_value=0), \
             patch("torch.distributed.get_world_size", return_value=1), \
             patch("torch.distributed.barrier"), \
             patch("torch.distributed.checkpoint.save"), \
             patch("torch.distributed.checkpoint.state_dict.get_state_dict", return_value=(model.state_dict(), None)):
            
            # Run save
            checkpoint.save(actor, iteration=0)
            
            # Verify Output
            iter_dir = save_path / "iter_0000001"
            print(f"Checking {iter_dir}")
            
            # Check adapter_config.json
            config_path = iter_dir / "adapter_config.json"
            assert config_path.exists(), "adapter_config.json not found"
            with open(config_path) as f:
                config = json.load(f)
                print("Config:", config)
                assert config["r"] == 4
                assert config["lora_alpha"] == 8
                assert config["peft_type"] == "LORA"
                
            # Check adapter_model.safetensors
            weights_path = iter_dir / "adapter_model.safetensors"
            assert weights_path.exists(), "adapter_model.safetensors not found"
            
            import safetensors.torch
            weights = safetensors.torch.load_file(weights_path)
            print("Weights keys:", weights.keys())
            assert "lora_A" in weights
            assert "lora_B" in weights
            assert "layer.weight" not in weights # Should only have LoRA keys
            
            print("Test PASSED!")
            
    finally:
        shutil.rmtree(temp_dir)

if __name__ == "__main__":
    test_hf_lora_checkpointing()