Function extract_lora_from_diff

tools/extract/lora_extractor.py:287–395

Extract LoRA weights from model difference Args: source_weights: Source model weights target_weights: Target model weights rank: LoRA rank diff_only: If True, save all weights as direct diff without LoRA decomposition Returns: LoRA weights dicti

(source_weights: Dict[str, torch.Tensor], target_weights: Dict[str, torch.Tensor], rank: int = 16, diff_only: bool = False)

def extract_lora_from_diff(source_weights: Dict[str, torch.Tensor], target_weights: Dict[str, torch.Tensor], rank: int = 16, diff_only: bool = False) -> Dict[str, torch.Tensor]:
    """
    Extract LoRA weights from model difference

    Args:
        source_weights: Source model weights
        target_weights: Target model weights
        rank: LoRA rank
        diff_only: If True, save all weights as direct diff without LoRA decomposition

    Returns:
        LoRA weights dictionary
    """
    print("Starting LoRA weight extraction...")
    if diff_only:
        print("Mode: Direct diff only (no LoRA decomposition)")
    else:
        print(f"Mode: Smart extraction - rank: {rank}")
    print(f"Source model weight count: {len(source_weights)}")
    print(f"Target model weight count: {len(target_weights)}")

    lora_weights = {}
    processed_count = 0
    diff_count = 0
    lora_count = 0
    similar_count = 0
    skipped_count = 0
    fail_count = 0

    # Find common keys between two models
    common_keys = set(source_weights.keys()) & set(target_weights.keys())
    source_only_keys = set(source_weights.keys()) - set(target_weights.keys())
    target_only_keys = set(target_weights.keys()) - set(source_weights.keys())

    if source_only_keys:
        print(f"Warning: Source model exclusive weight keys ({len(source_only_keys)} keys): {list(source_only_keys)[:5]}...")
    if target_only_keys:
        print(f"Warning: Target model exclusive weight keys ({len(target_only_keys)} keys): {list(target_only_keys)[:5]}...")

    print(f"Common weight keys count: {len(common_keys)}")

    # Process common keys, extract LoRA weights
    common_keys_sorted = sorted(common_keys)
    pbar = tqdm(common_keys_sorted, desc="Extracting LoRA weights", unit="layer")

    for key in pbar:
        source_tensor = source_weights[key]
        target_tensor = target_weights[key]

        # Update progress bar description
        short_key = key.split(".")[-2:] if "." in key else [key]
        pbar.set_postfix_str(f"Processing: {'.'.join(short_key)}")

        # Compute weight difference
        diff = _compute_weight_diff(source_tensor, target_tensor, key)

        if diff is None:
            # No change or shape mismatch