universal_distill_v4_2_0_aurora_fast.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Universal SVD-LoRA Distillation v4.2.0 (Cross-Arch Innovations) — Plug-and-Play Truncated SVD (Low-VRAM)

Key changes vs v3.7.16 (SVD-Optimized):
----------------------
0. Adds AURORA-SVD (Adaptive Unrolled Residual-Order Range Augmentation SVD) as a distillation-centric low-rank backend.
0. Fixes DRF-SVD backend definition/indentation so --svd-mode drf works as intended.
0. Randomized SVD now uses QR-normalized subspace iteration (Halko et al., Alg. 4.4) for better numerical stability.
0. Adaptive rank selection is now measured against ||Δ||_F^2 and grows rank geometrically instead of always computing SVD at max_rank.
1. Adds SVD-LLM style full-covariance whitening (Cholesky) as a calibration mode ("cov")
   to replace diagonal-only RMS scaling when requested.
2. Restores v3.7.15-stable TIES merge semantics by default and applies the TIES Trim
   step to single-delta cases by default (use --no-ties-trim-single to disable).
3. Preserves the low-VRAM streaming/offload design (cpu/disk offload + prefetch).

Output:
-------
Writes a PEFT LoRA adapter (adapter_model.safetensors + adapter_config.json).

License: Apache 2.0
"""

import os
import re
import json
import math
import time
import gc
import argparse
import warnings
import threading
import queue
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple, Any, Set, Union, Iterator
from enum import Enum
from contextlib import contextmanager
from pathlib import Path
import tempfile
import shutil

import torch
import torch.nn.functional as F
import numpy as np
from tqdm.auto import tqdm

try:
    from safetensors.torch import load_file, save_file, safe_open
except ImportError:
    raise RuntimeError("safetensors is required. pip install safetensors")

try:
    from transformers import AutoConfig
    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False
    warnings.warn("transformers not found. Architecture auto-detection disabled.")


# =============================================================================
#                              ENUMS & DATA CLASSES
# =============================================================================

class AttentionType(Enum):
    MHA = "multi_head_attention"
    MQA = "multi_query_attention"
    GQA = "grouped_query_attention"

class MLPType(Enum):
    STANDARD = "standard_ffn"
    SWIGLU = "swiglu"
    GEGLU = "geglu"
    MOE = "mixture_of_experts"

class OffloadStrategy(Enum):
    NONE = "none"
    CPU = "cpu"
    DISK = "disk"
    AUTO = "auto"


@dataclass
class ArchitectureInfo:
    """Detected architecture information."""
    attention_type: AttentionType = AttentionType.MHA
    mlp_type: MLPType = MLPType.STANDARD
    num_attention_heads: int = 32
    num_kv_heads: int = 32
    hidden_size: int = 4096
    intermediate_size: int = 11008
    num_hidden_layers: int = 32
    vocab_size: int = 32000
    num_experts: int = 0
    num_experts_per_tok: int = 0
    layer_prefix: str = "model.layers"
    attn_prefix: str = "self_attn"
    mlp_prefix: str = "mlp"
    q_proj_name: str = "q_proj"
    k_proj_name: str = "k_proj"
    v_proj_name: str = "v_proj"
    o_proj_name: str = "o_proj"
    gate_proj_name: str = "gate_proj"
    up_proj_name: str = "up_proj"
    down_proj_name: str = "down_proj"
    has_fused_qkv: bool = False
    qkv_proj_name: str = "qkv_proj"


@dataclass
class DistillConfig:
    """Configuration for Universal SVD-LoRA Distillation.

    v3.7 goals:
      - Truly low-cost, out-of-the-box SVD distillation via truncated / randomized SVD by default.
      - Preserve v3.2's low-VRAM, shard-wise processing and layer offloading characteristics.
      - Improve cross-architecture robustness via efficient SVD-based projection with bounded rank.
    """

    # Paths
    teacher_path: str
    student_path: str
    output_path: str

    # LoRA rank controls
    rank_default: int = 64
    rank_attn: Optional[int] = None
    rank_mlp: Optional[int] = None
    lora_alpha: Optional[int] = None

    alpha_mode: str = "auto"  # auto | rank | fixed
    use_adaptive_rank: bool = True
    energy_threshold: float = 0.95
    min_rank: int = 8
    max_rank: int = 256

    # Delta regularization / merge helpers
    use_dare: bool = False
    dare_drop_rate: float = 0.3
    use_ties: bool = True
    ties_density: float = 0.3
    ties_trim_single: bool = True

    # Optional MoE blending (kept for backwards-compat; not used by default)
    moe_merge_method: str = "none"
    max_experts_to_blend: int = 2

    # Layer mapping schedule
    map_schedule: str = "linear"     # linear | sigmoid
    sigmoid_k: float = 0.15

    interp_mode: str = "schedule"  # schedule | lsq

    # Runtime / memory controls
    num_gpus: int = 1
    gpu_memory_fraction: float = 0.95
    use_cuda_streams: bool = True
    num_cuda_streams: int = 3
    use_mixed_precision: bool = True
    svd_batch_size: int = 1

    offload_strategy: str = "cpu"    # cpu | disk | none
    offload_dir: Optional[str] = None
    prefetch_layers: int = 2
    use_pinned_memory: bool = True
    max_cpu_memory_gb: float = 64.0

    # SVD execution / batching
    chunk_threshold_mb: int = 512
    chunk_size_mb: int = 128
    micro_batch_size: int = 8

    # Module selection
    include_pattern: str = "self_attn|mlp"
    exclude_pattern: str = ""
    include_embed_lm_head: bool = False

    # Cross-architecture projection (teacher tensor -> student tensor shape)
    use_svd_projection: bool = True
    projection_rank: int = 256
    projection_min_rank: int = 8
    projection_adaptive_rank: bool = True
    projection_energy_threshold: float = 0.99
    projection_randomized_iter: int = 2
    projection_randomized_oversamples: int = 8

    # Cross-arch subspace-delta controls
    subspace_delta: str = "off"  # off|auto|diag|full
    subspace_offdiag_threshold: float = 0.55
    subspace_denom: str = "core"  # core|full
    subspace_norm_match: bool = True

    # SVD backend selection
    #   - auto: heuristic (full vs low-rank)
    #   - full: torch.linalg.svd
    #   - randomized: Halko Alg. 4.4 (stable subspace iteration)
    #   - krylov: Block Krylov subspace (stores intermediate iterates)
    #   - adaptive: Adaptive range finder (Halko Alg. 4.2) for energy-targeted rank
    svd_mode: str = "auto"  # auto | full | randomized | krylov | adaptive | drf | aurora
    svd_randomized_iter: int = 2
    svd_randomized_oversamples: int = 8
    svd_auto_min_dim: int = 1024
    svd_auto_full_rank_ratio: float = 0.6
    svd_auto_lowrank: str = "randomized"  # randomized | krylov | drf | aurora

    # DRF-SVD (Dual-Residual Folding SVD) knobs
    #   - steps: number of residual-folding refinements (>=0)
    #   - theta: residual scaling exponent in [0,1] (0=none, 1=full 1/σ weighting)
    #   - resid_eps: numerical stabilizer for 1/σ^theta
    svd_drf_steps: int = 1
    svd_drf_theta: float = 0.5
    svd_drf_resid_eps: float = 1e-8

    # AURORA-SVD (Adaptive Unrolled Residual-Order Range Augmentation SVD) knobs
    #   - steps: refinement rounds (>=0)
    #   - order: 1=first-order residual; 2=adds second-order residual response via A/A^T (recommended)
    #   - theta1/theta2: residual scaling exponents in [0,1]
    #   - resid_eps: numerical stabilizer for scaling
    #   - level2_keep: fraction of components kept for the level-2 term in (0,1]; 1.0 keeps all
    svd_aurora_steps: int = 1
    svd_aurora_order: int = 2
    svd_aurora_theta1: float = 0.5
    svd_aurora_theta2: float = 0.25
    svd_aurora_resid_eps: float = 1e-8
    svd_aurora_level2_keep: float = 1.0

    # Adaptive range finder (Alg. 4.2) knobs (used when svd_mode=adaptive)
    svd_adaptive_block_size: int = 32  # columns added per iteration
    svd_adaptive_n_test: int = 8       # test vectors for residual estimate

    # Delta stability guard (scales delta if ||delta||/||student|| exceeds this ratio; 0 disables)
    max_delta_ratio: float = 0.35


    # Calibration-aware SVD (optional; default: off)
    #
    # Purpose: collect activation statistics on the *student* model (calibration prompts)
    # and apply data-aware whitening before SVD so the truncated low-rank factors better
    # match data-distribution sensitivity.
    #
    # calibration_mode:
    #   - "none": disable
    #   - "rms":  diagonal RMS whitening using per-channel activation RMS (input/output)
    #   - "cov":  SVD-LLM style full-covariance input whitening via Cholesky factor S
    #             of E[xx^T], using PiSSA on (Δ·S) and then unwhitening with S^{-1}.
    calibration_mode: str = "none"  # choices: none | rms | cov
    calib_data: Optional[str] = None          # .txt (1 prompt/line), .json/.jsonl, or Alpaca JSON (.json/.jsonl)
    calib_format: str = "auto"                # auto | txt | jsonl | json | alpaca
    calib_alpaca_template: str = "classic"     # classic | plain (only used when calib_format=alpaca or auto-detected)
    calib_alpaca_include_output: bool = False # include output tokens in calibration forward (alpaca only)
    calib_max_samples: int = 128              # number of prompts (after filtering empties)
    calib_max_length: int = 256               # max tokens per prompt
    calib_padding: str = "longest"          # longest | max_length (padding strategy during tokenization)
    calib_use_attention_mask: bool = True   # exclude attention_mask==0 tokens from activation stats (padding tokens)
    calib_batch_size: int = 1                 # batch size for forward passes
    calib_device: str = "auto"                # auto | cuda | cpu
    calib_dtype: str = "auto"                 # auto | bf16 | fp16 | fp32
    calib_load: Optional[str] = None          # load precomputed stats (.safetensors)
    calib_save: Optional[str] = None          # save computed stats (.safetensors)
    calib_eps: float = 1e-6                   # numerical stabilizer for whitening
    calib_collect_in: bool = True             # collect input RMS
    calib_collect_out: bool = True            # collect output RMS

    # Full-covariance (SVD-LLM) calibration controls (only used when calib_mode=cov)
    #
    # Notes:
    #   - cov whitening is applied on the *input* side (per SVD-LLM). For very large
    #     input dims (e.g., MLP down_proj), collection can be gated via calib_cov_max_dim.
    #   - calibration can be chunked across layers to bound peak memory.
    calib_cov_max_dim: int = 8192             # skip covariance collection if in_features > this
    calib_cov_chunk_layers: int = 4           # number of layers per calibration pass (>=1)
    calib_cov_groups: str = "qkv,o,mlp"       # comma-list: qkv,o,mlp,down
    calib_cov_store_dtype: str = "fp16"       # fp16|bf16|fp32 (saved chol factor dtype)

    # Misc
    auto_adjust_dare: bool = True
    seed: int = 42
    verbose: bool = True

# =============================================================================
#                           GPU MEMORY MANAGEMENT  
# =============================================================================

class GPUMemoryManager:
    def __init__(self, device: Union[str, torch.device], 
                 memory_fraction: float = 0.85,
                 verbose: bool = True):
        self.device = torch.device(device) if isinstance(device, str) else device
        self.memory_fraction = memory_fraction
        self.verbose = verbose
        
        if self.device.type == "cuda":
            torch.cuda.set_device(self.device)
            self.total_memory = torch.cuda.get_device_properties(self.device).total_memory
            self.max_memory = int(self.total_memory * memory_fraction)
        else:
            self.total_memory = 0
            self.max_memory = 0
    
    def get_stats(self) -> Dict[str, float]:
        if self.device.type != "cuda":
            return {"total_gb": 0, "used_gb": 0, "free_gb": 0}
        
        allocated = torch.cuda.memory_allocated(self.device)
        return {
            "total_gb": self.total_memory / (1024**3),
            "used_gb": allocated / (1024**3),
            "free_gb": (self.max_memory - allocated) / (1024**3),
        }
    
    def can_allocate(self, size_bytes: int) -> bool:
        if self.device.type != "cuda":
            return True
        allocated = torch.cuda.memory_allocated(self.device)
        return (allocated + size_bytes) < self.max_memory
    
    def clear_cache(self):
        gc.collect()
        if self.device.type == "cuda":
            torch.cuda.empty_cache()


# =============================================================================
#                           CUDA STREAM MANAGER
# =============================================================================

class CUDAStreamManager:
    def __init__(self, device: torch.device, num_streams: int = 3):
        self.device = device
        self.num_streams = num_streams
        self.streams = []
        self.current_stream_idx = 0
        
        if device.type == "cuda":
            for _ in range(num_streams):
                self.streams.append(torch.cuda.Stream(device=device))
    
    @contextmanager
    def stream_context(self, stream_idx: Optional[int] = None):
        if not self.streams:
            yield
            return
        
        idx = stream_idx if stream_idx is not None else self.current_stream_idx
        self.current_stream_idx = (self.current_stream_idx + 1) % self.num_streams
        
        with torch.cuda.stream(self.streams[idx]):
            yield
    
    def synchronize_current(self):
        if self.streams:
            self.streams[self.current_stream_idx].synchronize()
    
    def synchronize_all(self):
        for stream in self.streams:
            stream.synchronize()


# =============================================================================
#                           LAYER OFFLOADER
# =============================================================================

class LayerOffloader:
    def __init__(self,
                 strategy: OffloadStrategy,
                 device: torch.device,
                 offload_dir: Optional[str] = None,
                 use_pinned: bool = True,
                 max_cpu_gb: float = 32.0,
                 prefetch_count: int = 2,
                 verbose: bool = True):
        self.strategy = strategy
        self.device = device
        self.use_pinned = use_pinned and device.type == "cuda"
        self.max_cpu_bytes = int(max_cpu_gb * 1024**3)
        self.prefetch_count = prefetch_count
        self.verbose = verbose
        
        self.gpu_cache: Dict[str, torch.Tensor] = {}
        self.cpu_cache: Dict[str, torch.Tensor] = {}
        self.disk_paths: Dict[str, str] = {}
        self.cpu_bytes_used = 0
        
        if offload_dir:
            self.offload_dir = offload_dir
        else:
            self.offload_dir = tempfile.mkdtemp(prefix="distill_offload_")
        os.makedirs(self.offload_dir, exist_ok=True)
        
        self.prefetch_queue = queue.Queue()
        self.stop_prefetch = threading.Event()
        self.prefetch_thread = None
        self.transfer_stream = None
        
        if device.type == "cuda" and strategy != OffloadStrategy.NONE:
            self.transfer_stream = torch.cuda.Stream(device=device)
            self.prefetch_thread = threading.Thread(target=self._prefetch_worker, daemon=True)
            self.prefetch_thread.start()
    
    def _prefetch_worker(self):
        while not self.stop_prefetch.is_set():
            try:
                key = self.prefetch_queue.get(timeout=0.1)
                if key not in self.gpu_cache:
                    self._load_to_gpu(key)
            except queue.Empty:
                continue
    
    def _to_pinned(self, tensor: torch.Tensor) -> torch.Tensor:
        if self.use_pinned and not tensor.is_pinned():
            pinned = torch.empty(tensor.shape, dtype=tensor.dtype, pin_memory=True)
            pinned.copy_(tensor)
            return pinned
        return tensor
    
    def store(self, key: str, tensor: torch.Tensor, priority: str = "cpu"):
        if self.strategy == OffloadStrategy.NONE:
            self.gpu_cache[key] = tensor
        else:
            size = tensor.numel() * tensor.element_size()
            if priority == "disk":
                self._store_to_disk(key, tensor)
            elif self.cpu_bytes_used + size <= self.max_cpu_bytes:
                cpu_tensor = tensor.cpu()
                self.cpu_cache[key] = self._to_pinned(cpu_tensor)
                self.cpu_bytes_used += size
            else:
                self._store_to_disk(key, tensor)
    
    def _store_to_disk(self, key: str, tensor: torch.Tensor):
        safe_key = re.sub(r'[^\w\-.]', '_', key)
        path = os.path.join(self.offload_dir, f"{safe_key}.pt")
        torch.save(tensor.cpu(), path)
        self.disk_paths[key] = path
    
    def _load_to_gpu(self, key: str) -> Optional[torch.Tensor]:
        if key in self.gpu_cache:
            return self.gpu_cache[key]
        
        tensor = None
        if key in self.cpu_cache:
            tensor = self.cpu_cache[key]
        elif key in self.disk_paths:
            tensor = torch.load(self.disk_paths[key])
        
        if tensor is not None:
            if self.transfer_stream:
                with torch.cuda.stream(self.transfer_stream):
                    gpu_tensor = tensor.to(self.device, non_blocking=True)
                self.transfer_stream.synchronize()
            else:
                gpu_tensor = tensor.to(self.device)
            
            self.gpu_cache[key] = gpu_tensor
            return gpu_tensor
        return None
    
    def get(self, key: str, prefetch_next: Optional[List[str]] = None) -> Optional[torch.Tensor]:
        if prefetch_next:
            for next_key in prefetch_next[:self.prefetch_count]:
                if next_key not in self.gpu_cache:
                    self.prefetch_queue.put(next_key)
        return self._load_to_gpu(key)
    
    def evict(self, key: str):
        if key in self.gpu_cache:
            del self.gpu_cache[key]
    
    def evict_all_gpu(self):
        self.gpu_cache.clear()
        if self.device.type == "cuda":
            torch.cuda.empty_cache()
    
    def cleanup(self):
        self.stop_prefetch.set()
        if self.prefetch_thread:
            self.prefetch_thread.join(timeout=1.0)
        self.gpu_cache.clear()
        self.cpu_cache.clear()
        if os.path.exists(self.offload_dir):
            shutil.rmtree(self.offload_dir, ignore_errors=True)


# =============================================================================
#                        GPU-ACCELERATED SVD
# =============================================================================

class GPUAcceleratedSVD:
    """SVD backend with GPU acceleration and low-memory low-rank fallbacks.

    Implemented low-rank backends:
      - Randomized SVD with numerically-stable subspace iteration (Halko et al., Alg. 4.4)
      - Block Krylov subspace SVD (stores intermediate iterates, improves accuracy per iteration)
      - Adaptive range finder (Halko et al., Alg. 4.2) for energy-targeted rank selection
      - DRF-SVD (Dual-Residual Folding SVD): residual-folding refinement for improved subspace capture
    """

    def __init__(
        self,
        device: torch.device,
        memory_manager: GPUMemoryManager,
        stream_manager: Optional[CUDAStreamManager],
        use_mixed_precision: bool = True,
        chunk_threshold_mb: int = 512,
        chunk_size_mb: int = 128,
        force_cpu: bool = False,
        verbose: bool = False,
        svd_mode: str = "auto",
        randomized_iter: int = 2,
        randomized_oversamples: int = 8,
        auto_min_dim: int = 1024,
        auto_full_rank_ratio: float = 0.6,
        auto_lowrank: str = "randomized",  # randomized|krylov
        adaptive_block_size: int = 32,
        adaptive_n_test: int = 8,
        drf_steps: int = 1,
        drf_theta: float = 0.5,
        drf_resid_eps: float = 1e-8,
        aurora_steps: int = 1,
        aurora_order: int = 2,
        aurora_theta1: float = 0.5,
        aurora_theta2: float = 0.25,
        aurora_resid_eps: float = 1e-8,
        aurora_level2_keep: float = 1.0,
    ):
        self.device = device
        self.memory_manager = memory_manager
        self.stream_manager = stream_manager
        self.use_mixed_precision = use_mixed_precision
        self.chunk_threshold_bytes = int(chunk_threshold_mb) * 1024 * 1024
        self.chunk_size_bytes = int(chunk_size_mb) * 1024 * 1024
        self.force_cpu = bool(force_cpu)
        self.verbose = bool(verbose)

        self.svd_mode = (svd_mode or "auto").lower()
        self.randomized_iter = int(randomized_iter)
        self.randomized_oversamples = int(randomized_oversamples)
        self.auto_min_dim = int(auto_min_dim)
        self.auto_full_rank_ratio = float(auto_full_rank_ratio)
        self.auto_lowrank = (auto_lowrank or "randomized").lower()

        self.adaptive_block_size = int(adaptive_block_size)
        self.adaptive_n_test = int(adaptive_n_test)
        # DRF-SVD knobs
        self.drf_steps = int(drf_steps)
        self.drf_theta = float(drf_theta)
        self.drf_resid_eps = float(drf_resid_eps)
        # AURORA-SVD knobs
        self.aurora_steps = int(aurora_steps)
        self.aurora_order = int(aurora_order)
        self.aurora_theta1 = float(aurora_theta1)
        self.aurora_theta2 = float(aurora_theta2)
        self.aurora_resid_eps = float(aurora_resid_eps)
        self.aurora_level2_keep = float(aurora_level2_keep)

        # Validate
        if self.svd_mode not in ("auto", "full", "randomized", "krylov", "adaptive", "drf", "aurora"):
            raise ValueError(
                f"Invalid svd_mode: {svd_mode} "
                f"(expected auto|full|randomized|krylov|adaptive|drf|aurora)"
            )
        if self.auto_lowrank not in ("randomized", "krylov", "drf", "aurora"):
            raise ValueError("auto_lowrank must be randomized|krylov|drf|aurora")
        if self.randomized_iter < 0:
            raise ValueError("randomized_iter must be >= 0")
        if self.randomized_oversamples < 0:
            raise ValueError("randomized_oversamples must be >= 0")
        if not (0.0 < self.auto_full_rank_ratio <= 1.0):
            raise ValueError("auto_full_rank_ratio must be in (0, 1]")
        if self.adaptive_block_size <= 0:
            raise ValueError("adaptive_block_size must be > 0")
        if self.adaptive_n_test <= 0:
            raise ValueError("adaptive_n_test must be > 0")
        if self.drf_steps < 0:
            raise ValueError("drf_steps must be >= 0")
        if not (0.0 <= self.drf_theta <= 1.0):
            raise ValueError("drf_theta must be in [0, 1]")
        if self.drf_resid_eps <= 0:
            raise ValueError("drf_resid_eps must be > 0")

        # AURORA validation
        if self.aurora_steps < 0:
            raise ValueError("aurora_steps must be >= 0")
        if self.aurora_order not in (1, 2):
            raise ValueError("aurora_order must be 1 or 2")
        if not (0.0 <= self.aurora_theta1 <= 1.0):
            raise ValueError("aurora_theta1 must be in [0, 1]")
        if not (0.0 <= self.aurora_theta2 <= 1.0):
            raise ValueError("aurora_theta2 must be in [0, 1]")
        if self.aurora_resid_eps <= 0:
            raise ValueError("aurora_resid_eps must be > 0")
        if not (0.0 <= self.aurora_level2_keep <= 1.0):
            raise ValueError("aurora_level2_keep must be in [0, 1]")


    @contextmanager
    def _stream_context(self):
        if self.stream_manager is None:
            yield
        else:
            with self.stream_manager.stream_context():
                yield

    # -------------------------------------------------------------------------
    # Public APIs
    # -------------------------------------------------------------------------

    def svd(self, tensor: torch.Tensor, rank: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Compute truncated SVD (U, S, Vh) with a bounded rank.

        Note:
          - svd_mode=adaptive is an energy-targeted algorithm; for fixed-rank calls it
            dispatches to the configured low-rank backend (auto_lowrank).
          - Use `svd_energy(...)` to run the Adaptive range finder (Alg. 4.2).
        """
        if tensor is None:
            raise ValueError("tensor is None")
        if tensor.dim() != 2:
            raise ValueError(f"SVD expects a 2D tensor, got dim={tensor.dim()}")

        m, n = tensor.shape
        min_dim = min(m, n)
        r = int(rank)
        if r <= 0:
            raise ValueError(f"rank must be > 0, got {rank}")
        r = min(r, min_dim)

        # Estimate size in fp32 (we cast for numerical stability)
        tensor_bytes = int(m) * int(n) * 4

        if self.force_cpu:
            return self._cpu_svd(tensor, r)

        mode = self.svd_mode

        # Heuristics for auto mode: decide full vs low-rank.
        if mode == "auto":
            if tensor_bytes > self.chunk_threshold_bytes:
                mode = self.auto_lowrank
            elif min_dim >= self.auto_min_dim and r < int(min_dim * self.auto_full_rank_ratio):
                mode = self.auto_lowrank
            else:
                mode = "full"

        # svd_mode=adaptive is only meaningful for `svd_energy`; for fixed-rank it becomes low-rank.
        if mode == "adaptive":
            mode = self.auto_lowrank

        if mode in ("randomized", "krylov", "drf", "aurora"):
            return self._lowrank_svd(
                tensor,
                r,
                algo=mode,
                n_oversamples=self.randomized_oversamples,
                n_iter=self.randomized_iter,
            )

        # Full SVD path
        try:
            return self._gpu_svd(tensor, r)
        except RuntimeError as e:
            if self.verbose:
                print(f"[SVD] GPU full SVD failed ({type(e).__name__}: {e}); falling back to low-rank SVD.")
            return self._lowrank_svd(
                tensor,
                r,
                algo=self.auto_lowrank,
                n_oversamples=self.randomized_oversamples,
                n_iter=self.randomized_iter,
            )

    def randomized_svd(
        self,
        tensor: torch.Tensor,
        rank: int,
        n_oversamples: Optional[int] = None,
        n_iter: Optional[int] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Explicit low-rank SVD (randomized/Krylov) with per-call parameters.

        Dispatch:
          - If svd_mode is krylov, uses Block Krylov.
          - If svd_mode is auto and auto_lowrank is krylov, uses Block Krylov.
          - Otherwise uses randomized subspace iteration (Alg. 4.4).
        """
        algo = "randomized"
        if self.svd_mode == "krylov":
            algo = "krylov"
        elif self.svd_mode == "drf":
            algo = "drf"
        elif self.svd_mode == "aurora":
            algo = "aurora"
        elif self.svd_mode == "auto":
            if self.auto_lowrank == "krylov":
                algo = "krylov"
            elif self.auto_lowrank == "drf":
                algo = "drf"
            elif self.auto_lowrank == "aurora":
                algo = "aurora"

        return self._lowrank_svd(
            tensor,
            int(rank),
            algo=algo,
            n_oversamples=self.randomized_oversamples if n_oversamples is None else int(n_oversamples),
            n_iter=self.randomized_iter if n_iter is None else int(n_iter),
        )

    def svd_energy(
        self,
        tensor: torch.Tensor,
        energy_threshold: float,
        min_rank: int,
        max_rank: int,
        block_size: Optional[int] = None,
        n_test: Optional[int] = None,
        n_iter: Optional[int] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Energy-targeted low-rank SVD using Adaptive range finder (Halko et al., Alg. 4.2).

        Parameters
        ----------
        energy_threshold:
            Target retained Frobenius energy fraction in (0,1]. The adaptive range finder
            controls the residual Frobenius energy via eps = sqrt(1 - energy_threshold).
        min_rank / max_rank:
            Hard rank bounds.
        block_size:
            Number of basis vectors added per adaptive iteration.
        n_test:
            Number of random test vectors for residual estimation.
        n_iter:
            Optional subspace-iteration steps applied to each new block before QR.
        """
        if tensor is None:
            raise ValueError("tensor is None")
        if tensor.dim() != 2:
            raise ValueError(f"SVD expects a 2D tensor, got dim={tensor.dim()}")

        with self._stream_context():
            x = tensor.to(self.device, non_blocking=True).float()
            m, n = x.shape
            min_dim = min(m, n)

            max_r = min(int(max_rank), min_dim)
            min_r = max(1, min(int(min_rank), max_r))

            tau = float(energy_threshold)
            tau = float(max(0.0, min(1.0, tau)))

            # Residual Frobenius tolerance: ||(I-QQ^T)A||_F <= eps ||A||_F
            eps = math.sqrt(max(0.0, 1.0 - tau))

            bs = self.adaptive_block_size if block_size is None else int(block_size)
            nt = self.adaptive_n_test if n_test is None else int(n_test)
            q = 0 if n_iter is None else int(n_iter)

            Q = self._adaptive_range_finder(
                x,
                eps=eps,
                max_rank=max_r,
                block_size=bs,
                n_test=nt,
                n_iter=q,
            )

            # Ensure at least min_rank basis vectors.
            if Q.shape[1] < min_r:
                add = min_r - Q.shape[1]
                omega = torch.randn(n, add, device=self.device, dtype=torch.float32)
                Y = x @ omega
                if Q.shape[1] > 0:
                    Y = Y - Q @ (Q.transpose(0, 1) @ Y)
                Qi, _ = torch.linalg.qr(Y, mode="reduced")
                Q = torch.cat([Q, Qi], dim=1)

            # Cap and re-orthonormalize.
            if Q.shape[1] > max_r:
                Q = Q[:, :max_r]
            Q, _ = torch.linalg.qr(Q, mode="reduced")

            B = Q.transpose(0, 1) @ x
            Ub, S, Vh = torch.linalg.svd(B, full_matrices=False)
            U = Q @ Ub
            return U, S, Vh

    # -------------------------------------------------------------------------
    # Full SVD backends
    # -------------------------------------------------------------------------

    def _gpu_svd(self, x: torch.Tensor, r: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Full SVD on GPU, then truncate to r."""
        with self._stream_context():
            x = x.to(self.device, non_blocking=True).float()
            U, S, Vh = torch.linalg.svd(x, full_matrices=False)
            return U[:, :r], S[:r], Vh[:r, :]

    def _cpu_svd(self, x: torch.Tensor, r: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Full SVD on CPU, then truncate to r."""
        x = x.cpu().float()
        U, S, Vh = torch.linalg.svd(x, full_matrices=False)
        return U[:, :r].to(self.device), S[:r].to(self.device), Vh[:r, :].to(self.device)

    # -------------------------------------------------------------------------
    # Low-rank SVD backends
    # -------------------------------------------------------------------------

    def _lowrank_svd(
        self,
        x: torch.Tensor,
        r: int,
        algo: str,
        n_oversamples: int = 8,
        n_iter: int = 2,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        if algo == "krylov":
            return self._krylov_svd(x, r, n_oversamples=n_oversamples, n_iter=n_iter)
        if algo == "drf":
            return self._drf_svd(x, r, n_oversamples=n_oversamples, n_iter=n_iter)
        if algo == "aurora":
            return self._aurora_svd(x, r, n_oversamples=n_oversamples, n_iter=n_iter)
        return self._randomized_svd(x, r, n_oversamples=n_oversamples, n_iter=n_iter)

    def _randomized_svd(
        self,
        x: torch.Tensor,
        r: int,
        n_oversamples: int = 8,
        n_iter: int = 2,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Randomized SVD using numerically-stable subspace iteration (Halko et al., Algorithm 4.4).

        This alternates applications of A and A^T with intermediate QR orthonormalization.
        """
        with self._stream_context():
            x = x.to(self.device, non_blocking=True).float()
            m, n = x.shape

            r = min(int(r), min(m, n))
            p = max(0, int(n_oversamples))
            k = min(min(m, n), r + p)
            if k <= 0:
                raise ValueError("randomized_svd: computed k <= 0")

            omega = torch.randn(n, k, device=self.device, dtype=torch.float32)
            Y = x @ omega
            Q, _ = torch.linalg.qr(Y, mode="reduced")

            q = max(0, int(n_iter))
            for _ in range(q):
                Y_tilde = x.transpose(0, 1) @ Q
                Q_tilde, _ = torch.linalg.qr(Y_tilde, mode="reduced")
                Y = x @ Q_tilde
                Q, _ = torch.linalg.qr(Y, mode="reduced")

            B = Q.transpose(0, 1) @ x
            Ub, S, Vh = torch.linalg.svd(B, full_matrices=False)
            U = Q @ Ub
            return U[:, :r], S[:r], Vh[:r, :]

    def _krylov_svd(
        self,
        x: torch.Tensor,
        r: int,
        n_oversamples: int = 8,
        n_iter: int = 2,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Block Krylov subspace SVD.

        Builds the Krylov subspace span{AΩ, (AA^T)AΩ, ..., (AA^T)^q AΩ} without forming AA^T,
        using alternating multiplies by A and A^T. Intermediate orthonormal blocks are stored
        and concatenated before solving the projected SVD.
        """
        with self._stream_context():
            x = x.to(self.device, non_blocking=True).float()
            m, n = x.shape

            r = min(int(r), min(m, n))
            p = max(0, int(n_oversamples))
            k = min(min(m, n), r + p)
            if k <= 0:
                raise ValueError("krylov_svd: computed k <= 0")

            omega = torch.randn(n, k, device=self.device, dtype=torch.float32)
            Y = x @ omega
            Q, _ = torch.linalg.qr(Y, mode="reduced")

            blocks = [Q]
            q = max(0, int(n_iter))
            for _ in range(q):
                # One stable subspace iteration step, but keep each block.
                Y_tilde = x.transpose(0, 1) @ Q
                Q_tilde, _ = torch.linalg.qr(Y_tilde, mode="reduced")
                Y = x @ Q_tilde
                Q, _ = torch.linalg.qr(Y, mode="reduced")
                blocks.append(Q)

            Qbar = torch.cat(blocks, dim=1)
            Qbar, _ = torch.linalg.qr(Qbar, mode="reduced")

            B = Qbar.transpose(0, 1) @ x
            Ub, S, Vh = torch.linalg.svd(B, full_matrices=False)
            U = Qbar @ Ub
            return U[:, :r], S[:r], Vh[:r, :]


    def _drf_svd(
        self,
        x: torch.Tensor,
        r: int,
        n_oversamples: int = 8,
        n_iter: int = 2,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        DRF-SVD (Dual-Residual Folding SVD) — a residual-folding refinement on top of a
        single randomized range capture.
    
        Motivation (distillation-centric):
          For distillation deltas, the spectrum is often moderately ill-conditioned and
          the LoRA rank budget is tight. Standard randomized SVD accuracy can be limited
          by incomplete capture of directions that are *correlated* with the current
          approximate singular space but live in its orthogonal complement. DRF-SVD
          explicitly augments both left and right subspaces with *scaled residual*
          directions from the current approximation and then performs a Rayleigh–Ritz
          compression back to the target rank.
    
        Algorithm sketch (one refinement step):
          1) Compute an initial k0=(r+p) low-rank SVD approximation via randomized SVD.
          2) Form residual blocks:
                R_L = A V - U diag(S)
                R_R = A^T U - V diag(S)
             then scale them by 1/(S^theta) to stabilize small/large singular components.
          3) Fold residual blocks into the subspaces:
                Q = orth([U, R_L]),  P = orth([V, R_R])
          4) Solve the projected SVD of M = Q^T A P and truncate back to k0.
          5) Repeat for `drf_steps` rounds (default 1), then return the top-r slice.
    
        Notes:
          - This is *not* power iteration: residual folding adds directions proportional to
            the current first-order optimality violations for each singular triplet.
          - Numerical stability: QR is used for orthonormalization; residual scaling uses
            an epsilon floor to avoid division by 0.
        """
        with self._stream_context():
            A = x.to(self.device, non_blocking=True).float()
            m, n = A.shape
            min_dim = min(m, n)
            r = min(int(r), min_dim)
            p = max(0, int(n_oversamples))
    
            # Internal working rank: keep (r+p) components for refinement, then truncate to r.
            k0 = min(min_dim, r + p)
            if k0 <= 0:
                raise ValueError("drf_svd: computed k0 <= 0")
    
            # Initial capture (use oversamples=0 here to make k0 the true working rank).
            U, S, Vh = self._randomized_svd(A, k0, n_oversamples=0, n_iter=max(0, int(n_iter)))
            # Ensure contiguous
            U = U.contiguous()
            S = S.contiguous()
            Vh = Vh.contiguous()
    
            steps = max(0, int(getattr(self, "drf_steps", 1)))
            theta = float(getattr(self, "drf_theta", 0.5))
            theta = float(max(0.0, min(1.0, theta)))
            eps = float(getattr(self, "drf_resid_eps", 1e-8))
    
            for _ in range(steps):
                # Current right singular basis
                V = Vh.transpose(0, 1).contiguous()  # (n, k0)
    
                # Residual blocks (m,k0) and (n,k0)
                AV = A @ V
                ATU = A.transpose(0, 1) @ U
    
                # R_L = A V - U diag(S);   R_R = A^T U - V diag(S)
                US = U * S.unsqueeze(0)
                VS = V * S.unsqueeze(0)
                R_L = AV - US
                R_R = ATU - VS
    
                # Scale residuals by 1/(S^theta) to reduce dominance of the largest modes.
                if theta > 0.0:
                    denom = torch.pow(torch.clamp(S, min=eps), theta).unsqueeze(0)
                    R_L = R_L / denom
                    R_R = R_R / denom
    
                # Fold residuals into subspaces and re-orthonormalize.
                Q = torch.cat([U, R_L], dim=1)
                P = torch.cat([V, R_R], dim=1)
    
                Q, _ = torch.linalg.qr(Q, mode="reduced")
                P, _ = torch.linalg.qr(P, mode="reduced")
    
                # Project and solve small SVD
                # M = Q^T A P, computed as (Q^T (A P)) for efficiency.
                AP = A @ P
                M = Q.transpose(0, 1) @ AP
    
                Uc, Sc, Vhc = torch.linalg.svd(M, full_matrices=False)
    
                # Keep k0 components for potential further refinement.
                kk = min(int(Sc.numel()), int(k0))
                if kk <= 0:
                    break
    
                U = (Q @ Uc[:, :kk]).contiguous()
                S = Sc[:kk].contiguous()
                Vh = (Vhc[:kk, :] @ P.transpose(0, 1)).contiguous()
    
                # Defensive re-orthonormalization drift guard (cheap; kk is small).
                U, _ = torch.linalg.qr(U, mode="reduced")
                Vt = Vh.transpose(0, 1).contiguous()
                Vt, _ = torch.linalg.qr(Vt, mode="reduced")
                Vh = Vt.transpose(0, 1).contiguous()
    
            # Return top-r
            U = U[:, :r]
            S = S[:r]
            Vh = Vh[:r, :]
            return U, S, Vh
    
    def _aurora_svd(
        self,
        x: torch.Tensor,
        r: int,
        n_oversamples: int = 8,
        n_iter: int = 2,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        AURORA-SVD (Adaptive Unrolled Residual-Order Range Augmentation SVD).

        Distillation-centric goal:
          Improve low-rank approximation quality under a tight LoRA rank budget by
          augmenting the approximation subspace with *structured residual orders*,
          then performing a Rayleigh–Ritz compression back to the working rank.

        Core idea (per refinement round):
          - Start from a working-rank approximation A ≈ U diag(S) V^T.
          - Form first-order residual blocks:
                R_L = A V − U diag(S)
                R_R = A^T U − V diag(S)
          - Optionally form a second-order residual response (order=2):
                T_L = A (scaled R_R),   T_R = A^T (scaled R_L)
          - Build enlarged subspaces Q, P and compute the best rank-(r+p) approximation
            inside span(Q) × span(P) by SVD of the projected core Q^T A P.

        Implementation notes (v4.2.0 speed path):
          - Uses *block-anchored orthonormalization*: keep current U/V blocks intact and
            orthonormalize residual blocks against them, avoiding full QR on concatenated
            bases while preserving the augmented subspaces (up to floating-point effects).
          - Avoids redundant re-orthonormalization of U/V: U and V remain orthonormal by
            construction (QR on residual blocks + orthogonalization against anchors; and
            SVD of the projected core).
          - Reuses A@V when constructing A@P (since V is an explicit leading block of P).
        """
        with self._stream_context():
            A = x.to(self.device, non_blocking=True).float()
            m, n = A.shape
            min_dim = min(m, n)

            r = min(int(r), min_dim)
            p = max(0, int(n_oversamples))
            k0 = min(min_dim, r + p)
            if k0 <= 0:
                raise ValueError("aurora_svd: computed k0 <= 0")

            # Initial capture (use oversamples=0 so k0 is the true working rank).
            U, S, Vh = self._randomized_svd(A, k0, n_oversamples=0, n_iter=max(0, int(n_iter)))
            U = U.contiguous()
            S = S.contiguous()
            Vh = Vh.contiguous()

            steps = max(0, int(getattr(self, "aurora_steps", 1)))
            order = int(getattr(self, "aurora_order", 2))
            order = 1 if order <= 1 else 2

            theta1 = float(getattr(self, "aurora_theta1", 0.5))
            theta2 = float(getattr(self, "aurora_theta2", 0.25))
            theta1 = float(max(0.0, min(1.0, theta1)))
            theta2 = float(max(0.0, min(1.0, theta2)))

            eps = float(getattr(self, "aurora_resid_eps", 1e-8))
            eps = float(max(1e-30, eps))

            level2_keep = float(getattr(self, "aurora_level2_keep", 1.0))
            level2_keep = float(max(0.0, min(1.0, level2_keep)))

            use_amp = bool(self.use_mixed_precision) and (self.device.type == "cuda")
            if use_amp:
                # Prefer bf16 when supported; otherwise fp16.
                amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
            else:
                amp_dtype = None  # type: ignore

            def _mm(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
                """Matrix multiply with an optional AMP speed path; returns fp32 tensor."""
                if use_amp:
                    with torch.autocast(device_type="cuda", dtype=amp_dtype):
                        return (X @ Y).float()
                return X @ Y

            def _orth_block(block: torch.Tensor, basis: torch.Tensor) -> torch.Tensor:
                """Orthonormalize `block` against an existing orthonormal `basis`."""
                if block.numel() == 0:
                    return block
                if basis.numel() != 0 and basis.shape[1] > 0:
                    block = block - basis @ (basis.transpose(0, 1) @ block)
                # If the block is (near) zero after projection, QR will produce garbage; handle safely.
                if torch.linalg.norm(block, ord="fro") < 1e-20:
                    return block[:, :0]
                Qb, _ = torch.linalg.qr(block, mode="reduced")
                return Qb

            for _ in range(steps):
                V = Vh.transpose(0, 1).contiguous()  # (n, k0)

                # First-order residuals
                AV = _mm(A, V)
                ATU = _mm(A.transpose(0, 1), U)

                US = U * S.unsqueeze(0)
                VS = V * S.unsqueeze(0)

                R_L = AV - US  # (m, k0)
                R_R = ATU - VS # (n, k0)

                # Residual scaling for stability / tail emphasis
                if theta1 > 0.0:
                    denom1 = torch.pow(torch.clamp(S, min=eps), theta1).unsqueeze(0)
                    R_Ls = R_L / denom1
                    R_Rs = R_R / denom1
                else:
                    R_Ls = R_L
                    R_Rs = R_R

                # Build orthonormal bases with anchored blocks (U / V).
                Q_blocks: List[torch.Tensor] = [U]
                P_blocks: List[torch.Tensor] = [V]

                Q1 = _orth_block(R_Ls, U)
                if Q1.numel() != 0 and Q1.shape[1] > 0:
                    Q_blocks.append(Q1)

                P1 = _orth_block(R_Rs, V)
                if P1.numel() != 0 and P1.shape[1] > 0:
                    P_blocks.append(P1)

                if order >= 2:
                    # Select a subset of components for the level-2 term (optional).
                    if level2_keep <= 0.0:
                        idx = None
                    elif level2_keep >= 1.0 or k0 <= 1:
                        idx = None
                    else:
                        # Residual energy per component: ||R_L[:,i]||^2 + ||R_R[:,i]||^2
                        eL = torch.sum(R_L.float() * R_L.float(), dim=0)
                        eR = torch.sum(R_R.float() * R_R.float(), dim=0)
                        e = (eL + eR).contiguous()
                        keep = max(1, int(round(level2_keep * k0)))
                        keep = min(keep, k0)
                        _, topi = torch.topk(e, k=keep, largest=True, sorted=False)
                        idx = topi

                    if idx is None:
                        R_L2 = R_Ls
                        R_R2 = R_Rs
                        S2 = S
                    else:
                        R_L2 = R_Ls[:, idx]
                        R_R2 = R_Rs[:, idx]
                        S2 = S[idx]

                    # Second-order residual response (apply A / A^T)
                    T_L = _mm(A, R_R2)
                    T_R = _mm(A.transpose(0, 1), R_L2)

                    # Orthogonalize level-2 blocks against current bases.
                    Q_basis = torch.cat(Q_blocks, dim=1)
                    P_basis = torch.cat(P_blocks, dim=1)
                    T_L = T_L - Q_basis @ (Q_basis.transpose(0, 1) @ T_L)
                    T_R = T_R - P_basis @ (P_basis.transpose(0, 1) @ T_R)

                    # Optional additional scaling for the level-2 response.
                    if theta2 > 0.0:
                        denom2 = torch.pow(torch.clamp(S2, min=eps), theta2).unsqueeze(0)
                        T_L = T_L / denom2
                        T_R = T_R / denom2

                    Q2 = _orth_block(T_L, Q_basis)
                    if Q2.numel() != 0 and Q2.shape[1] > 0:
                        Q_blocks.append(Q2)

                    P2 = _orth_block(T_R, P_basis)
                    if P2.numel() != 0 and P2.shape[1] > 0:
                        P_blocks.append(P2)

                # Final orthonormal bases (already orthonormal by construction)
                Q = torch.cat(Q_blocks, dim=1)
                P = torch.cat(P_blocks, dim=1)

                # Compute A@P with reuse of A@V (V is the explicit leading block of P).
                if len(P_blocks) == 1:
                    AP = AV
                else:
                    P_rest = torch.cat(P_blocks[1:], dim=1)
                    AP_rest = _mm(A, P_rest)
                    AP = torch.cat([AV, AP_rest], dim=1)

                # Project and solve small SVD (Rayleigh–Ritz)
                M = Q.transpose(0, 1) @ AP
                Uc, Sc, Vhc = torch.linalg.svd(M, full_matrices=False)

                kk = min(int(Sc.numel()), int(k0))
                if kk <= 0:
                    break

                U = (Q @ Uc[:, :kk]).contiguous()
                S = Sc[:kk].contiguous()
                Vh = (Vhc[:kk, :] @ P.transpose(0, 1)).contiguous()

            U = U[:, :r]
            S = S[:r]
            Vh = Vh[:r, :]
            return U, S, Vh
    # -------------------------------------------------------------------------
    # Adaptive range finder (Alg. 4.2)
    # -------------------------------------------------------------------------

    def _adaptive_range_finder(
        self,
        x: torch.Tensor,
        eps: float,
        max_rank: int,
        block_size: int,
        n_test: int,
        n_iter: int = 0,
    ) -> torch.Tensor:
        """
        Adaptive randomized range finder (Halko et al., Alg. 4.2).

        Maintains an estimate of the residual via random test vectors:
          H = (I - QQ^T) A G
        and stops when ||H||_F <= eps ||A||_F or max_rank is reached.
        """
        m, n = x.shape
        max_r = min(int(max_rank), min(m, n))
        if max_r <= 0:
            return torch.zeros((m, 0), device=self.device, dtype=torch.float32)

        bs = max(1, int(block_size))
        nt = max(1, int(n_test))
        q = max(0, int(n_iter))
        eps = float(max(0.0, eps))

        normA = torch.linalg.norm(x, ord="fro")
        if normA < 1e-20:
            return torch.zeros((m, 0), device=self.device, dtype=torch.float32)

        G = torch.randn(n, nt, device=self.device, dtype=torch.float32)
        H = x @ G

        Q: Optional[torch.Tensor] = None
        k = 0

        # Current residual estimate
        res = torch.linalg.norm(H, ord="fro")
        tol = eps * normA

        while (res > tol) and (k < max_r):
            b = min(bs, max_r - k)

            omega = torch.randn(n, b, device=self.device, dtype=torch.float32)
            Y = x @ omega  # (m, b)

            # Optional stabilization / spectrum-boost for each added block.
            if q > 0:
                # Stable subspace iteration restricted to the block.
                Qi, _ = torch.linalg.qr(Y, mode="reduced")
                for _ in range(q):
                    Z = x.transpose(0, 1) @ Qi
                    Qz, _ = torch.linalg.qr(Z, mode="reduced")
                    Y = x @ Qz
                    Qi, _ = torch.linalg.qr(Y, mode="reduced")
                Y = Qi

            # Orthogonalize against current basis.
            if Q is not None and Q.shape[1] > 0:
                Y = Y - Q @ (Q.transpose(0, 1) @ Y)

            Qi, _ = torch.linalg.qr(Y, mode="reduced")

            # Double-orthogonalize for numerical safety.
            if Q is not None and Q.shape[1] > 0:
                Qi = Qi - Q @ (Q.transpose(0, 1) @ Qi)
                Qi, _ = torch.linalg.qr(Qi, mode="reduced")

            if Qi.numel() == 0 or Qi.shape[1] == 0:
                break

            Q = Qi if Q is None else torch.cat([Q, Qi], dim=1)
            k = Q.shape[1]

            # Update residual sketch: H <- (I - QiQi^T) H
            H = H - Qi @ (Qi.transpose(0, 1) @ H)
            res = torch.linalg.norm(H, ord="fro")

        if Q is None:
            return torch.zeros((m, 0), device=self.device, dtype=torch.float32)

        # Final re-orthonormalization (cheap for small k, avoids drift)
        Q, _ = torch.linalg.qr(Q, mode="reduced")
        if Q.shape[1] > max_r:
            Q = Q[:, :max_r]
        return Q



# =============================================================================
#                           HELPER FUNCTIONS
# =============================================================================

NORM_TOKENS = ["layernorm", "rmsnorm", "layer_norm", "rms_norm", "ln_", "norm"]
EMBED_TOKENS = ["embed", "wte", "wpe", "lm_head", "output.weight"]


def seed_all(seed: int):
    torch.manual_seed(seed)
    np.random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def read_index_map(model_path: str) -> Dict[str, str]:
    """Read weight index map from model directory."""
    idx_files = [
        "model.safetensors.index.json",
        "pytorch_model.bin.index.json",
    ]
    
    for idx_file in idx_files:
        path = os.path.join(model_path, idx_file)
        if os.path.exists(path):
            with open(path, 'r') as f:
                data = json.load(f)
                return data.get("weight_map", {})
    
    # Single file models
    for ext in [".safetensors", ".bin"]:
        for name in ["model", "pytorch_model"]:
            path = os.path.join(model_path, f"{name}{ext}")
            if os.path.exists(path):
                if ext == ".safetensors":
                    with safe_open(path, framework="pt") as f:
                        return {k: f"{name}{ext}" for k in f.keys()}
    
    return {}


def scan_layers(keys: List[str]) -> Dict[str, List[int]]:
    """Scan model keys to find layer structure."""
    layers = defaultdict(set)
    
    for key in keys:
        m = re.search(r'(model\.layers|transformer\.h|decoder\.layers)\.(\d+)\.', key)
        if m:
            prefix = m.group(1)
            idx = int(m.group(2))
            layers[prefix].add(idx)
    
    return {k: sorted(list(v)) for k, v in layers.items()}


def split_key(key: str) -> Optional[Tuple[str, str, int, str]]:
    """Split a model key into components."""
    patterns = [
        r'^(model\.layers)\.(\d+)\.(.+)$',
        r'^(transformer\.h)\.(\d+)\.(.+)$',
        r'^(decoder\.layers)\.(\d+)\.(.+)$',
    ]
    
    for pattern in patterns:
        m = re.match(pattern, key)
        if m:
            prefix = m.group(1)
            idx = int(m.group(2))
            rest = m.group(3)
            token_name = rest.split('.')[0]
            return prefix, token_name, idx, rest
    
    return None


# =============================================================================
#                      ARCHITECTURE DETECTION
# =============================================================================

def detect_architecture_from_config(model_path: str) -> ArchitectureInfo:
    """Detect architecture from transformers config."""
    info = ArchitectureInfo()
    
    if not HAS_TRANSFORMERS:
        return info
    
    try:
        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
        
        info.hidden_size = getattr(config, "hidden_size", 4096)
        info.intermediate_size = getattr(config, "intermediate_size", 11008)
        info.num_hidden_layers = getattr(config, "num_hidden_layers", 32)
        info.num_attention_heads = getattr(config, "num_attention_heads", 32)
        info.vocab_size = getattr(config, "vocab_size", 32000)
        
        num_kv = getattr(config, "num_key_value_heads", None)
        if num_kv is None:
            num_kv = info.num_attention_heads
        info.num_kv_heads = num_kv
        
        if num_kv == info.num_attention_heads:
            info.attention_type = AttentionType.MHA
        elif num_kv == 1:
            info.attention_type = AttentionType.MQA
        else:
            info.attention_type = AttentionType.GQA
        
        info.num_experts = getattr(config, "num_local_experts", 0)
        if info.num_experts == 0:
            info.num_experts = getattr(config, "num_experts", 0)
        info.num_experts_per_tok = getattr(config, "num_experts_per_tok", 2)
        
        if info.num_experts > 0:
            info.mlp_type = MLPType.MOE
        elif hasattr(config, "hidden_act"):
            act = config.hidden_act.lower() if isinstance(config.hidden_act, str) else ""
            if "silu" in act or "swish" in act:
                info.mlp_type = MLPType.SWIGLU
            elif "gelu" in act:
                info.mlp_type = MLPType.GEGLU
    
    except Exception as e:
        warnings.warn(f"Config detection failed: {e}")
    
    return info


def detect_architecture_from_weights(keys: List[str]) -> ArchitectureInfo:
    """Detect architecture from weight keys."""
    info = ArchitectureInfo()
    
    key_str = " ".join(keys[:500])
    
    if ".experts." in key_str or "block_sparse_moe" in key_str:
        info.mlp_type = MLPType.MOE
        expert_indices = set()
        for k in keys:
            m = re.search(r'experts\.(\d+)', k)
            if m:
                expert_indices.add(int(m.group(1)))
        if expert_indices:
            info.num_experts = max(expert_indices) + 1
    
    if "c_attn" in key_str or "query_key_value" in key_str or "qkv_proj" in key_str:
        info.has_fused_qkv = True
    
    layer_indices = set()
    for k in keys:
        m = re.search(r'layers\.(\d+)', k) or re.search(r'\.h\.(\d+)', k)
        if m:
            layer_indices.add(int(m.group(1)))
    if layer_indices:
        info.num_hidden_layers = max(layer_indices) + 1
    
    return info


# =============================================================================
#                           DARE + TIES
# =============================================================================

def apply_dare(delta: torch.Tensor, drop_rate: float = 0.7, 
               rescale: bool = True) -> torch.Tensor:
    """DARE: Drop And REscale for delta parameters."""
    if delta.dim() != 2 or drop_rate <= 0:
        return delta
    
    mask = torch.rand_like(delta.float()) > drop_rate
    sparse_delta = delta * mask
    
    if rescale and drop_rate < 1.0:
        sparse_delta = sparse_delta / (1.0 - drop_rate)
    
    return sparse_delta


def apply_ties_single(delta: torch.Tensor, density: float = 0.3) -> torch.Tensor:
    """
    Apply TIES-style magnitude trimming to a single delta.
    
    This is the "Trim" step of TIES: keep only the top-k parameters by magnitude.
    Useful for sparsification even when there's only one delta.
    """
    if delta.dim() != 2 or density >= 1.0:
        return delta
    
    dtype = delta.dtype
    device = delta.device
    
    # Flatten and find threshold
    flat = delta.float().view(-1)
    k = max(1, int(flat.numel() * density))
    
    # Get threshold value (k-th largest magnitude)
    threshold = torch.topk(flat.abs(), k, largest=True).values[-1]
    
    # Create mask and apply
    mask = flat.abs() >= threshold
    trimmed = flat * mask.float()
    
    return trimmed.view(delta.shape).to(dtype)


def apply_ties(deltas: List[torch.Tensor], density: float = 0.3) -> torch.Tensor:
    """TIES-Merging: Trim, Elect Sign, Disjoint Merge."""
    if not deltas:
        return torch.zeros(1)

    if len(deltas) == 1:
        return deltas[0]

    device = deltas[0].device
    dtype = deltas[0].dtype
    shape = deltas[0].shape
    deltas = [d if d.shape == shape else torch.zeros(shape, device=device, dtype=dtype) 
              for d in deltas]

    stacked = torch.stack([d.float() for d in deltas], dim=0)

    k = max(1, int(stacked[0].numel() * density))
    abs_flat = stacked.abs().view(len(deltas), -1)

    trimmed = []
    for i in range(len(deltas)):
        threshold = torch.topk(abs_flat[i], k, largest=True).values[-1]
        mask = abs_flat[i] >= threshold
        trimmed_flat = stacked[i].view(-1) * mask.float()
        trimmed.append(trimmed_flat.view(shape))

    stacked_trimmed = torch.stack(trimmed, dim=0)

    signs = torch.sign(stacked_trimmed)
    sign_sum = signs.sum(dim=0)
    elected_sign = torch.sign(sign_sum)
    elected_sign = torch.where(elected_sign == 0, torch.ones_like(elected_sign), elected_sign)

    merged = torch.zeros(shape, device=device, dtype=dtype)
    counts = torch.zeros(shape, device=device, dtype=dtype)

    for t in stacked_trimmed:
        match = (torch.sign(t) == elected_sign) | (t == 0)
        merged += torch.where(match, t, torch.zeros_like(t))
        counts += match.float()

    counts = torch.clamp(counts, min=1)
    return (merged / counts).to(dtype)



def lsq_mix_weight_from_deltas(
    delta_floor: torch.Tensor,
    delta_ceil: torch.Tensor,
    fallback: float,
    eps: float = 1e-12,
) -> float:
    """Compute convex mixing weight w in [0,1] minimizing ||(1-w)Δ_floor + wΔ_ceil||_F.

    Weight-space least-squares mixing rule:

      argmin_w || Δ_floor + w(Δ_ceil - Δ_floor) ||_F^2

      w* = - <Δ_floor, Δ_ceil - Δ_floor> / ||Δ_ceil - Δ_floor||_F^2

    If the denominator is too small or inputs are invalid, returns fallback.
    """
    try:
        if delta_floor is None or delta_ceil is None:
            return float(fallback)

        if delta_floor.shape != delta_ceil.shape:
            return float(fallback)

        d0 = delta_floor.float()
        d1 = delta_ceil.float()
        d = d1 - d0

        denom = float(torch.sum(d * d).item())
        if not (denom > eps) or not math.isfinite(denom):
            return float(fallback)

        numer = float((-torch.sum(d0 * d)).item())
        if not math.isfinite(numer):
            return float(fallback)

        w = numer / denom
        if not math.isfinite(w):
            return float(fallback)

        if w < 0.0:
            w = 0.0
        elif w > 1.0:
            w = 1.0
        return float(w)
    except Exception:
        return float(fallback)


# =============================================================================
#                           ADAPTIVE RANK
# =============================================================================

def compute_adaptive_rank(
    singular_values: torch.Tensor,
    energy_threshold: float = 0.95,
    min_rank: int = 8,
    max_rank: int = 256,
    total_energy: Optional[Union[float, torch.Tensor]] = None,
) -> int:
    """
    Compute an adaptive rank based on retained Frobenius-energy.

    Parameters
    ----------
    singular_values:
        1D tensor of (estimated) singular values, typically sorted in descending order.
    energy_threshold:
        Target fraction of Frobenius energy to retain.
    min_rank / max_rank:
        Hard bounds on the returned rank.
    total_energy:
        If provided, this is interpreted as ||A||_F^2 for the *original* matrix A.
        This is the preferred mode when `singular_values` is truncated (e.g. you only
        computed the top-k singular values). If not provided, `total_energy` is
        computed as sum(s_i^2) over the supplied `singular_values` (legacy behavior).

    Notes
    -----
    For truncated spectra, using `total_energy=None` can systematically overestimate the
    retained-energy ratio (because missing tail energy is treated as 0). Passing
    `total_energy=||A||_F^2` avoids this.
    """
    s = singular_values.detach().float()
    if s.numel() == 0:
        return max(0, int(min_rank))

    # Guard against tiny negatives from numerical noise (SVD should be non-negative).
    s = torch.clamp(s, min=0.0)

    if total_energy is None:
        total = torch.sum(s ** 2)
    else:
        if torch.is_tensor(total_energy):
            total = total_energy.to(device=s.device, dtype=s.dtype)
        else:
            total = torch.tensor(float(total_energy), device=s.device, dtype=s.dtype)

    if total < 1e-20:
        # Degenerate matrix: keep at least min_rank (and at least 1 if possible).
        return max(1, int(min_rank))

    cumulative = torch.cumsum(s ** 2, dim=0)
    target = float(energy_threshold) * total

    hit = (cumulative >= target)
    if bool(hit.any()):
        r = int(hit.nonzero(as_tuple=True)[0][0].item() + 1)
    else:
        r = int(s.numel())

    r = max(int(min_rank), min(r, int(max_rank), int(s.numel())))
    return r

# =============================================================================
#                  IMPROVED DIMENSION PROJECTION (FIX #2)
# =============================================================================

def svd_projection(
    src: torch.Tensor,
    target_shape: Tuple[int, ...],
    svd_engine: "GPUAcceleratedSVD",
    config: "DistillConfig",
    device: torch.device,
    ref: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """Project a tensor to `target_shape`.

    v3.7.6 fix (student-subspace alignment):
      The legacy (pool/crop/interpolate) projection loses too much information and can
      destabilize merging (gibberish / repetition), especially for cross-architecture
      distillation.

      If `ref` (student tensor) is provided and both `src/ref` are 2D matrices, we:
        1) Compute a truncated SVD of teacher: Wt ≈ Ut diag(St) Vt^T
        2) Compute a truncated SVD of student: Ws ≈ Us diag(Ss) Vs^T
        3) Build teacher low-rank factors Bt, At (PiSSA-style: sqrt split)
        4) Resize factors to the student's row/col sizes (Bt -> rows, At -> cols)
        5) Project the resized teacher into the *student* singular subspaces:
             core = Us^T (Bt_rs @ At_rs) Vs
             Wproj = Us @ core @ Vs^T

      This explicitly converts teacher information into the student's low-rank subspace
      (the "common" space) and avoids the destructive "pool/crop" behavior.

    Notes:
      - For non-2D tensors, falls back to safe interpolation/crop.
      - For same-shape matrices, returns `src` unchanged.
    """
    if tuple(src.shape) == tuple(target_shape):
        return src

    # -------- 1D vectors: linear interpolation --------
    if src.dim() == 1 and len(target_shape) == 1:
        out_len = int(target_shape[0])
        x = torch.nan_to_num(src.float(), nan=0.0, posinf=0.0, neginf=0.0)
        if x.numel() == out_len:
            return x.to(dtype=src.dtype, device=src.device)
        x3 = x.view(1, 1, -1)
        y3 = F.interpolate(x3, size=out_len, mode="linear", align_corners=False)
        return y3.view(out_len).to(dtype=src.dtype, device=src.device)

    # -------- 2D matrices: student-subspace alignment (preferred) --------
    if src.dim() == 2 and len(target_shape) == 2:
        tm, tn = int(target_shape[0]), int(target_shape[1])
        m, n = int(src.shape[0]), int(src.shape[1])

        # Bound the working rank
        k_cap = int(getattr(config, 'projection_rank', 256))
        k_max = min(m, n, tm, tn, k_cap)
        if k_max <= 0:
            return torch.zeros((tm, tn), device=src.device, dtype=src.dtype)

        # Resize helpers
        def _resize_rows(mat: torch.Tensor, new_rows: int) -> torch.Tensor:
            if mat.shape[0] == new_rows:
                return mat
            # (rows, k) -> (1, k, rows) -> interpolate -> (new_rows, k)
            x = mat.transpose(0, 1).unsqueeze(0)
            y = F.interpolate(x, size=new_rows, mode="linear", align_corners=False)
            return y.squeeze(0).transpose(0, 1)

        def _resize_cols(mat: torch.Tensor, new_cols: int) -> torch.Tensor:
            if mat.shape[1] == new_cols:
                return mat
            # (k, cols) -> (1, k, cols) -> interpolate -> (k, new_cols)
            x = mat.unsqueeze(0)
            y = F.interpolate(x, size=new_cols, mode="linear", align_corners=False)
            return y.squeeze(0)

        # Fast fallback: if no ref matrix is provided
        if ref is None or (not isinstance(ref, torch.Tensor)) or ref.dim() != 2:
            # Teacher-only SVD factor resize (legacy), then reconstruct
            src_f = torch.nan_to_num(src.to(device).float(), nan=0.0, posinf=0.0, neginf=0.0)
            try:
                Ut, St, Vht = svd_engine.randomized_svd(
                    src_f,
                    k_max,
                    n_oversamples=getattr(config, 'projection_randomized_oversamples', 8),
                    n_iter=getattr(config, 'projection_randomized_iter', 2),
                )
                if getattr(config, 'projection_adaptive_rank', True):
                    k_t = compute_adaptive_rank(
                        St,
                        energy_threshold=getattr(config, 'projection_energy_threshold', 0.99),
                        min_rank=getattr(config, 'projection_min_rank', 8),
                        max_rank=k_max,
                    )
                    k = max(1, min(k_t, k_max))
                    Ut, St, Vht = Ut[:, :k], St[:k], Vht[:k, :]
                else:
                    k = k_max

                sqrtS = torch.sqrt(torch.clamp(St, min=0.0)).to(dtype=Ut.dtype)
                Bt = Ut * sqrtS.unsqueeze(0)          # (m, k)
                At = sqrtS.unsqueeze(1) * Vht         # (k, n)

                Bt_rs = _resize_rows(Bt, tm)
                At_rs = _resize_cols(At, tn)

                out = Bt_rs @ At_rs
                return out.to(dtype=src.dtype, device=src.device)
            except Exception:
                # Last resort: bilinear interpolation of the dense matrix
                dense = torch.nan_to_num(src.to(device).float(), nan=0.0, posinf=0.0, neginf=0.0)
                out = F.interpolate(dense.unsqueeze(0).unsqueeze(0), size=(tm, tn), mode="bilinear", align_corners=False)
                return out.squeeze(0).squeeze(0).to(dtype=src.dtype, device=src.device)

        # Ensure ref has the expected target shape
        ref = ref.to(device)
        if tuple(ref.shape) != (tm, tn):
            # We do not attempt to project ref; fallback to teacher-only projection
            return svd_projection(src, (tm, tn), svd_engine, config, device, ref=None)

        # Compute truncated SVDs (teacher + student)
        src_f = torch.nan_to_num(src.to(device).float(), nan=0.0, posinf=0.0, neginf=0.0)
        ref_f = torch.nan_to_num(ref.to(device).float(), nan=0.0, posinf=0.0, neginf=0.0)

        try:
            Ut, St, Vht = svd_engine.randomized_svd(
                src_f,
                k_max,
                n_oversamples=getattr(config, 'projection_randomized_oversamples', 8),
                n_iter=getattr(config, 'projection_randomized_iter', 2),
            )
            Us, Ss, Vhs = svd_engine.randomized_svd(
                ref_f,
                k_max,
                n_oversamples=getattr(config, 'projection_randomized_oversamples', 8),
                n_iter=getattr(config, 'projection_randomized_iter', 2),
            )

            if getattr(config, 'projection_adaptive_rank', True):
                k_t = compute_adaptive_rank(
                    St,
                    energy_threshold=getattr(config, 'projection_energy_threshold', 0.99),
                    min_rank=getattr(config, 'projection_min_rank', 8),
                    max_rank=k_max,
                )
                k_s = compute_adaptive_rank(
                    Ss,
                    energy_threshold=getattr(config, 'projection_energy_threshold', 0.99),
                    min_rank=getattr(config, 'projection_min_rank', 8),
                    max_rank=k_max,
                )
                k = max(1, min(k_t, k_s, k_max))
            else:
                k = k_max

            Ut, St, Vht = Ut[:, :k], St[:k], Vht[:k, :]
            Us, Ss, Vhs = Us[:, :k], Ss[:k], Vhs[:k, :]

            # Teacher low-rank factors (PiSSA sqrt split)
            sqrtSt = torch.sqrt(torch.clamp(St, min=0.0)).to(dtype=Ut.dtype)
            Bt = Ut * sqrtSt.unsqueeze(0)            # (m, k)
            At = sqrtSt.unsqueeze(1) * Vht           # (k, n)

            # Resize teacher factors to student dims
            Bt_rs = _resize_rows(Bt, tm)             # (tm, k)
            At_rs = _resize_cols(At, tn)             # (k, tn)

            # Student right singular vectors
            Vs = Vhs.transpose(0, 1).contiguous()    # (tn, k)

            # core = Us^T (Bt_rs @ At_rs) Vs  (k×k)
            # Compute as (Us^T Bt_rs) @ (At_rs Vs) for efficiency.
            left = Us.transpose(0, 1) @ Bt_rs        # (k, k)
            right = At_rs @ Vs                       # (k, k)
            core = left @ right                      # (k, k)

            out = Us @ core @ Vs.transpose(0, 1)     # (tm, tn)
            return out.to(dtype=src.dtype, device=src.device)

        except Exception:
            dense = torch.nan_to_num(src.to(device).float(), nan=0.0, posinf=0.0, neginf=0.0)
            out = F.interpolate(dense.unsqueeze(0).unsqueeze(0), size=(tm, tn), mode="bilinear", align_corners=False)
            return out.squeeze(0).squeeze(0).to(dtype=src.dtype, device=src.device)

    # -------- Generic fallback for other ranks/shapes --------
    # Fallback: flatten + interpolate length-wise, then reshape/crop.
    tgt_numel = int(torch.tensor(target_shape).prod().item())
    x = torch.nan_to_num(src.float().view(1, 1, -1), nan=0.0, posinf=0.0, neginf=0.0)
    y = F.interpolate(x, size=tgt_numel, mode="linear", align_corners=False)
    out = y.view(*target_shape)
    return out.to(dtype=src.dtype, device=src.device)

# =============================================================================
#           CROSS-ARCH SUBSPACE-DELTA (FIX: preserve student residual space)
# =============================================================================

def _resize_2d_rows_linear(mat: torch.Tensor, target_rows: int) -> torch.Tensor:
    """Resize a 2D matrix along rows using 1D linear interpolation (preserves columns)."""
    if mat.dim() != 2:
        return mat
    if mat.shape[0] == target_rows:
        return mat
    if target_rows <= 0:
        return mat[:0]
    # Treat columns as channels, interpolate over rows
    x = mat.transpose(0, 1).unsqueeze(0)  # (1, cols, rows)
    y = F.interpolate(x, size=target_rows, mode="linear", align_corners=False)
    return y.squeeze(0).transpose(0, 1).contiguous()


def _resize_2d_cols_linear(mat: torch.Tensor, target_cols: int) -> torch.Tensor:
    """Resize a 2D matrix along cols using 1D linear interpolation (preserves rows)."""
    if mat.dim() != 2:
        return mat
    if mat.shape[1] == target_cols:
        return mat
    if target_cols <= 0:
        return mat[:, :0]
    # Treat rows as channels, interpolate over cols
    x = mat.unsqueeze(0)  # (1, rows, cols)
    y = F.interpolate(x, size=target_cols, mode="linear", align_corners=False)
    return y.squeeze(0).contiguous()


def compute_student_svd_basis(
    student_mat: torch.Tensor,
    svd_engine: "GPUAcceleratedSVD",
    config: "DistillConfig",
    device: torch.device,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Compute the student's truncated SVD basis (U, S, V) used as the common subspace.
    Returns:
      U: (out, k), S: (k,), V: (in, k) where V = Vh^T
    """
    if student_mat.dim() != 2:
        raise ValueError("student_mat must be 2D")
    tm, tn = int(student_mat.shape[0]), int(student_mat.shape[1])
    k_cap = min(int(getattr(config, "projection_rank", 256)), tm, tn)
    if k_cap <= 0:
        raise ValueError("projection_rank too small for matrix")

    ref_f = torch.nan_to_num(student_mat.to(device).float(), nan=0.0, posinf=0.0, neginf=0.0)
    U, S, Vh = svd_engine.randomized_svd(
        ref_f,
        k_cap,
        n_oversamples=int(getattr(config, "projection_randomized_oversamples", 8)),
        n_iter=int(getattr(config, "projection_randomized_iter", 2)),
    )

    k = k_cap
    if getattr(config, "projection_adaptive_rank", True):
        k = compute_adaptive_rank(
            S,
            energy_threshold=float(getattr(config, "projection_energy_threshold", 0.99)),
            min_rank=int(getattr(config, "projection_min_rank", 8)),
            max_rank=k_cap,
        )
    k = max(1, min(int(k), int(S.numel()), k_cap))

    U = U[:, :k].contiguous()
    S = S[:k].contiguous()
    Vh = Vh[:k, :].contiguous()
    V = Vh.transpose(0, 1).contiguous()
    return U, S, V


def teacher_core_in_student_basis(
    teacher_mat: torch.Tensor,
    U_s: torch.Tensor,
    V_s: torch.Tensor,
    svd_engine: "GPUAcceleratedSVD",
    config: "DistillConfig",
    device: torch.device,
) -> torch.Tensor:
    """
    Compute teacher's representation in the student's SVD basis:
        core_t = U_s^T * Wt_resized * V_s
    where Wt_resized is constructed from teacher truncated SVD factors (PiSSA sqrt-split)
    and then row/col-resized to the student's matrix shape.

    Returns:
      core_t: (k, k) in float32 on `device`.
    """
    if teacher_mat.dim() != 2:
        raise ValueError("teacher_mat must be 2D")
    if U_s.dim() != 2 or V_s.dim() != 2:
        raise ValueError("U_s and V_s must be 2D")

    tm, tn = int(U_s.shape[0]), int(V_s.shape[0])
    k_s = int(U_s.shape[1])

    src_f = torch.nan_to_num(teacher_mat.to(device).float(), nan=0.0, posinf=0.0, neginf=0.0)
    r_cap = min(k_s, int(src_f.shape[0]), int(src_f.shape[1]))
    if r_cap <= 0:
        return torch.zeros((k_s, k_s), device=device, dtype=torch.float32)

    U_t, S_t, Vh_t = svd_engine.randomized_svd(
        src_f,
        r_cap,
        n_oversamples=int(getattr(config, "projection_randomized_oversamples", 8)),
        n_iter=int(getattr(config, "projection_randomized_iter", 2)),
    )

    r = r_cap
    if getattr(config, "projection_adaptive_rank", True):
        r = compute_adaptive_rank(
            S_t,
            energy_threshold=float(getattr(config, "projection_energy_threshold", 0.99)),
            min_rank=int(getattr(config, "projection_min_rank", 8)),
            max_rank=r_cap,
        )
    r = max(1, min(int(r), int(S_t.numel()), r_cap))

    U_t = U_t[:, :r]
    S_t = S_t[:r]
    Vh_t = Vh_t[:r, :]

    sqrt_S = torch.sqrt(torch.clamp(S_t.float(), min=1e-10)).to(dtype=U_t.dtype)
    B_t = (U_t * sqrt_S.unsqueeze(0)).contiguous()              # (m_t, r)
    A_t = (sqrt_S.unsqueeze(1) * Vh_t).contiguous()             # (r, n_t)

    B_rs = _resize_2d_rows_linear(B_t, tm)                       # (tm, r)
    A_rs = _resize_2d_cols_linear(A_t, tn)                       # (r, tn)

    # core = U_s^T (B_rs A_rs) V_s = (U_s^T B_rs) (A_rs V_s)
    left = U_s.transpose(0, 1) @ B_rs                            # (k, r)
    right = A_rs @ V_s                                           # (r, k)
    core = left @ right                                          # (k, k)
    return core.to(dtype=torch.float32)


def extract_lora_from_core_delta(
    delta_core: torch.Tensor,
    U_s: torch.Tensor,
    V_s: torch.Tensor,
    rank_hint: int,
    config: "DistillConfig",
    out_dtype: torch.dtype,
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], int]:
    """
    Build LoRA factors directly from a low-rank subspace delta:
        delta = U_s * delta_core * V_s^T

    We SVD delta_core (k x k) and lift it into (out x in) factors without forming the full matrix.
    """
    if delta_core.dim() != 2:
        return None, None, 0
    k = int(delta_core.shape[0])
    if k <= 0:
        return None, None, 0

    dc = torch.nan_to_num(delta_core.float(), nan=0.0, posinf=0.0, neginf=0.0)
    if torch.linalg.vector_norm(dc) < 1e-8:
        return None, None, 0

    # Full SVD on small core (k <= projection_rank)
    try:
        Uc, Sc, Vh = torch.linalg.svd(dc, full_matrices=False)
    except Exception:
        return None, None, 0

    if Sc.numel() == 0 or torch.sum(Sc ** 2) < 1e-12:
        return None, None, 0

    if getattr(config, "use_adaptive_rank", True):
        actual_rank = compute_adaptive_rank(
            Sc,
            energy_threshold=float(getattr(config, "energy_threshold", 0.95)),
            min_rank=int(getattr(config, "min_rank", 8)),
            max_rank=min(int(getattr(config, "max_rank", 256)), k),
        )
    else:
        actual_rank = int(rank_hint)
        actual_rank = max(int(getattr(config, "min_rank", 8)),
                          min(actual_rank, int(getattr(config, "max_rank", 256)), k))

    actual_rank = max(1, min(actual_rank, k, int(Sc.numel())))

    Uc_r = Uc[:, :actual_rank]
    Sc_r = Sc[:actual_rank]
    Vh_r = Vh[:actual_rank, :]

    sqrt_S = torch.sqrt(torch.clamp(Sc_r, min=1e-10))

    # Lift back to full matrix factors
    # B: (out, r) = (U_s @ Uc_r) * sqrt(S)
    # A: (r, in)  = sqrt(S) * (Vh_r @ V_s^T)
    U_lift = (U_s.float() @ Uc_r).contiguous()                   # (out, r)
    V_lift_t = (Vh_r @ V_s.float().transpose(0, 1)).contiguous() # (r, in)

    lora_B = (U_lift * sqrt_S.unsqueeze(0)).to(dtype=out_dtype).contiguous()
    lora_A = (sqrt_S.unsqueeze(1) * V_lift_t).to(dtype=out_dtype).contiguous()

    return lora_A, lora_B, actual_rank
def extract_lora_pissa(
    delta: torch.Tensor,
    rank: int,
    svd_engine: GPUAcceleratedSVD,
    config: DistillConfig,
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], int]:
    """
    Extract LoRA matrices using PiSSA (sqrt-split SVD factors).

    SVD improvements (no backward-compat constraints assumed):
      1) Adaptive-rank selection is measured against the *true* Frobenius energy ||Δ||_F^2,
         not just the energy of the (possibly truncated) singular-value vector.
      2) When adaptive rank is enabled, the SVD rank is grown geometrically until the
         energy target is met or max_rank is reached. This avoids always computing SVD at
         max_rank (which can be prohibitively expensive when max_rank is large).
    """
    if delta.dim() != 2:
        return None, None, 0

    out_features, in_features = delta.shape
    min_dim = min(out_features, in_features)

    # Sanitize numerical pathologies early.
    delta = torch.nan_to_num(delta, nan=0.0, posinf=1e6, neginf=-1e6)

    # Frobenius norm / energy (used for exact energy accounting in adaptive rank).
    delta_norm = torch.norm(delta.float())
    if delta_norm < 1e-10:
        return None, None, 0
    total_energy = delta_norm.float() ** 2  # ||Δ||_F^2

    # Respect hard bounds.
    max_rank = int(min(getattr(config, "max_rank", 256), min_dim))
    min_rank = int(max(1, min(getattr(config, "min_rank", 8), max_rank)))

    # If adaptive rank is disabled, do a single SVD at the requested rank (bounded).
    if not bool(getattr(config, "use_adaptive_rank", True)):
        actual_rank = int(max(min_rank, min(int(rank), max_rank)))
        try:
            U, S, Vh = svd_engine.svd(delta, actual_rank)
        except Exception as e:
            print(f"  SVD failed completely, skipping: {type(e).__name__}")
            return None, None, 0

        if S.numel() == 0 or torch.all(S < 1e-10):
            return None, None, 0

        S = torch.clamp(S, min=0.0)
        sqrtS = torch.sqrt(S).to(dtype=U.dtype)
        lora_B = U[:, :actual_rank] * sqrtS.unsqueeze(0)
        lora_A = sqrtS.unsqueeze(1) * Vh[:actual_rank, :]
        return lora_A.float(), lora_B.float(), actual_rank

    # Adaptive-rank path: grow the SVD rank until the retained energy meets the target.
    energy_threshold = float(getattr(config, "energy_threshold", 0.95))
    energy_threshold = float(max(0.0, min(1.0, energy_threshold)))
    target_energy = total_energy * energy_threshold

    # Fast-path: Adaptive range finder (Halko et al., Alg. 4.2) + one projected SVD.
    if str(getattr(config, "svd_mode", "auto")).lower() == "adaptive":
        try:
            U, S, Vh = svd_engine.svd_energy(
                delta,
                energy_threshold=energy_threshold,
                min_rank=min_rank,
                max_rank=max_rank,
                block_size=int(getattr(config, "svd_adaptive_block_size", 32)),
                n_test=int(getattr(config, "svd_adaptive_n_test", 8)),
                n_iter=int(getattr(config, "svd_randomized_iter", 0)),
            )
        except Exception as e:
            if bool(getattr(config, "verbose", False)):
                print(f"  Adaptive SVD failed, falling back to geometric SVD: {type(e).__name__}")
        else:
            if S.numel() == 0 or torch.all(S < 1e-10):
                return None, None, 0

            actual_rank = compute_adaptive_rank(
                S,
                energy_threshold=energy_threshold,
                min_rank=min_rank,
                max_rank=max_rank,
                total_energy=total_energy,
            )

            U = U[:, :actual_rank]
            S = S[:actual_rank]
            Vh = Vh[:actual_rank, :]

            S = torch.clamp(S, min=0.0)
            sqrtS = torch.sqrt(S).to(dtype=U.dtype)
            lora_B = U * sqrtS.unsqueeze(0)
            lora_A = sqrtS.unsqueeze(1) * Vh

            return lora_A.float(), lora_B.float(), actual_rank

    # Start from the user-provided rank (bounded), then grow geometrically as needed.
    r_try = int(max(min_rank, min(int(rank), max_rank)))
    r_try = max(1, min(r_try, max_rank))

    U: Optional[torch.Tensor] = None
    S: Optional[torch.Tensor] = None
    Vh: Optional[torch.Tensor] = None

    while True:
        try:
            U, S, Vh = svd_engine.svd(delta, r_try)
        except Exception as e:
            print(f"  SVD failed at rank {r_try}, skipping: {type(e).__name__}")
            return None, None, 0

        if S.numel() == 0 or torch.all(S < 1e-10):
            return None, None, 0

        # Compute retained energy for prefixes of the singular spectrum we actually computed.
        s = torch.clamp(S.detach().float(), min=0.0)
        cum_energy = torch.cumsum(s ** 2, dim=0)

        hit = (cum_energy >= target_energy)
        if bool(hit.any()):
            r_needed = int(hit.nonzero(as_tuple=True)[0][0].item() + 1)
            actual_rank = int(max(min_rank, min(r_needed, max_rank)))
            break

        # Not enough energy captured by top-r_try singular values.
        if r_try >= max_rank:
            actual_rank = int(max_rank)
            break

        # Grow rank geometrically (fast convergence, bounded by max_rank).
        r_next = int(max(r_try + 1, math.ceil(r_try * 2.0)))
        r_try = int(min(max_rank, r_next))

    # Slice to actual_rank (may be smaller than the last computed r_try).
    U = U[:, :actual_rank]
    S = S[:actual_rank]
    Vh = Vh[:actual_rank, :]

    S = torch.clamp(S, min=0.0)
    sqrtS = torch.sqrt(S).to(dtype=U.dtype)
    lora_B = U * sqrtS.unsqueeze(0)
    lora_A = sqrtS.unsqueeze(1) * Vh

    return lora_A.float(), lora_B.float(), actual_rank

# =============================================================================
#                 CALIBRATION-AWARE (DATA-AWARE) SVD  (OPTIONAL)
# =============================================================================
#
# Goal:
#   Collect activation statistics on the STUDENT model (calibration prompts) and apply
#   data-aware whitening before SVD so the truncated low-rank factors better match the
#   target input distribution.
#
# Modes:
#   - rms: diagonal whitening with per-channel activation RMS (input/output).
#   - cov: SVD-LLM style full-covariance whitening on the input side using a Cholesky
#          factor S of E[xx^T]. PiSSA is applied on (Δ·S) and then unwhitened by S^{-1}.
#
# Notes:
#   - Optional; default disabled ("none").
#   - Calibration only loads the STUDENT model; teacher weights are not needed.
#   - Stats are stored in safetensors (see save/load helpers below).


@dataclass
class CalibStatsEntry:
    """Calibration stats for a module or a deterministic shared-input group."""

    # Diagonal (ASVD-style) stats
    in_rms: Optional[torch.Tensor] = None    # (in_features,)
    out_rms: Optional[torch.Tensor] = None   # (out_features,)

    # Full-covariance (SVD-LLM) stats
    in_chol: Optional[torch.Tensor] = None   # (in_features, in_features) lower-triangular


def _dtype_from_str(dtype_str: str, device: torch.device) -> torch.dtype:
    ds = (dtype_str or "auto").lower()
    if ds == "bf16":
        return torch.bfloat16
    if ds == "fp16":
        return torch.float16
    if ds == "fp32":
        return torch.float32
    # auto
    if device.type == "cuda":
        try:
            if torch.cuda.is_bf16_supported():
                return torch.bfloat16
        except Exception:
            pass
        return torch.float16
    return torch.float32


def _infer_calib_device(calib_device: str) -> torch.device:
    cd = (calib_device or "auto").lower()
    if cd == "cuda":
        return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    if cd == "cpu":
        return torch.device("cpu")
    # auto
    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def _flatten_text_field(x: Any) -> str:
    """Best-effort: normalize a field that may be str / list[str] / other into a single string."""
    if x is None:
        return ""
    if isinstance(x, str):
        return x
    if isinstance(x, (list, tuple)):
        parts: List[str] = []
        for v in x:
            if v is None:
                continue
            if isinstance(v, str):
                s = v.strip()
            else:
                s = str(v).strip()
            if s:
                parts.append(s)
        return "\n".join(parts)
    try:
        return str(x)
    except Exception:
        return ""


def _looks_like_alpaca_record(obj: Any) -> bool:
    if not isinstance(obj, dict):
        return False
    # Classic Alpaca keys: instruction / input / output (some datasets use lists for input/output)
    if "instruction" in obj and ("input" in obj or "output" in obj):
        return True
    # Some variants add system_prompt or system
    if ("system_prompt" in obj or "system" in obj) and ("instruction" in obj or "input" in obj):
        return True
    return False


def _format_alpaca_prompt(obj: Dict[str, Any],
                          template: str = "classic",
                          include_output: bool = False) -> str:
    """
    Build a calibration prompt from an Alpaca-style record.

    Supported fields (best-effort):
      - system_prompt / system / system_message (optional)
      - instruction / prompt / query (required-ish)
      - input / context (optional)
      - output / response (optional; included only if include_output=True)
    """
    sys_p = _flatten_text_field(obj.get("system_prompt", obj.get("system", obj.get("system_message", "")))).strip()
    instruction = _flatten_text_field(obj.get("instruction", obj.get("prompt", obj.get("query", "")))).strip()
    inp = _flatten_text_field(obj.get("input", obj.get("context", ""))).strip()
    out = _flatten_text_field(obj.get("output", obj.get("response", ""))).strip()

    template = (template or "classic").lower().strip()
    prefix = (sys_p + "\n\n") if sys_p else ""

    if template == "plain":
        parts: List[str] = []
        if sys_p:
            parts.append(sys_p)
        if instruction:
            parts.append(instruction)
        if inp:
            parts.append(inp)
        if include_output and out:
            parts.append(out)
        return "\n\n".join(parts).strip()

    # classic
    if inp:
        prompt = f"{prefix}### Instruction:\n{instruction}\n\n### Input:\n{inp}\n\n### Response:\n"
    else:
        prompt = f"{prefix}### Instruction:\n{instruction}\n\n### Response:\n"
    if include_output and out:
        prompt = prompt + out
    return prompt


def _stream_json_array_first_n(path: str, n: int) -> Optional[List[Any]]:
    """
    Streaming parse for a top-level JSON array file: returns first n items without loading whole file.
    Returns None if the file does not look like a JSON array.
    """
    try:
        decoder = json.JSONDecoder()
        buf = ""
        items: List[Any] = []
        with open(path, "r", encoding="utf-8") as f:
            # Read until we find the opening '['
            while True:
                if not buf:
                    chunk = f.read(65536)
                    if not chunk:
                        return []
                    buf += chunk
                buf = buf.lstrip()
                if buf.startswith("["):
                    buf = buf[1:]
                    break
                # Not an array
                return None

            while len(items) < n:
                buf = buf.lstrip()
                if buf.startswith("]"):
                    break
                # Ensure buffer has enough content for a full JSON object
                while True:
                    try:
                        obj, idx = decoder.raw_decode(buf)
                        items.append(obj)
                        buf = buf[idx:]
                        buf = buf.lstrip()
                        if buf.startswith(","):
                            buf = buf[1:]
                        break
                    except json.JSONDecodeError:
                        chunk = f.read(65536)
                        if not chunk:
                            return items
                        buf += chunk
        return items
    except Exception:
        return None


def _load_calib_texts(path: Optional[str],
                      max_samples: int,
                      fmt: str = "auto",
                      alpaca_template: str = "classic",
                      alpaca_include_output: bool = False) -> List[str]:
    """
    Load calibration texts.

    Supported formats:
      - txt:   1 prompt per line
      - jsonl: 1 JSON per line; supports {"text": "..."} or Alpaca records
      - json:  JSON array/object; supports "text" or Alpaca records (streamed for arrays)
      - alpaca: force Alpaca parsing/formatting for .json/.jsonl
    """
    if not path:
        # Minimal built-in prompts (language-mixed) as a fallback to keep the feature plug-and-play.
        return [
            "你好，請用一句話介紹你自己。",
            "請解釋什麼是知識蒸餾（Knowledge Distillation）。",
            "請把下面句子翻譯成英文：人工智慧正在改變世界。",
            "Write a short paragraph about neural networks.",
            "Summarize the following in one sentence: Large language models learn patterns from data.",
            "List three key points about transformer attention.",
        ][:max_samples]

    fmt = (fmt or "auto").lower().strip()
    p_lower = str(path).lower()
    if fmt == "auto":
        if p_lower.endswith(".jsonl"):
            fmt = "jsonl"
        elif p_lower.endswith(".json"):
            fmt = "json"
        else:
            fmt = "txt"

    # For alpaca, we accept both .json and .jsonl
    force_alpaca = (fmt == "alpaca")
    if force_alpaca:
        # decide underlying file type
        if p_lower.endswith(".jsonl"):
            fmt_under = "jsonl"
        else:
            fmt_under = "json"
    else:
        fmt_under = fmt

    def _extract_text(obj: Any) -> Optional[str]:
        if obj is None:
            return None
        if isinstance(obj, str):
            s = obj.strip()
            return s if s else None
        if isinstance(obj, dict):
            if force_alpaca or _looks_like_alpaca_record(obj):
                t = _format_alpaca_prompt(obj, template=alpaca_template, include_output=alpaca_include_output)
                t = (t or "").strip()
                return t if t else None
            # json/jsonl generic: prefer "text"
            t = obj.get("text", None)
            if isinstance(t, str) and t.strip():
                return t.strip()
            # fallback: first string-ish value
            for v in obj.values():
                if isinstance(v, str) and v.strip():
                    return v.strip()
            return None
        if isinstance(obj, (list, tuple)):
            # list[str] -> join
            s = _flatten_text_field(obj).strip()
            return s if s else None
        # last resort
        try:
            s = str(obj).strip()
            return s if s else None
        except Exception:
            return None

    texts: List[str] = []

    try:
        if fmt_under == "txt":
            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    texts.append(line)
                    if len(texts) >= max_samples:
                        break
            return texts[:max_samples] if max_samples > 0 else texts

        if fmt_under == "jsonl":
            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        obj = json.loads(line)
                    except Exception:
                        continue
                    t = _extract_text(obj)
                    if t:
                        texts.append(t)
                    if len(texts) >= max_samples:
                        break
            if texts:
                return texts[:max_samples] if max_samples > 0 else texts
            # If empty, fall through to fallback prompts
            return _load_calib_texts(None, max_samples, "txt")

        if fmt_under == "json":
            # First try jsonl-style streaming (many datasets are .json but actually JSONL)
            try:
                with open(path, "r", encoding="utf-8") as f:
                    for line in f:
                        line = line.strip()
                        if not line:
                            continue
                        try:
                            obj = json.loads(line)
                        except Exception:
                            # not jsonl
                            texts = []
                            break
                        t = _extract_text(obj)
                        if t:
                            texts.append(t)
                        if len(texts) >= max_samples:
                            break
                if texts:
                    return texts[:max_samples] if max_samples > 0 else texts
            except Exception:
                pass

            # Streaming parse JSON array for first N items
            recs = _stream_json_array_first_n(path, max_samples if max_samples > 0 else 0)
            if recs is not None:
                for obj in recs:
                    t = _extract_text(obj)
                    if t:
                        texts.append(t)
                    if len(texts) >= max_samples:
                        break
                if texts:
                    return texts[:max_samples] if max_samples > 0 else texts

            # Fallback: full json.load (may be heavy for huge files; use jsonl if possible)
            with open(path, "r", encoding="utf-8") as f:
                obj = json.load(f)

            # Unwrap common containers
            records: List[Any] = []
            if isinstance(obj, list):
                records = obj
            elif isinstance(obj, dict):
                # common keys
                for k in ("data", "train", "records", "samples", "items"):
                    v = obj.get(k, None)
                    if isinstance(v, list):
                        records = v
                        break
                if not records:
                    records = [obj]
            else:
                records = [obj]

            for r in records:
                t = _extract_text(r)
                if t:
                    texts.append(t)
                if len(texts) >= max_samples:
                    break

            if texts:
                return texts[:max_samples] if max_samples > 0 else texts
            return _load_calib_texts(None, max_samples, "txt")

    except Exception:
        # If load fails, fall back to built-in prompts (do not hard-fail).
        return _load_calib_texts(None, max_samples, "txt")

    return texts[:max_samples] if max_samples > 0 else texts


class _RMSAccumulator:
    __slots__ = ("in_sumsq", "out_sumsq", "in_count", "out_count")

    def __init__(self, in_features: int, out_features: int, collect_in: bool, collect_out: bool):
        self.in_sumsq = torch.zeros(in_features, dtype=torch.float32) if collect_in else None
        self.out_sumsq = torch.zeros(out_features, dtype=torch.float32) if collect_out else None
        self.in_count = 0
        self.out_count = 0

    def add_in(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
        if self.in_sumsq is None:
            return
        # x: (..., in_features)
        x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0).float()
        x = x.reshape(-1, x.shape[-1])

        if attention_mask is not None and torch.is_tensor(attention_mask):
            m = attention_mask
            if m.dim() > 1:
                m = m.reshape(-1)
            else:
                m = m.reshape(-1)

            if m.numel() == x.shape[0]:
                m = m.to(device=x.device, dtype=x.dtype)
                v = (x * x * m.unsqueeze(1)).sum(dim=0)
                cnt = int(m.sum().item())
            else:
                v = (x * x).sum(dim=0)
                cnt = int(x.shape[0])
        else:
            v = (x * x).sum(dim=0)
            cnt = int(x.shape[0])

        if cnt <= 0:
            return
        self.in_sumsq += v.detach().cpu()
        self.in_count += cnt
    def add_out(self, y: torch.Tensor, attention_mask: Optional[torch.Tensor] = None):
        if self.out_sumsq is None:
            return
        # y: (..., out_features)
        y = torch.nan_to_num(y, nan=0.0, posinf=0.0, neginf=0.0).float()
        y = y.reshape(-1, y.shape[-1])

        if attention_mask is not None and torch.is_tensor(attention_mask):
            m = attention_mask
            if m.dim() > 1:
                m = m.reshape(-1)
            else:
                m = m.reshape(-1)

            if m.numel() == y.shape[0]:
                m = m.to(device=y.device, dtype=y.dtype)
                v = (y * y * m.unsqueeze(1)).sum(dim=0)
                cnt = int(m.sum().item())
            else:
                v = (y * y).sum(dim=0)
                cnt = int(y.shape[0])
        else:
            v = (y * y).sum(dim=0)
            cnt = int(y.shape[0])

        if cnt <= 0:
            return
        self.out_sumsq += v.detach().cpu()
        self.out_count += cnt
    def finalize(self, eps: float) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
        in_rms = None
        out_rms = None
        if self.in_sumsq is not None:
            denom = max(1, self.in_count)
            in_rms = torch.sqrt(self.in_sumsq / float(denom)).clamp(min=eps).contiguous()
        if self.out_sumsq is not None:
            denom = max(1, self.out_count)
            out_rms = torch.sqrt(self.out_sumsq / float(denom)).clamp(min=eps).contiguous()
        return in_rms, out_rms


class _CovAccumulator:
    """Accumulate full input covariance: sum(x^T x) and token count."""

    __slots__ = ("dim", "gram", "count")

    def __init__(self, dim: int, device: torch.device):
        self.dim = int(dim)
        self.gram = torch.zeros((self.dim, self.dim), device=device, dtype=torch.float32)
        self.count = 0

    def add(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> None:
        if x is None:
            return
        if x.dim() > 2:
            x = x.reshape(-1, x.shape[-1])
        if x.numel() == 0:
            return
        x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0).to(dtype=torch.float32)
        if x.shape[-1] != self.dim:
            return

        if attn_mask is not None:
            m = attn_mask
            if m.dim() > 1:
                m = m.reshape(-1)
            if m.numel() == x.shape[0]:
                m = m.to(device=x.device, dtype=torch.float32)
                x = x * m.unsqueeze(1)
                cnt = int(torch.sum(m).item())
            else:
                cnt = int(x.shape[0])
        else:
            cnt = int(x.shape[0])

        # x^T x
        self.gram += x.transpose(0, 1) @ x
        self.count += max(0, cnt)


def _parse_layer_index(name: str) -> Optional[int]:
    """Return layer index if name matches common HF layer naming."""
    m = re.search(r'(?:model\.layers|transformer\.h|decoder\.layers)\.(\d+)\.', name)
    if not m:
        return None
    try:
        return int(m.group(1))
    except Exception:
        return None


def _cov_group_key_for_module(module_name: str, student_arch: ArchitectureInfo) -> Optional[str]:
    """Map a module name to a deterministic shared-input covariance key."""
    if not module_name or module_name.count(".") < 2:
        return None
    leaf = module_name.split(".")[-1]
    base = ".".join(module_name.split(".")[:-1])

    qn = getattr(student_arch, "q_proj_name", "q_proj")
    kn = getattr(student_arch, "k_proj_name", "k_proj")
    vn = getattr(student_arch, "v_proj_name", "v_proj")
    on = getattr(student_arch, "o_proj_name", "o_proj")
    gn = getattr(student_arch, "gate_proj_name", "gate_proj")
    un = getattr(student_arch, "up_proj_name", "up_proj")
    dn = getattr(student_arch, "down_proj_name", "down_proj")

    if leaf in (qn, kn, vn):
        return f"{base}.qkv_in"
    if leaf == on:
        return f"{base}.o_in"
    if leaf in (gn, un, "w1", "w3"):
        return f"{base}.up_in"
    if leaf in (dn, "w2"):
        return f"{base}.down_in"
    return None


def _resolve_cov_entry_key(module_name: str, calib_stats: Dict[str, CalibStatsEntry], student_arch: ArchitectureInfo) -> Optional[str]:
    """Resolve the key to use for covariance whitening (exact module or shared-input group)."""
    if module_name in calib_stats and calib_stats[module_name].in_chol is not None:
        return module_name
    gk = _cov_group_key_for_module(module_name, student_arch)
    if gk and gk in calib_stats and calib_stats[gk].in_chol is not None:
        return gk
    return None


def _safe_cholesky(cov: torch.Tensor, eps: float, max_tries: int = 6) -> torch.Tensor:
    """Robust Cholesky with escalating jitter."""
    cov = (cov + cov.transpose(0, 1)) * 0.5
    d = int(cov.shape[0])
    eye = torch.eye(d, device=cov.device, dtype=cov.dtype)
    jitter = float(eps)
    for _ in range(max_tries):
        try:
            return torch.linalg.cholesky(cov + jitter * eye)
        except Exception:
            jitter *= 10.0
    # Last-resort: attempt cholesky_ex, then jitter more
    L, info = torch.linalg.cholesky_ex(cov + jitter * eye)
    if int(info) == 0:
        return L
    raise RuntimeError("Cholesky failed (covariance not SPD even after jitter).")


def collect_calibration_stats_rms(student_path: str,
                                  student_arch: ArchitectureInfo,
                                  config: DistillConfig) -> Dict[str, CalibStatsEntry]:
    """Collect per-channel activation RMS stats on the STUDENT model."""
    if getattr(config, "calibration_mode", "none") == "none":
        return {}

    if not HAS_TRANSFORMERS:
        raise RuntimeError("Calibration mode requires transformers installed.")

    from transformers import AutoTokenizer, AutoModelForCausalLM

    device = _infer_calib_device(getattr(config, "calib_device", "auto"))
    dtype = _dtype_from_str(getattr(config, "calib_dtype", "auto"), device)

    texts = _load_calib_texts(
        getattr(config, "calib_data", None),
        int(getattr(config, "calib_max_samples", 128)),
        getattr(config, "calib_format", "auto"),
        getattr(config, "calib_alpaca_template", "classic"),
        bool(getattr(config, "calib_alpaca_include_output", False)),
    )
    if not texts:
        return {}

    tok = AutoTokenizer.from_pretrained(student_path, trust_remote_code=True, use_fast=True)
    if tok.pad_token_id is None:
        # Make padding well-defined for batch calibration
        tok.pad_token = tok.eos_token

    # Load model with best-effort low-memory strategy
    model = None
    try:
        if device.type == "cuda":
            # Prefer device_map="auto" (requires accelerate). Fall back to full cuda load.
            try:
                model = AutoModelForCausalLM.from_pretrained(
                    student_path,
                    trust_remote_code=True,
                    torch_dtype=dtype,
                    device_map="auto",
                    low_cpu_mem_usage=True,
                )
            except Exception:
                model = AutoModelForCausalLM.from_pretrained(
                    student_path,
                    trust_remote_code=True,
                    torch_dtype=dtype,
                    low_cpu_mem_usage=True,
                ).to(device)
        else:
            model = AutoModelForCausalLM.from_pretrained(
                student_path,
                trust_remote_code=True,
                torch_dtype=dtype,
                low_cpu_mem_usage=True,
            ).to(device)
    except Exception as e:
        raise RuntimeError(f"Failed to load student model for calibration: {type(e).__name__}: {e}")

    model.eval()

    # Select target linear modules for stats collection
    include_re = None
    try:
        pat = getattr(config, "include_pattern", "self_attn|mlp") or "self_attn|mlp"
        include_re = re.compile(pat)
    except Exception:
        include_re = re.compile("self_attn|mlp")

    target_leaf_names = set([
        student_arch.q_proj_name, student_arch.k_proj_name, student_arch.v_proj_name, student_arch.o_proj_name,
        student_arch.gate_proj_name, student_arch.up_proj_name, student_arch.down_proj_name,
        # Fallback names commonly seen in SwiGLU variants
        "w1", "w2", "w3",
    ])

    acc: Dict[str, _RMSAccumulator] = {}
    mask_ref: Dict[str, Optional[torch.Tensor]] = {"mask": None}
    hooks: List[Any] = []

    def _should_hook(name: str, module: torch.nn.Module) -> bool:
        if not isinstance(module, torch.nn.Linear):
            return False
        if include_re is not None and (not include_re.search(name)):
            return False
        leaf = name.split(".")[-1]
        return leaf in target_leaf_names

    # Register hooks
    for name, module in model.named_modules():
        if not _should_hook(name, module):
            continue
        try:
            in_f = int(getattr(module, "in_features", 0))
            out_f = int(getattr(module, "out_features", 0))
            if in_f <= 0 or out_f <= 0:
                continue
            acc[name] = _RMSAccumulator(in_f, out_f,
                                        collect_in=bool(getattr(config, "calib_collect_in", True)),
                                        collect_out=bool(getattr(config, "calib_collect_out", True)))
        except Exception:
            continue

        def _pre_hook(mod, inputs, name=name):
            if name not in acc:
                return
            try:
                x = inputs[0]
                if isinstance(x, (tuple, list)):
                    x = x[0]
                if x is None or not torch.is_tensor(x) or x.dim() == 0:
                    return
                acc[name].add_in(x, mask_ref["mask"] if config.calib_use_attention_mask else None)
            except Exception:
                return

        def _fwd_hook(mod, inputs, output, name=name):
            if name not in acc:
                return
            try:
                y = output
                if isinstance(y, (tuple, list)):
                    y = y[0]
                if y is None or not torch.is_tensor(y) or y.dim() == 0:
                    return
                acc[name].add_out(y, mask_ref["mask"] if config.calib_use_attention_mask else None)
            except Exception:
                return

        hooks.append(module.register_forward_pre_hook(_pre_hook))
        hooks.append(module.register_forward_hook(_fwd_hook))

    if not acc:
        # Nothing matched; cleanup
        for h in hooks:
            try:
                h.remove()
            except Exception:
                pass
        del model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return {}

    # Determine input device
    try:
        input_device = next(model.parameters()).device
    except Exception:
        input_device = device

    bs = max(1, int(getattr(config, "calib_batch_size", 1)))
    max_len = max(8, int(getattr(config, "calib_max_length", 256)))

    with torch.inference_mode():
        for i in range(0, len(texts), bs):
            batch_texts = texts[i:i+bs]
            enc = tok(
                batch_texts,
                return_tensors="pt",
                padding=("max_length" if (config.calib_padding == "max_length" and max_len is not None) else True),
                truncation=True,
                max_length=max_len,
            )
            enc = {k: v.to(input_device) for k, v in enc.items()}
            mask_ref["mask"] = enc.get("attention_mask", None)
            try:
                _ = model(**enc, use_cache=False)
            except Exception:
                # Skip problematic batch
                continue

    # Cleanup hooks
    for h in hooks:
        try:
            h.remove()
        except Exception:
            pass

    stats: Dict[str, CalibStatsEntry] = {}
    eps = float(getattr(config, "calib_eps", 1e-6))
    for name, a in acc.items():
        in_rms, out_rms = a.finalize(eps)
        stats[name] = CalibStatsEntry(in_rms=in_rms, out_rms=out_rms)

    # Release model memory
    del model
    try:
        import gc
        gc.collect()
    except Exception:
        pass
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return stats


def collect_calibration_stats_cov(student_path: str,
                                  student_arch: ArchitectureInfo,
                                  config: DistillConfig) -> Dict[str, CalibStatsEntry]:
    """Collect full-covariance (SVD-LLM) input whitening stats on the STUDENT model.

    Produces in_chol factors keyed by deterministic shared-input group keys (see
    _cov_group_key_for_module). These factors are consumed by extract_lora_pissa_calibrated
    when calib_mode=cov.
    """
    if not HAS_TRANSFORMERS:
        raise RuntimeError("Calibration mode requires transformers installed.")

    from transformers import AutoTokenizer, AutoModelForCausalLM

    device = _infer_calib_device(getattr(config, "calib_device", "auto"))
    dtype = _dtype_from_str(getattr(config, "calib_dtype", "auto"), device)

    texts = _load_calib_texts(
        getattr(config, "calib_data", None),
        int(getattr(config, "calib_max_samples", 128)),
        getattr(config, "calib_format", "auto"),
        getattr(config, "calib_alpaca_template", "classic"),
        bool(getattr(config, "calib_alpaca_include_output", False)),
    )
    if not texts:
        return {}

    tok = AutoTokenizer.from_pretrained(student_path, trust_remote_code=True, use_fast=True)
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token

    torch_dtype = dtype if device.type == "cuda" else torch.float32
    try:
        model = AutoModelForCausalLM.from_pretrained(
            student_path,
            torch_dtype=torch_dtype,
            device_map=("auto" if (device.type == "cuda" and getattr(config, "calib_device", "auto") == "auto") else None),
            low_cpu_mem_usage=True,
            trust_remote_code=True,
        )
    except Exception as e:
        raise RuntimeError(f"Failed to load student model for calibration: {e}")

    model.eval()

    include_pat = getattr(config, "include_pattern", "")
    exclude_pat = getattr(config, "exclude_pattern", "")
    include_re = re.compile(include_pat) if include_pat else None
    exclude_re = re.compile(exclude_pat) if exclude_pat else None

    # Which covariance groups to collect
    groups_raw = str(getattr(config, "calib_cov_groups", "qkv,o,mlp")).split(",")
    want_groups = {g.strip().lower() for g in groups_raw if g.strip()}
    max_dim = int(getattr(config, "calib_cov_max_dim", 8192))
    chunk_layers = max(1, int(getattr(config, "calib_cov_chunk_layers", 4)))
    store_dtype_str = str(getattr(config, "calib_cov_store_dtype", "fp16")).lower()
    store_dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}.get(store_dtype_str, torch.float16)
    eps = float(getattr(config, "calib_eps", 1e-6))
    use_attn_mask = bool(getattr(config, "calib_use_attention_mask", True))

    # Build per-layer representative module selection.
    module_dict: Dict[str, torch.nn.Module] = dict(model.named_modules())

    qn = getattr(student_arch, "q_proj_name", "q_proj")
    kn = getattr(student_arch, "k_proj_name", "k_proj")
    vn = getattr(student_arch, "v_proj_name", "v_proj")
    on = getattr(student_arch, "o_proj_name", "o_proj")
    gn = getattr(student_arch, "gate_proj_name", "gate_proj")
    un = getattr(student_arch, "up_proj_name", "up_proj")
    dn = getattr(student_arch, "down_proj_name", "down_proj")

    priorities: Dict[str, List[str]] = {
        "qkv": [qn, kn, vn],
        "o": [on],
        "mlp": [gn, un, "w1", "w3"],
        "down": [dn, "w2"],
    }

    layer_sel: Dict[int, Dict[str, str]] = defaultdict(dict)
    layer_pri: Dict[int, Dict[str, int]] = defaultdict(dict)

    for name, mod in module_dict.items():
        if not isinstance(mod, torch.nn.Linear):
            continue
        if include_re and not include_re.search(name):
            continue
        if exclude_re and exclude_re.search(name):
            continue

        li = _parse_layer_index(name)
        if li is None:
            continue

        leaf = name.split(".")[-1]
        for g in want_groups:
            if g not in priorities:
                continue
            plist = priorities[g]
            if leaf not in plist:
                continue
            pidx = plist.index(leaf)
            if (g not in layer_sel[li]) or (pidx < layer_pri[li].get(g, 10**9)):
                layer_sel[li][g] = name
                layer_pri[li][g] = pidx

    layer_indices = sorted(layer_sel.keys())
    if not layer_indices:
        # Nothing matched; return empty.
        del model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return {}

    # Determine the device for feeding input tensors.
    try:
        input_device = next(model.parameters()).device
    except Exception:
        input_device = device

    max_len = int(getattr(config, "calib_max_length", 256)) if getattr(config, "calib_max_length", None) else None
    bs = int(getattr(config, "calib_batch_size", 1))

    stats: Dict[str, CalibStatsEntry] = {}

    # Chunked passes to bound peak memory.
    mask_ref: Dict[str, Optional[torch.Tensor]] = {"mask": None}
    for c0 in range(0, len(layer_indices), chunk_layers):
        chunk = layer_indices[c0:c0 + chunk_layers]

        acc: Dict[str, _CovAccumulator] = {}
        hooks: List[Any] = []

        # Register hooks for this chunk.
        for li in chunk:
            for g in want_groups:
                if g not in priorities:
                    continue
                mod_name = layer_sel.get(li, {}).get(g, None)
                if not mod_name:
                    continue
                mod = module_dict.get(mod_name, None)
                if mod is None or not isinstance(mod, torch.nn.Linear):
                    continue

                in_dim = int(getattr(mod, "in_features", 0) or 0)
                if in_dim <= 0 or in_dim > max_dim:
                    continue

                gk = _cov_group_key_for_module(mod_name, student_arch)
                if not gk:
                    continue
                if gk in acc:
                    continue

                try:
                    dev = next(mod.parameters()).device
                except Exception:
                    dev = input_device

                acc[gk] = _CovAccumulator(in_dim, dev)

                def _make_pre_hook(key: str):
                    def _pre_hook(_mod, inputs):
                        if not inputs:
                            return
                        x = inputs[0]
                        if not torch.is_tensor(x):
                            return
                        m = mask_ref.get("mask", None)
                        if not use_attn_mask:
                            m = None
                        acc[key].add(x, m)
                    return _pre_hook

                hooks.append(mod.register_forward_pre_hook(_make_pre_hook(gk)))

        if not acc:
            # Nothing to collect in this chunk.
            for h in hooks:
                try:
                    h.remove()
                except Exception:
                    pass
            continue

        # Run calibration forward passes.
        with torch.inference_mode():
            for i in range(0, len(texts), bs):
                batch_texts = texts[i:i + bs]
                enc = tok(
                    batch_texts,
                    return_tensors="pt",
                    padding=("max_length" if (getattr(config, "calib_padding", "longest") == "max_length" and max_len is not None) else True),
                    truncation=True,
                    max_length=max_len,
                )
                enc = {k: v.to(input_device) for k, v in enc.items()}
                mask_ref["mask"] = enc.get("attention_mask", None)
                try:
                    _ = model(**enc, use_cache=False)
                except Exception:
                    continue

        # Cleanup hooks for this chunk.
        for h in hooks:
            try:
                h.remove()
            except Exception:
                pass

        # Finalize and store cholesky factors.
        for gk, a in acc.items():
            if a.count <= 0:
                continue
            cov = a.gram / float(max(1, a.count))
            try:
                L = _safe_cholesky(cov, eps=eps)
            except Exception:
                continue

            stats[gk] = CalibStatsEntry(in_chol=L.detach().to(dtype=store_dtype).cpu().contiguous())

        # Release chunk accumulators.
        del acc
        try:
            import gc
            gc.collect()
        except Exception:
            pass
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    # Release model memory
    del model
    try:
        import gc
        gc.collect()
    except Exception:
        pass
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return stats


def collect_calibration_stats(student_path: str,
                              student_arch: ArchitectureInfo,
                              config: DistillConfig) -> Dict[str, CalibStatsEntry]:
    """Dispatch calibration collection based on config.calibration_mode."""
    mode = str(getattr(config, "calibration_mode", "none")).lower()
    if mode == "rms":
        return collect_calibration_stats_rms(student_path, student_arch, config)
    if mode == "cov":
        return collect_calibration_stats_cov(student_path, student_arch, config)
    return {}


def save_calibration_stats(stats: Dict[str, CalibStatsEntry], path: str) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    out: Dict[str, torch.Tensor] = {}
    for name, e in stats.items():
        if e.in_rms is not None:
            out[f"{name}.in_rms"] = e.in_rms.detach().cpu().contiguous()
        if e.out_rms is not None:
            out[f"{name}.out_rms"] = e.out_rms.detach().cpu().contiguous()
        if e.in_chol is not None:
            out[f"{name}.in_chol"] = e.in_chol.detach().cpu().contiguous()
    if not out:
        return
    save_file(out, path, metadata={"format": "ud_calib_stats_v2"})


def load_calibration_stats(path: str) -> Dict[str, CalibStatsEntry]:
    tensors = load_file(path)
    stats: Dict[str, CalibStatsEntry] = {}
    tmp: Dict[str, Dict[str, torch.Tensor]] = defaultdict(dict)
    for k, v in tensors.items():
        if k.endswith(".in_rms"):
            tmp[k[:-len(".in_rms")]]["in"] = v
        elif k.endswith(".out_rms"):
            tmp[k[:-len(".out_rms")]]["out"] = v
        elif k.endswith(".in_chol"):
            tmp[k[:-len(".in_chol")]]["in_chol"] = v
    for name, d in tmp.items():
        stats[name] = CalibStatsEntry(
            in_rms=d.get("in", None),
            out_rms=d.get("out", None),
            in_chol=d.get("in_chol", None),
        )
    return stats


def _delta_whiten_diag(delta: torch.Tensor,
                      in_rms: Optional[torch.Tensor],
                      out_rms: Optional[torch.Tensor],
                      eps: float) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
    """
    Return whitened delta and the (clamped) scaling vectors used.

    delta: (out, in)
    in_rms:  (in,)
    out_rms: (out,)
    """
    dw = delta
    in_s = None
    out_s = None

    if out_rms is not None:
        out_s = out_rms.to(device=delta.device, dtype=torch.float32).clamp(min=eps)
        dw = dw / out_s.unsqueeze(1)

    if in_rms is not None:
        in_s = in_rms.to(device=delta.device, dtype=torch.float32).clamp(min=eps)
        dw = dw / in_s.unsqueeze(0)

    return dw, in_s, out_s


def extract_lora_pissa_calibrated(
    delta: torch.Tensor,
    rank: int,
    svd_engine: GPUAcceleratedSVD,
    config: DistillConfig,
    module_name: str,
    calib_stats: Dict[str, CalibStatsEntry],
    student_arch: ArchitectureInfo,
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], int]:
    """Calibration-aware PiSSA.

    - rms: diagonal RMS whitening (ASVD-style), then unwhiten LoRA factors.
    - cov: SVD-LLM style full-covariance input whitening via Cholesky factor S of E[xx^T].
           PiSSA is applied on (Δ·S) and then unwhitened by S^{-1} on the right factor.
    """
    if not calib_stats:
        return extract_lora_pissa(delta, rank, svd_engine, config)

    mode = str(getattr(config, "calibration_mode", "none")).lower()
    eps = float(getattr(config, "calib_eps", 1e-6))

    if mode == "rms":
        entry = calib_stats.get(module_name, None)
        if entry is None:
            return extract_lora_pissa(delta, rank, svd_engine, config)
        if entry.in_rms is None and entry.out_rms is None:
            return extract_lora_pissa(delta, rank, svd_engine, config)

        dw, in_s, out_s = _delta_whiten_diag(delta.float(), entry.in_rms, entry.out_rms, eps)

        # Run PiSSA on whitened delta
        lora_A_w, lora_B_w, actual_rank = extract_lora_pissa(dw.to(delta.dtype), rank, svd_engine, config)
        if lora_A_w is None or lora_B_w is None or actual_rank <= 0:
            return None, None, 0

        # Unwhiten factors back to original space: Δ ≈ (D_out·B_w) (A_w·D_in)
        lora_A = lora_A_w.float()
        lora_B = lora_B_w.float()
        if out_s is not None:
            lora_B = lora_B * out_s.unsqueeze(1)
        if in_s is not None:
            lora_A = lora_A * in_s.unsqueeze(0)

        return lora_A.to(dtype=delta.dtype).contiguous(), lora_B.to(dtype=delta.dtype).contiguous(), actual_rank

    if mode == "cov":
        key = _resolve_cov_entry_key(module_name, calib_stats, student_arch)
        if not key:
            return extract_lora_pissa(delta, rank, svd_engine, config)

        S = calib_stats[key].in_chol
        if S is None:
            return extract_lora_pissa(delta, rank, svd_engine, config)

        # Ensure S is on the same device.
        S = S.to(device=delta.device, dtype=torch.float32)

        # PiSSA on (Δ·S)
        dw = (delta.float() @ S).to(dtype=delta.dtype)
        lora_A_w, lora_B_w, actual_rank = extract_lora_pissa(dw, rank, svd_engine, config)
        if lora_A_w is None or lora_B_w is None or actual_rank <= 0:
            return None, None, 0

        # Unwhiten on the right: A = A_w · S^{-1}
        A_w = lora_A_w.float().contiguous()
        B = lora_B_w.float().contiguous()

        # Solve (S^T) X^T = A_w^T  =>  X = A_w · S^{-1}
        try:
            A = torch.linalg.solve_triangular(S.transpose(0, 1), A_w.transpose(0, 1), upper=True).transpose(0, 1)
        except Exception:
            # Fallback: generic solve (slower, but robust)
            A = torch.linalg.solve(S.transpose(0, 1), A_w.transpose(0, 1)).transpose(0, 1)

        return A.to(dtype=delta.dtype).contiguous(), B.to(dtype=delta.dtype).contiguous(), actual_rank

    return extract_lora_pissa(delta, rank, svd_engine, config)


# =============================================================================
#                        LAYER MAPPING
# =============================================================================

def teacher_idx_from_student_idx(s_idx: int, s_total: int, t_total: int,
                                  schedule: str = "sigmoid",
                                  k: float = 0.15) -> Tuple[int, float]:
    """Map student layer index to teacher layer index with interpolation weight."""
    if t_total <= 1 or s_total <= 1:
        return 0, 0.0
    
    s_norm = s_idx / (s_total - 1)
    
    if schedule == "sigmoid":
        x = 12 * (s_norm - 0.5)
        t_norm = 1 / (1 + math.exp(-k * x))
    else:
        t_norm = s_norm
    
    t_float = t_norm * (t_total - 1)
    t_floor = int(t_float)
    t_floor = max(0, min(t_floor, t_total - 1))
    interp_w = t_float - t_floor
    
    return t_floor, interp_w


def create_layer_mapping(teacher_info: ArchitectureInfo,
                         student_info: ArchitectureInfo) -> Dict[str, str]:
    """Create mapping from student layer names to teacher layer names."""
    mapping = {}
    
    attn_map = [
        (student_info.q_proj_name, teacher_info.q_proj_name),
        (student_info.k_proj_name, teacher_info.k_proj_name),
        (student_info.v_proj_name, teacher_info.v_proj_name),
        (student_info.o_proj_name, teacher_info.o_proj_name),
    ]
    
    mlp_map = [
        (student_info.gate_proj_name, teacher_info.gate_proj_name),
        (student_info.up_proj_name, teacher_info.up_proj_name),
        (student_info.down_proj_name, teacher_info.down_proj_name),
    ]
    
    for s_name, t_name in attn_map + mlp_map:
        if s_name != t_name:
            mapping[s_name] = t_name
    
    return mapping


# =============================================================================
#                          PEFT FORMAT OUTPUT
# =============================================================================

def to_peft_key(param_key: str) -> str:
    """Convert parameter key to PEFT-compatible format."""
    key = param_key
    if key.endswith(".weight"):
        key = key[:-len(".weight")]
    
    base_prefix = "base_model.model."
    if key.startswith(base_prefix):
        return key
    if key.startswith("model."):
        return base_prefix + key
    elif key.startswith("transformer."):
        return base_prefix + key
    else:
        return base_prefix + "model." + key


def save_peft_adapter(lora_weights: Dict[str, torch.Tensor],
                      output_dir: str,
                      config: DistillConfig,
                      student_arch: ArchitectureInfo,
                      rank_info: Dict[str, int]):
    """Save LoRA weights in PEFT-compatible format."""
    os.makedirs(output_dir, exist_ok=True)
    
    weights_path = os.path.join(output_dir, "adapter_model.safetensors")
    # safetensors requires CPU + contiguous tensors (no views).
    safe_weights: Dict[str, torch.Tensor] = {}
    for k, v in lora_weights.items():
        if not isinstance(v, torch.Tensor):
            continue
        t = v.detach()
        if t.device.type != "cpu":
            t = t.to("cpu")
        # Ensure contiguous storage (also breaks view/reference sharing).
        t = t.contiguous()
        safe_weights[k] = t
    lora_weights = safe_weights

    save_file(lora_weights, weights_path)
    
    target_modules = set()
    for key in lora_weights.keys():
        module = re.sub(r'\.lora_[AB](?:\.[^\.]+)?\.weight$', '', key)
        module = re.sub(r'^base_model\.model\.', '', module)
        parts = module.split('.')
        for part in reversed(parts):
            if part in ['q_proj', 'k_proj', 'v_proj', 'o_proj',
                       'gate_proj', 'up_proj', 'down_proj',
                       'c_attn', 'c_proj', 'c_fc', 'query_key_value',
                       'w1', 'w2', 'w3', 'fc1', 'fc2', 'dense']:
                target_modules.add(part)
                break
    
    if rank_info:
        rank_counts = defaultdict(int)
        for r in rank_info.values():
            rank_counts[r] += 1
        effective_rank = max(rank_counts.items(), key=lambda x: x[1])[0]
    else:
        effective_rank = config.rank_default

    def _rank_key_for_config(peft_key: str) -> str:
        # PEFT rank_pattern/alpha_pattern match base model module names (e.g. 'model.layers.0...'),
        # not the wrapped state_dict prefix (e.g. 'base_model.model.').
        prefix = "base_model.model."
        return peft_key[len(prefix):] if peft_key.startswith(prefix) else peft_key

    

    # alpha_pattern: keep per-module scaling stable when rank_pattern varies
    alpha_pattern: Dict[str, int] = {}
    alpha_mode = getattr(config, "alpha_mode", "auto")
    if alpha_mode == "auto":
        alpha_mode = "rank" if config.lora_alpha is None else "fixed"
    if alpha_mode == "rank" and rank_info:
        alpha_pattern = {}
        for _k, _v in rank_info.items():
            _iv = int(_v)
            alpha_pattern[_rank_key_for_config(_k)] = _iv
    rank_pattern: Dict[str, int] = {}
    if rank_info:
        for _k, _v in rank_info.items():
            _iv = int(_v)
            rank_pattern[_rank_key_for_config(_k)] = _iv

    adapter_config = {
        "alpha_pattern": alpha_pattern,
        "auto_mapping": None,
        "base_model_name_or_path": config.student_path,
        "bias": "none",
        "fan_in_fan_out": False,
        "inference_mode": True,
        "init_lora_weights": True,
        "layer_replication": None,
        "layers_pattern": None,
        "layers_to_transform": None,
        "loftq_config": {},
        "lora_alpha": config.lora_alpha or effective_rank,
        "lora_dropout": 0.0,
        "megatron_config": None,
        "megatron_core": "megatron.core",
        "modules_to_save": None,
        "peft_type": "LORA",
        "r": effective_rank,
        "rank_pattern": rank_pattern,
        "revision": None,
        "target_modules": sorted(list(target_modules)),
        "task_type": "CAUSAL_LM",
        "use_dora": False,
        "use_rslora": False,
    }
    
    config_path = os.path.join(output_dir, "adapter_config.json")
    with open(config_path, 'w') as f:
        json.dump(adapter_config, f, indent=2)
    
    print(f"\n✅ Saved PEFT adapter to: {output_dir}")
    print(f"   - Weights: {weights_path}")
    print(f"   - Config: {config_path}")
    print(f"   - Effective rank: {effective_rank}")
    print(f"   - Target modules: {sorted(list(target_modules))}")


# =============================================================================
#                      ACCELERATED DISTILLER
# =============================================================================

class AcceleratedDistiller:
    """GPU-accelerated distillation with memory management."""
    
    def __init__(self, config: DistillConfig, device: torch.device,
                 teacher_arch: ArchitectureInfo, student_arch: ArchitectureInfo):
        self.config = config
        self.device = device
        self.teacher_arch = teacher_arch
        self.student_arch = student_arch
        
        # Check for cross-architecture distillation
        self.is_cross_arch = (
            teacher_arch.hidden_size != student_arch.hidden_size or
            teacher_arch.num_hidden_layers != student_arch.num_hidden_layers
        )
        
        if self.is_cross_arch and config.verbose:
            print(f"\n⚠️  Cross-architecture distillation detected!")
            print(f"   Teacher: {teacher_arch.num_hidden_layers} layers, {teacher_arch.hidden_size} hidden")
            print(f"   Student: {student_arch.num_hidden_layers} layers, {student_arch.hidden_size} hidden")
        
        # FIX #3: Auto-adjust DARE for cross-architecture
        if self.is_cross_arch and config.auto_adjust_dare and config.use_dare:
            # Reduce drop rate for cross-arch to preserve more info
            old_rate = config.dare_drop_rate
            config.dare_drop_rate = min(0.1, config.dare_drop_rate)
            if config.verbose and old_rate != config.dare_drop_rate:
                print(f"   Auto-adjusted DARE drop_rate: {old_rate} → {config.dare_drop_rate}")
        
        self.memory_manager = GPUMemoryManager(
            device, config.gpu_memory_fraction, config.verbose
        )
        
        self.stream_manager = CUDAStreamManager(
            device, config.num_cuda_streams
        ) if config.use_cuda_streams and device.type == "cuda" else None
        
        if config.offload_strategy == "auto":
            stats = self.memory_manager.get_stats()
            if stats["total_gb"] < 16:
                strategy = OffloadStrategy.CPU
            else:
                strategy = OffloadStrategy.NONE
        else:
            strategy = OffloadStrategy[config.offload_strategy.upper()]
        
        self.offloader = LayerOffloader(
            strategy=strategy,
            device=device,
            offload_dir=config.offload_dir,
            use_pinned=config.use_pinned_memory,
            max_cpu_gb=config.max_cpu_memory_gb,
            prefetch_count=config.prefetch_layers,
            verbose=config.verbose
        )
        
        self.svd_engine = GPUAcceleratedSVD(
            device=device,
            memory_manager=self.memory_manager,
            stream_manager=self.stream_manager,
            use_mixed_precision=config.use_mixed_precision,
            chunk_threshold_mb=config.chunk_threshold_mb,
            chunk_size_mb=config.chunk_size_mb,
            verbose=config.verbose,
            svd_mode=getattr(config, 'svd_mode', 'auto'),
            randomized_iter=getattr(config, 'svd_randomized_iter', 2),
            randomized_oversamples=getattr(config, 'svd_randomized_oversamples', 8),
            auto_min_dim=getattr(config, 'svd_auto_min_dim', 1024),
            auto_full_rank_ratio=getattr(config, 'svd_auto_full_rank_ratio', 0.6),
            auto_lowrank=getattr(config, 'svd_auto_lowrank', 'randomized'),
            adaptive_block_size=getattr(config, 'svd_adaptive_block_size', 32),
            adaptive_n_test=getattr(config, 'svd_adaptive_n_test', 8),
            drf_steps=int(getattr(config, 'svd_drf_steps', 1)),
            drf_theta=float(getattr(config, 'svd_drf_theta', 0.5)),
            drf_resid_eps=float(getattr(config, 'svd_drf_resid_eps', 1e-8)),
            aurora_steps=int(getattr(config, 'svd_aurora_steps', 1)),
            aurora_order=int(getattr(config, 'svd_aurora_order', 2)),
            aurora_theta1=float(getattr(config, 'svd_aurora_theta1', 0.5)),
            aurora_theta2=float(getattr(config, 'svd_aurora_theta2', 0.25)),
            aurora_resid_eps=float(getattr(config, 'svd_aurora_resid_eps', 1e-8)),
            aurora_level2_keep=float(getattr(config, 'svd_aurora_level2_keep', 1.0)),
        )
        
        self.lora_weights: Dict[str, torch.Tensor] = {}
        self.rank_info: Dict[str, int] = {}
        # MoE->Dense synthesis cache (teacher MoE -> student dense)
        self._teacher_moe_style: Optional[str] = None
        self._teacher_has_mlp_experts: bool = False
        self._teacher_has_block_sparse_moe: bool = False
        self._moe_dense_cache: Dict[Tuple[str, int, str], torch.Tensor] = {}
        self.stats = defaultdict(int)
        self.calib_stats = None  # optional: {module_name: (in_rms, out_rms)}
    
    def load_tensors_with_prefetch(self,
                                   keys: List[str],
                                   model_folder: str,
                                   weight_map: Dict[str, str],
                                   next_keys: Optional[List[str]] = None) -> Dict[str, torch.Tensor]:
        """Load tensors with optional prefetching."""
        shards = defaultdict(list)
        for k in keys:
            if k in weight_map:
                shards[weight_map[k]].append(k)
        
        tensors = {}
        
        for shard, shard_keys in shards.items():
            path = os.path.join(model_folder, shard)
            with safe_open(path, framework="pt") as f:
                for k in shard_keys:
                    try:
                        t = f.get_tensor(k)
                        scale_key = f"{k}.scales"
                        if scale_key in f.keys():
                            scale = f.get_tensor(scale_key)
                            t = t.float() * scale.float()
                            t = torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0)
                        tensors[k] = t
                    except Exception:
                        continue
        
        if self.stream_manager:
            with self.stream_manager.stream_context():
                for k in tensors:
                    tensors[k] = tensors[k].to(self.device, non_blocking=True)
            self.stream_manager.synchronize_current()
        else:
            for k in tensors:
                tensors[k] = tensors[k].to(self.device)
        
        return tensors
    

    # -----------------------------
    # MoE -> Dense support (teacher MoE, student dense)
    # -----------------------------
    def _init_teacher_moe_style(self, teacher_weight_map: Dict[str, str]) -> None:
        """One-time detection of teacher MoE key layout."""
        if self._teacher_moe_style is not None:
            return
        # Detect common HF MoE layouts
        self._teacher_has_mlp_experts = any(".mlp.experts." in k for k in teacher_weight_map.keys())
        self._teacher_has_block_sparse_moe = any("block_sparse_moe.experts." in k for k in teacher_weight_map.keys())

        if self._teacher_has_mlp_experts and self._teacher_has_block_sparse_moe:
            self._teacher_moe_style = "both"
        elif self._teacher_has_mlp_experts:
            self._teacher_moe_style = "mlp_experts"
        elif self._teacher_has_block_sparse_moe:
            self._teacher_moe_style = "block_sparse_moe"
        else:
            self._teacher_moe_style = "none"

    def _estimate_moe_expert_weights(
        self,
        layer_prefix: str,
        layer_idx: int,
        teacher_weight_map: Dict[str, str],
        num_experts: int,
    ) -> torch.Tensor:
        """Estimate per-expert routing weights using router matrix norms (data-free)."""
        # Default: uniform
        w = torch.full((num_experts,), 1.0 / max(1, num_experts), dtype=torch.float32)

        # Candidate router keys (common HF names)
        router_keys: List[str] = []
        if self._teacher_has_mlp_experts:
            router_keys.append(f"{layer_prefix}.{layer_idx}.mlp.gate.weight")
        if self._teacher_has_block_sparse_moe:
            router_keys.append(f"{layer_prefix}.{layer_idx}.block_sparse_moe.gate.weight")
        # fallback variants
        if self._teacher_has_mlp_experts:
            router_keys.append(f"{layer_prefix}.{layer_idx}.mlp.router.weight")
        if self._teacher_has_block_sparse_moe:
            router_keys.append(f"{layer_prefix}.{layer_idx}.block_sparse_moe.router.weight")

        router_key = next((k for k in router_keys if k in teacher_weight_map), None)
        if router_key is None:
            return w

        try:
            t = self.load_tensors_with_prefetch([router_key], self.config.teacher_path, teacher_weight_map)[router_key]
            t = torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0).float()
            if t.dim() == 2 and t.shape[0] == num_experts:
                scores = torch.linalg.vector_norm(t, ord=2, dim=1)
                # Softmax on scores for stability
                w = torch.softmax(scores, dim=0)
        except Exception:
            # Keep uniform
            pass

        return w

    def _synthesize_moe_dense_part(
        self,
        layer_prefix: str,
        layer_idx: int,
        part: str,
        teacher_weight_map: Dict[str, str],
    ) -> Optional[torch.Tensor]:
        """Synthesize a dense MLP matrix from teacher MoE experts for a single layer and part.

        part is one of: 'gate_proj', 'up_proj', 'down_proj'
        """
        self._init_teacher_moe_style(teacher_weight_map)
        if self._teacher_moe_style in (None, "none"):
            return None

        num_experts = int(getattr(self.teacher_arch, "num_experts", 0) or 0)
        if num_experts <= 0:
            # Try infer from keys (robust fallback)
            max_e = -1
            rx = re.compile(rf"^{re.escape(layer_prefix)}\.{layer_idx}\.(?:mlp\.experts|block_sparse_moe\.experts)\.(\d+)\.")
            for k in teacher_weight_map.keys():
                m = rx.match(k)
                if m:
                    max_e = max(max_e, int(m.group(1)))
            num_experts = max_e + 1
        if num_experts <= 0:
            return None

        # Cache key
        cache_key = (layer_prefix, int(layer_idx), str(part))
        if cache_key in self._moe_dense_cache:
            return self._moe_dense_cache[cache_key]

        # Map dense part -> possible expert param names
        part_candidates: List[str] = []
        if part == "gate_proj":
            part_candidates = ["gate_proj", "w1"]
        elif part == "up_proj":
            part_candidates = ["up_proj", "w3"]
        elif part == "down_proj":
            part_candidates = ["down_proj", "w2"]
        else:
            return None

        # Estimate routing weights (data-free, router norm)
        w = self._estimate_moe_expert_weights(layer_prefix, layer_idx, teacher_weight_map, num_experts)

        # Top-k experts (use existing config knob max_experts_to_blend; if <=0, keep all)
        topk = int(self.config.max_experts_to_blend or 0)
        if topk <= 0:
            topk = num_experts
        topk = min(topk, num_experts)

        # Select experts
        if topk < num_experts:
            vals, idxs = torch.topk(w, k=topk, largest=True)
            sel_experts = idxs.tolist()
            sel_w = (vals / (vals.sum() + 1e-12)).tolist()
        else:
            sel_experts = list(range(num_experts))
            sel_w = (w / (w.sum() + 1e-12)).tolist()

        # Build actual keys to load for selected experts
        expert_keys: List[str] = []
        expert_key_for: Dict[int, str] = {}

        def _first_existing(cands: List[str]) -> Optional[str]:
            for kk in cands:
                if kk in teacher_weight_map:
                    return kk
            return None

        for e in sel_experts:
            cands: List[str] = []
            if self._teacher_has_mlp_experts:
                for pn in part_candidates:
                    cands.append(f"{layer_prefix}.{layer_idx}.mlp.experts.{e}.{pn}.weight")
            if self._teacher_has_block_sparse_moe:
                for pn in part_candidates:
                    cands.append(f"{layer_prefix}.{layer_idx}.block_sparse_moe.experts.{e}.{pn}.weight")
            kk = _first_existing(cands)
            if kk is None:
                continue
            expert_key_for[e] = kk
            expert_keys.append(kk)

        if not expert_keys:
            return None

        # Load expert weights (streamed accumulation)
        tensors = self.load_tensors_with_prefetch(expert_keys, self.config.teacher_path, teacher_weight_map)

        acc: Optional[torch.Tensor] = None
        total_w = 0.0
        for e, we in zip(sel_experts, sel_w):
            kk = expert_key_for.get(e)
            if kk is None:
                continue
            t = tensors.get(kk, None)
            if t is None:
                continue
            # Accumulate in fp32 on CPU (keeps VRAM low, maximizes precision)
            t = torch.nan_to_num(t, nan=0.0, posinf=0.0, neginf=0.0).float()
            if acc is None:
                acc = t.mul(float(we))
            else:
                acc.add_(t, alpha=float(we))
            total_w += float(we)

        if acc is None or total_w <= 0.0:
            return None

        # Normalize in case some experts were missing
        if abs(total_w - 1.0) > 1e-3:
            acc.div_(total_w)

        self._moe_dense_cache[cache_key] = acc
        return acc

    def _maybe_synthesize_teacher_from_moe(
        self,
        layer_prefix: str,
        t_floor: int,
        t_ceil: int,
        mapped_rest: str,
        teacher_weight_map: Dict[str, str],
    ) -> Dict[str, torch.Tensor]:
        """If teacher is MoE and student expects dense MLP weights, synthesize teacher matrices."""
        if self.teacher_arch.mlp_type != MLPType.MOE:
            return {}
        if self.student_arch.mlp_type == MLPType.MOE:
            return {}

        # Only for standard MLP projections
        part = None
        if ".mlp.gate_proj." in f".{mapped_rest}.":
            part = "gate_proj"
        elif ".mlp.up_proj." in f".{mapped_rest}.":
            part = "up_proj"
        elif ".mlp.down_proj." in f".{mapped_rest}.":
            part = "down_proj"
        else:
            return {}

        out: Dict[str, torch.Tensor] = {}
        tf = self._synthesize_moe_dense_part(layer_prefix, int(t_floor), part, teacher_weight_map)
        if tf is not None:
            out[f"{layer_prefix}.{t_floor}.<moe_synth>.mlp.{part}.weight"] = tf
        if int(t_ceil) != int(t_floor):
            tc = self._synthesize_moe_dense_part(layer_prefix, int(t_ceil), part, teacher_weight_map)
            if tc is not None:
                out[f"{layer_prefix}.{t_ceil}.<moe_synth>.mlp.{part}.weight"] = tc

        return out


    def process_layer(
        self,
        student_key: str,
        student_tensor: torch.Tensor,
        teacher_tensors: Dict[str, torch.Tensor],
        interp_w: float,
    ) -> Optional[Tuple[torch.Tensor, torch.Tensor, int]]:
        """Process a single layer: compute delta, apply pruning, extract LoRA."""
        if not teacher_tensors:
            return None

        teacher_keys = list(teacher_tensors.keys())
        student_on_device = student_tensor.to(self.device, non_blocking=True)
        target_shape = student_tensor.shape

        # ---------------------------------------------------------------------
        # Cross-architecture safe path:
        #   When projection is needed (teacher matrix shape != student matrix shape),
        #   do NOT compute delta = (projected_teacher - student_full).
        #   That subtracts the student's residual space and pushes the weight toward
        #   a low-rank matrix, which is a common cause of degenerate/infinite output.
        #
        #   Instead:
        #     1) compute student's truncated SVD basis (U_s, S_s, V_s)
        #     2) compute teacher core in that basis: core_t = U_s^T Wt_resized V_s
        #     3) compute student core: core_s = diag(S_s)
        #     4) delta_core = core_t - core_s   (only inside the shared subspace)
        #     5) extract LoRA directly from delta_core without forming the full delta
        # ---------------------------------------------------------------------
        try_subspace = (
            self.config.use_svd_projection
            and self.is_cross_arch
            and (getattr(self.config, "subspace_delta", "off") != "off")
            and student_on_device.dim() == 2
            and len(target_shape) == 2
            and any(tuple(t.shape) != tuple(target_shape) for t in teacher_tensors.values())
        )

        if try_subspace:
            try:
                # Student basis (common space)
                U_s, S_s, V_s = compute_student_svd_basis(
                    student_on_device, self.svd_engine, self.config, self.device
                )
                k = int(S_s.numel())
                diag_Ss = torch.diag(S_s).to(device=self.device, dtype=torch.float32)

                # Compute teacher cores and build delta in core space
                if len(teacher_keys) == 1:
                    core_t = teacher_core_in_student_basis(
                        teacher_tensors[teacher_keys[0]],
                        U_s, V_s,
                        self.svd_engine, self.config, self.device
                    )
                    delta_core = core_t - diag_Ss
                    if self.config.use_dare:
                        delta_core = apply_dare(delta_core, self.config.dare_drop_rate)
                    if self.config.use_ties and getattr(self.config, "ties_trim_single", True):
                        delta_core = apply_ties_single(delta_core, self.config.ties_density)
                else:
                    core_floor = teacher_core_in_student_basis(
                        teacher_tensors[teacher_keys[0]],
                        U_s, V_s,
                        self.svd_engine, self.config, self.device
                    )
                    core_ceil = teacher_core_in_student_basis(
                        teacher_tensors[teacher_keys[1]],
                        U_s, V_s,
                        self.svd_engine, self.config, self.device
                    )
                    delta_floor = core_floor - diag_Ss
                    delta_ceil = core_ceil - diag_Ss

                    mix_w = float(interp_w)
                    if getattr(self.config, "interp_mode", "schedule") == "lsq":
                        mix_w = lsq_mix_weight_from_deltas(delta_floor, delta_ceil, mix_w)

                    if self.config.use_ties:
                        deltas = [delta_floor * (1.0 - mix_w), delta_ceil * mix_w]
                        if self.config.use_dare:
                            deltas = [apply_dare(d, self.config.dare_drop_rate) for d in deltas]
                        delta_core = apply_ties(deltas, self.config.ties_density)
                    else:
                        core_t = core_floor * (1.0 - mix_w) + core_ceil * mix_w
                        delta_core = core_t - diag_Ss
                        if self.config.use_dare:
                            delta_core = apply_dare(delta_core, self.config.dare_drop_rate)

                # Subspace-delta stabilization / mode control
                mode = getattr(self.config, "subspace_delta", "off")

                # Optional: match teacher core scale to student core scale
                if getattr(self.config, "subspace_norm_match", True):
                    core_s_norm = torch.norm(diag_Ss.float()) + 1e-8
                    core_t = (delta_core + diag_Ss).to(dtype=torch.float32)
                    core_t_norm = torch.norm(core_t) + 1e-8
                    scale_t = (core_s_norm / core_t_norm).clamp(0.25, 4.0)
                    core_t = core_t * scale_t
                    delta_core = (core_t - diag_Ss).to(dtype=delta_core.dtype)

                # Optional: suppress off-diagonal energy (improves stability)
                if mode in ("diag", "auto"):
                    diag_part = torch.diag(torch.diag(delta_core))
                    off_part = delta_core - diag_part
                    off_frac = float((torch.norm(off_part.float()) / (torch.norm(delta_core.float()) + 1e-8)).item())
                    thr = float(getattr(self.config, "subspace_offdiag_threshold", 0.55))
                    if mode == "diag" or (mode == "auto" and off_frac > thr):
                        delta_core = diag_part

                # Delta ratio clipping (use ||delta_core|| since U/V are orthonormal)
                if getattr(self.config, "max_delta_ratio", 0.0) and self.config.max_delta_ratio > 0:
                    denom_mode = getattr(self.config, "subspace_denom", "core")
                    if denom_mode == "full":
                        denom = torch.norm(student_on_device.float()) + 1e-8
                    else:
                        denom = torch.norm(diag_Ss.float()) + 1e-8
                    ratio = torch.norm(delta_core.float()) / denom
                    if ratio > self.config.max_delta_ratio:
                        scale_t = (self.config.max_delta_ratio / (ratio + 1e-12)).clamp(max=1.0)
                        scale = float(scale_t.item())
                        delta_core = delta_core * scale
                        print(
                            f"    [delta-clip] scaled by {scale:.4f} (ratio {float(ratio.item()):.4f} -> {self.config.max_delta_ratio:.4f})"
                        )

                # Rank selection hint (still respected when adaptive rank is off)
                is_attn = any(proj in student_key for proj in ["q_proj", "k_proj", "v_proj", "o_proj"])
                is_mlp = any(proj in student_key for proj in ["gate_proj", "up_proj", "down_proj"])

                rank = self.config.rank_default
                if is_attn and self.config.rank_attn is not None:
                    rank = self.config.rank_attn
                elif is_mlp and self.config.rank_mlp is not None:
                    rank = self.config.rank_mlp

                # Output dtype: match student for plug-and-play inference
                out_dtype = student_on_device.dtype if self.config.use_mixed_precision else torch.float32

                lora_A, lora_B, actual_rank = extract_lora_from_core_delta(
                    delta_core, U_s, V_s, int(rank), self.config, out_dtype
                )
                if lora_A is None or lora_B is None or actual_rank <= 0:
                    return None

                # Ensure shapes are correct
                if lora_A.shape[1] != target_shape[1] or lora_B.shape[0] != target_shape[0]:
                    # Fallback if something unexpected happened
                    return None

                return lora_A, lora_B, int(actual_rank)

            except Exception:
                # Fall back to legacy behavior
                pass

        # ---------------------------------------------------------------------
        # Legacy path: teacher and student already share shape, or projection is off.
        # ---------------------------------------------------------------------
        deltas = []
        for teacher_key in teacher_keys:
            teacher_tensor = teacher_tensors[teacher_key].to(self.device)

            # Cross-architecture projection (legacy): only used when caller didn't
            # trigger the safer subspace path above.
            if self.config.use_svd_projection and self.is_cross_arch:
                teacher_tensor = svd_projection(
                    teacher_tensor, target_shape, self.svd_engine, self.config, self.device, ref=student_on_device
                )
            else:
                teacher_tensor = self._project_tensor(teacher_tensor, target_shape)

            delta = teacher_tensor.to(self.device) - student_on_device

            deltas.append(delta)

        # Apply interpolation or TIES between floor/ceil deltas
        mix_w = float(interp_w)
        if len(deltas) == 2 and getattr(self.config, "interp_mode", "schedule") == "lsq":
            mix_w = lsq_mix_weight_from_deltas(deltas[0], deltas[1], mix_w)
        if len(deltas) == 1:
            final_delta = deltas[0]
            if self.config.use_dare:
                final_delta = apply_dare(final_delta, self.config.dare_drop_rate)
            if self.config.use_ties and getattr(self.config, "ties_trim_single", True):
                final_delta = apply_ties_single(final_delta, self.config.ties_density)
        else:
            if self.config.use_ties:
                weighted_deltas = [deltas[0] * (1 - mix_w), deltas[1] * mix_w]
                if self.config.use_dare:
                    weighted_deltas = [apply_dare(d, self.config.dare_drop_rate) for d in weighted_deltas]
                final_delta = apply_ties(weighted_deltas, self.config.ties_density)
            else:
                final_delta = deltas[0] * (1 - mix_w) + deltas[1] * mix_w
                if self.config.use_dare:
                    final_delta = apply_dare(final_delta, self.config.dare_drop_rate)


        # Final delta ratio clipping (after interpolation/TIES/DARE)
        if getattr(self.config, "max_delta_ratio", 0.0) and self.config.max_delta_ratio > 0:
            denom = torch.norm(student_on_device.float()) + 1e-8
            ratio = torch.norm(final_delta.float()) / denom
            if ratio > self.config.max_delta_ratio:
                scale_t = (self.config.max_delta_ratio / (ratio + 1e-12)).clamp(max=1.0)
                scale = float(scale_t.item())
                final_delta = final_delta * scale
                print(
                    f"    [delta-clip] scaled by {scale:.4f} (ratio {float(ratio.item()):.4f} -> {self.config.max_delta_ratio:.4f})"
                )

        # Extract LoRA from final delta
        is_attn = any(proj in student_key for proj in ["q_proj", "k_proj", "v_proj", "o_proj"])
        is_mlp = any(proj in student_key for proj in ["gate_proj", "up_proj", "down_proj"])

        rank = self.config.rank_default
        if is_attn and self.config.rank_attn is not None:
            rank = self.config.rank_attn
        elif is_mlp and self.config.rank_mlp is not None:
            rank = self.config.rank_mlp

        # LoRA extraction (PiSSA). Optionally apply calibration-aware diagonal whitening.
        module_name = student_key[:-len(".weight")] if student_key.endswith(".weight") else student_key
        if getattr(self.config, "calibration_mode", "none") != "none" and self.calib_stats:
            lora_A, lora_B, actual_rank = extract_lora_pissa_calibrated(
                final_delta, rank, self.svd_engine, self.config, module_name, self.calib_stats, self.student_arch
            )
        else:
            lora_A, lora_B, actual_rank = extract_lora_pissa(final_delta, rank, self.svd_engine, self.config)

        if lora_A is None:
            return None

        return lora_A, lora_B, actual_rank
    def _project_tensor(self, src: torch.Tensor, 
                        target_shape: Tuple[int, ...]) -> torch.Tensor:
        """Simple projection (fallback)."""
        if src.shape == target_shape:
            return src
        
        if src.dim() == 1:
            out = torch.zeros(target_shape, device=src.device, dtype=src.dtype)
            copy_len = min(src.numel(), target_shape[0])
            out[:copy_len] = src[:copy_len]
            return out
        
        if src.dim() == 2:
            out = torch.zeros(target_shape, device=src.device, dtype=src.dtype)
            copy_rows = min(src.shape[0], target_shape[0])
            copy_cols = min(src.shape[1], target_shape[1])
            out[:copy_rows, :copy_cols] = src[:copy_rows, :copy_cols]
            return out
        
        return torch.zeros(target_shape, device=src.device, dtype=src.dtype)
    
    def run(self,
            student_keys: List[str],
            teacher_weight_map: Dict[str, str],
            student_weight_map: Dict[str, str],
            teacher_layers_map: Dict[str, List[int]],
            student_layers_map: Dict[str, List[int]],
            teacher_arch: ArchitectureInfo,
            student_arch: ArchitectureInfo,
            layer_name_mapping: Dict[str, str]) -> Tuple[Dict[str, torch.Tensor], Dict[str, int]]:
        """Run accelerated distillation."""
        s_layers = max(len(v) for v in student_layers_map.values()) if student_layers_map else 32
        t_layers = max(len(v) for v in teacher_layers_map.values()) if teacher_layers_map else 32
        
        # FIX #4: Skip embeddings for cross-arch with different vocab/hidden
        skip_embed_warning_shown = False
        
        processable_keys = []
        for k in student_keys:
            if not k.endswith(".weight"):
                continue
            if any(k.endswith(s) for s in [".bias"]):
                continue
            if any(tok in k.lower() for tok in NORM_TOKENS):
                continue
            
            # Special handling for embeddings in cross-arch
            if any(tok in k for tok in EMBED_TOKENS):
                if not self.config.include_embed_lm_head:
                    continue
                if self.is_cross_arch:
                    if not skip_embed_warning_shown and self.config.verbose:
                        print(f"\n⚠️  Skipping embeddings for cross-architecture distillation")
                        print(f"   (Teacher hidden={teacher_arch.hidden_size}, Student hidden={student_arch.hidden_size})")
                        skip_embed_warning_shown = True
                    continue
            
            if self.config.include_pattern and not re.search(self.config.include_pattern, k):
                continue
            if self.config.exclude_pattern and re.search(self.config.exclude_pattern, k):
                continue
            processable_keys.append(k)
        
        keys_by_shard = defaultdict(list)
        for k in processable_keys:
            if k in student_weight_map:
                keys_by_shard[student_weight_map[k]].append(k)
        
        pbar = tqdm(total=len(processable_keys),
                    desc="Distilling",
                    dynamic_ncols=True,
                    disable=not self.config.verbose)
        
        for shard_file, shard_keys in sorted(keys_by_shard.items()):
            pbar.set_description(f"Processing {os.path.basename(shard_file)}")
            
            for i in range(0, len(shard_keys), self.config.micro_batch_size):
                micro_keys = shard_keys[i:i + self.config.micro_batch_size]
                next_keys = shard_keys[i + self.config.micro_batch_size:i + 2*self.config.micro_batch_size]
                
                student_tensors = self.load_tensors_with_prefetch(
                    micro_keys, self.config.student_path, student_weight_map, next_keys
                )
                
                for sk in micro_keys:
                    pbar.update(1)
                    
                    st = student_tensors.get(sk)
                    if st is None:
                        continue
                    
                    seg = split_key(sk)
                    if seg is None:
                        continue
                    
                    s_prefix, token_name, s_idx, rest = seg
                    
                    t_floor, interp_w = teacher_idx_from_student_idx(
                        s_idx, s_layers, t_layers, 
                        self.config.map_schedule, self.config.sigmoid_k
                    )
                    t_ceil = min(t_floor + 1, t_layers - 1)
                    
                    mapped_rest = rest
                    for s_name, t_name in layer_name_mapping.items():
                        mapped_rest = mapped_rest.replace(s_name, t_name)
                    
                    t_prefix = s_prefix
                    key_floor = f"{t_prefix}.{t_floor}.{mapped_rest}"
                    key_ceil = f"{t_prefix}.{t_ceil}.{mapped_rest}"
                    
                    teacher_keys_to_load = []
                    for tk in [key_floor, key_ceil]:
                        if tk in teacher_weight_map:
                            teacher_keys_to_load.append(tk)
                    
                    teacher_tensors: Dict[str, torch.Tensor] = {}
                    if teacher_keys_to_load:
                        teacher_tensors = self.load_tensors_with_prefetch(
                            teacher_keys_to_load, self.config.teacher_path, teacher_weight_map
                        )
                    else:
                        teacher_tensors = self._maybe_synthesize_teacher_from_moe(
                            s_prefix, int(t_floor), int(t_ceil), mapped_rest, teacher_weight_map
                        )
                        if not teacher_tensors:
                            continue

                    result = self.process_layer(sk, st, teacher_tensors, interp_w)
                    
                    if result is not None:
                        lora_A, lora_B, actual_rank = result
                        
                        peft_key = to_peft_key(sk)
                        
                        self.lora_weights[f"{peft_key}.lora_A.default.weight"] = lora_A.detach().cpu().contiguous()
                        self.lora_weights[f"{peft_key}.lora_B.default.weight"] = lora_B.detach().cpu().contiguous()

                        self.rank_info[peft_key] = actual_rank
                        self.stats["processed"] += 1
                    
                    del teacher_tensors
                
                del student_tensors
                if i % (self.config.micro_batch_size * 4) == 0:
                    self.memory_manager.clear_cache()
        
        pbar.close()
        
        self.memory_manager.clear_cache()
        self.offloader.cleanup()
        
        return self.lora_weights, self.rank_info


# =============================================================================
#                              MAIN
# =============================================================================

def build_argparser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="Universal SVD-LoRA Distillation v4.2.0 (Cross-Arch Innovations) — Plug-and-Play Truncated SVD (Low-VRAM)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
v3.7 highlights:
  - Truncated / randomized SVD by default in 'auto' mode (avoids full SVD on large matrices).
  - Efficient SVD-based projection for cross-architecture distillation with bounded projection rank.
  - Keeps shard-wise streaming + offload (low VRAM) behavior from v3.2.

Examples:
  # Cross-architecture distillation (e.g., 72B -> 36B)
  python %(prog)s --teacher /path/to/teacher --student /path/to/student --output /path/to/lora \
      --adaptive-rank --min-rank 256 --max-rank 2048 --energy-threshold 0.95 \
      --svd-mode auto --projection-rank 256 --svd-projection \
      --ties --ties-density 0.3 --no-dare \
      --include "self_attn|mlp" --offload-strategy cpu --mixed-precision
"""
    )

    # Required
    p.add_argument("--teacher", required=True, help="Teacher model folder (HF safetensors)")
    p.add_argument("--student", required=True, help="Student model folder (HF safetensors)")
    p.add_argument("--output", required=True, help="Output directory for PEFT adapter")

    # Rank controls
    p.add_argument("--rank", type=int, default=64, help="Default LoRA rank (if not using rank-attn/mlp)")
    p.add_argument("--rank-attn", type=int, default=None, help="Override LoRA rank for attention projections")
    p.add_argument("--rank-mlp", type=int, default=None, help="Override LoRA rank for MLP projections")
    p.add_argument("--lora-alpha", type=int, default=None, help="LoRA alpha; default=r (scale=1)")
    p.add_argument("--alpha-mode", type=str, default="auto", choices=["auto","rank","fixed"], help="LoRA alpha policy: auto=rank if --lora-alpha not set; rank=alpha_pattern=r (scale~1); fixed=global alpha only")

    p.add_argument("--adaptive-rank", action="store_true", default=True, help="Enable per-layer adaptive rank")
    p.add_argument("--no-adaptive-rank", action="store_false", dest="adaptive_rank")
    p.add_argument("--energy-threshold", type=float, default=0.95, help="Energy threshold for adaptive rank (0-1)")
    p.add_argument("--min-rank", type=int, default=8, help="Minimum rank when adaptive rank is enabled")
    p.add_argument("--max-rank", type=int, default=256, help="Maximum rank when adaptive rank is enabled")

    # Layer mapping
    p.add_argument("--map-schedule", choices=["linear", "sigmoid"], default="linear",
                   help="Teacher->student layer mapping schedule")
    p.add_argument("--sigmoid-k", type=float, default=0.15, help="Sigmoid schedule slope")
    p.add_argument("--interp-mode", choices=["schedule", "lsq"], default="schedule",
                   help="Teacher floor/ceil mix: schedule=use mapping schedule weight; lsq=per-matrix least-squares mix weight in student space (clamped 0..1).")

    # Merge / regularization
    # DARE (off by default in v3.7)
    p.add_argument("--dare", action="store_true", default=False, help="Enable DARE (dropout on delta)")
    p.add_argument("--no-dare", action="store_false", dest="dare")
    p.add_argument("--dare-drop-rate", type=float, default=0.3, help="DARE dropout probability")
    p.add_argument("--auto-adjust-dare", action="store_true", default=True,
                   help="Auto-adjust DARE for cross-architecture (when DARE enabled)")
    p.add_argument("--no-auto-adjust-dare", action="store_false", dest="auto_adjust_dare")

    # TIES (on by default)
    p.add_argument("--ties", action="store_true", default=True, help="Enable TIES merging when combining deltas")
    p.add_argument("--no-ties", action="store_false", dest="ties")
    p.add_argument("--ties-density", type=float, default=0.3, help="Keep ratio for TIES trimming (0-1)")
    p.add_argument("--ties-trim-single", action="store_true", default=True,
                   help="Apply the TIES Trim step to single-delta cases (default: on; matches v3.7.15).")
    p.add_argument("--no-ties-trim-single", action="store_false", dest="ties_trim_single",
                   help="Disable TIES Trim for single-delta cases.")

    # Stability guard
    p.add_argument("--max-delta-ratio", type=float, default=0.35,
                   help="If >0, scales delta when ||delta||/||student|| exceeds this value")

    # SVD backend controls
    svd_group = p.add_argument_group("SVD Options")
    svd_group.add_argument("--svd-mode", choices=["auto", "full", "randomized", "krylov", "adaptive", "drf", "aurora"], default="auto",
                           help="SVD backend: auto|full|randomized|krylov|adaptive|drf|aurora")
    svd_group.add_argument("--svd-rand-iter", type=int, default=2,
                           help="Low-rank iterations: randomized=subspace iters (Alg. 4.4), krylov=depth, adaptive=per-block iters")
    svd_group.add_argument("--svd-rand-oversamples", type=int, default=8,
                           help="Low-rank oversamples (k = rank + oversamples) for randomized/krylov")
    svd_group.add_argument("--svd-auto-min-dim", type=int, default=1024,
                           help="In auto mode, prefer low-rank SVD when min(m,n) >= this")
    svd_group.add_argument("--svd-auto-full-rank-ratio", type=float, default=0.6,
                           help="In auto mode, use full SVD when rank >= ratio*min(m,n)")
    svd_group.add_argument("--svd-auto-lowrank", choices=["randomized", "krylov", "drf", "aurora"], default="randomized",
                           help="In auto mode (when low-rank is selected), choose randomized, krylov, drf, or aurora")
    svd_group.add_argument("--svd-adaptive-block-size", type=int, default=32,
                           help="(adaptive) columns added per adaptive iteration (Alg. 4.2)")
    svd_group.add_argument("--svd-adaptive-n-test", type=int, default=8,
                           help="(adaptive) test vectors for residual estimate (Alg. 4.2)")

    svd_group.add_argument("--svd-drf-steps", type=int, default=1,
                       help="(drf) Number of dual-residual folding refinement rounds (>=0).")
    svd_group.add_argument("--svd-drf-theta", type=float, default=0.5,
                       help="(drf) Residual scaling exponent in [0,1]. 0=no scaling; 1=full 1/sigma scaling.")
    svd_group.add_argument("--svd-drf-resid-eps", type=float, default=1e-8,
                       help="(drf) Numerical floor for sigma in residual scaling.")

    svd_group.add_argument("--svd-aurora-steps", type=int, default=1,
                       help="(aurora) Number of AURORA refinement rounds (>=0).")
    svd_group.add_argument("--svd-aurora-order", type=int, choices=[1, 2], default=2,
                       help="(aurora) Residual order: 1=first-order residual only; 2=adds second-order residual response via A/A^T.")
    svd_group.add_argument("--svd-aurora-theta1", type=float, default=0.5,
                       help="(aurora) Level-1 residual scaling exponent in [0,1]. 0=no scaling; 1=full 1/sigma scaling.")
    svd_group.add_argument("--svd-aurora-theta2", type=float, default=0.25,
                       help="(aurora) Level-2 residual scaling exponent in [0,1].")
    svd_group.add_argument("--svd-aurora-resid-eps", type=float, default=1e-8,
                       help="(aurora) Numerical floor for sigma in residual scaling.")
    svd_group.add_argument("--svd-aurora-level2-keep", type=float, default=1.0,
                       help="(aurora) Fraction (0..1] of components kept for the level-2 term; <1 keeps only the most residual-dominant components.")

    # Cross-architecture projection controls
    cross_group = p.add_argument_group("Cross-Architecture Projection")
    cross_group.add_argument("--svd-projection", action="store_true", default=True,
                             help="Enable SVD-based teacher->student tensor projection")
    cross_group.add_argument("--no-svd-projection", action="store_false", dest="svd_projection")
    cross_group.add_argument("--projection-rank", type=int, default=256,
                             help="Max rank used for SVD projection (bounded for cost control)")
    cross_group.add_argument("--projection-min-rank", type=int, default=8,
                             help="Min rank for projection when projection-adaptive-rank is enabled")
    cross_group.add_argument("--projection-adaptive-rank", action="store_true", default=True,
                             help="Enable energy-based adaptive rank for projection SVD")
    cross_group.add_argument("--no-projection-adaptive-rank", action="store_false", dest="projection_adaptive_rank")
    cross_group.add_argument("--projection-energy-threshold", type=float, default=0.99,
                             help="Energy threshold for projection adaptive rank (0-1)")
    cross_group.add_argument("--projection-rand-iter", type=int, default=2,
                             help="Randomized SVD power iterations for projection")
    cross_group.add_argument("--projection-rand-oversamples", type=int, default=8,
                             help="Randomized SVD oversamples for projection")


    cross_group.add_argument("--subspace-delta", type=str, default="off",
                             choices=["off", "auto", "diag", "full"],
                             help="Cross-arch: subspace-delta mode. off=disable (use full-delta), diag=diagonal core only, full=full core, auto=diag if off-diagonal energy is high")
    cross_group.add_argument("--subspace-offdiag-threshold", type=float, default=0.55,
                             help="subspace-delta auto: if off-diagonal Frobenius fraction > threshold, fall back to diag")
    cross_group.add_argument("--subspace-denom", type=str, default="core", choices=["core", "full"],
                             help="subspace-delta delta-clip denominator: core=||diag(S_student)||, full=||W_student||")
    cross_group.add_argument("--no-subspace-norm-match", action="store_true",
                             help="Disable subspace core norm matching (scale teacher core to student core norm)")
    # Performance / memory
    perf = p.add_argument_group("Performance / Memory")
    perf.add_argument("--num-gpus", type=int, default=1)
    perf.add_argument("--gpu-memory-fraction", type=float, default=0.95)
    perf.add_argument("--no-cuda-streams", action="store_true", help="Disable CUDA streams")
    perf.add_argument("--cuda-streams", type=int, default=3)
    perf.add_argument("--svd-batch-size", type=int, default=1)
    perf.add_argument("--micro-batch", type=int, default=8, help="Tensors per IO micro-batch")
    perf.add_argument("--prefetch-layers", type=int, default=2)

    perf.add_argument("--offload-strategy", choices=["cpu", "disk", "none"], default="cpu")
    perf.add_argument("--offload-dir", type=str, default=None)
    perf.add_argument("--pinned-memory", action="store_true", default=True)
    perf.add_argument("--no-pinned-memory", action="store_false", dest="pinned_memory")
    perf.add_argument("--max-cpu-memory", type=float, default=96.0)

    perf.add_argument("--chunk-threshold", type=int, default=512, help="MB: use randomized SVD above this size")
    perf.add_argument("--chunk-size", type=int, default=128, help="MB: chunk size (reserved for future use)")

    # Precision
    p.add_argument("--mixed-precision", action="store_true", default=True)
    p.add_argument("--no-mixed-precision", action="store_false", dest="mixed_precision")

    # Module selection
    p.add_argument("--include", type=str, default="self_attn|mlp",
                   help="Regex include filter for tensor keys")
    p.add_argument("--exclude", type=str, default="",
                   help="Regex exclude filter for tensor keys")
    p.add_argument("--include-embed", action="store_true", default=False,
                   help="Include embedding and lm_head weights (only safe when shapes match)")


    # Calibration-aware SVD (optional)
    calib = p.add_argument_group("Calibration-aware SVD (optional)")
    calib.add_argument("--calib-mode", choices=["none", "rms", "cov"], default="none",
                       help="Enable calibration-aware whitening before SVD. none=disable, rms=diagonal RMS whitening, cov=full-covariance (SVD-LLM-style) whitening.")
    calib.add_argument("--calib-data", type=str, default=None,
                       help="Calibration prompts file: .txt (1 prompt/line) or .jsonl ({\"text\":...}). If omitted, a small built-in prompt set is used.")
    calib.add_argument("--calib-format", choices=["auto", "txt", "jsonl", "json", "alpaca"], default="auto",
                       help="Calibration file format (auto uses extension).")
    calib.add_argument("--calib-alpaca-template", choices=["classic", "plain"], default="classic",
                       help="Alpaca prompt template for calibration (only for calib-format=alpaca/auto).")
    calib.add_argument("--calib-alpaca-include-output", action="store_true",
                       help="Include Alpaca output field in calibration forward (alpaca only).")
    calib.add_argument("--calib-max-samples", type=int, default=128,
                       help="Max number of calibration prompts to use.")
    calib.add_argument("--calib-max-length", type=int, default=256,
                       help="Max token length per prompt (tokenizer truncation).")
    calib.add_argument("--calib-padding", type=str, default="longest", choices=["longest","max_length"],
                       help="Padding strategy for calibration tokenization (longest minimizes padding).")
    calib.add_argument("--calib-no-attention-mask", action="store_true",
                       help="Disable attention_mask-based padding exclusion in calibration stats.")
    calib.add_argument("--calib-batch-size", type=int, default=1,
                       help="Batch size for calibration forward passes.")
    calib.add_argument("--calib-device", choices=["auto", "cuda", "cpu"], default="auto",
                       help="Device placement for calibration model forward (auto=cuda if available).")
    calib.add_argument("--calib-dtype", choices=["auto", "bf16", "fp16", "fp32"], default="auto",
                       help="dtype for calibration forward (auto: bf16 if supported on cuda else fp16; cpu->fp32).")
    calib.add_argument("--calib-load", type=str, default=None,
                       help="Load precomputed calibration stats (.safetensors). Skips forward calibration.")
    calib.add_argument("--calib-save", type=str, default=None,
                       help="Save calibration stats (.safetensors) for reuse.")
    calib.add_argument("--calib-eps", type=float, default=1e-6,
                       help="Epsilon for whitening (clamp min).")
    calib.add_argument("--calib-cov-chunk-layers", type=int, default=4,
                       help="(cov mode) Number of layers per calibration pass (controls memory).")
    calib.add_argument("--calib-cov-max-dim", type=int, default=8192,
                       help="(cov mode) Max in_features dimension to collect full covariance (skip larger dims).")
    calib.add_argument("--calib-cov-groups", type=str, default="qkv,mlp,o",
                       help="(cov mode) Comma-separated groups to collect: qkv,o,mlp,down.")
    calib.add_argument("--calib-cov-store-dtype", choices=["fp16", "bf16", "fp32"], default="fp16",
                       help="(cov mode) dtype for storing Cholesky factors in calib-save.")
    calib.add_argument("--calib-no-in", action="store_true",
                       help="Disable input RMS collection (not recommended).")
    calib.add_argument("--calib-no-out", action="store_true",
                       help="Disable output RMS collection (reduces overhead; input-only whitening).")
    # MoE (kept for compatibility)
    p.add_argument("--moe-method", type=str, default="none")
    p.add_argument("--max-experts", type=int, default=2)

    # Misc
    p.add_argument("--seed", type=int, default=1234)
    p.add_argument("--quiet", action="store_true")

    return p


def main():
    args = build_argparser().parse_args()
    
    config = DistillConfig(
        teacher_path=args.teacher,
        student_path=args.student,
        output_path=args.output,
        rank_default=args.rank,
        rank_attn=args.rank_attn,
        rank_mlp=args.rank_mlp,
        lora_alpha=args.lora_alpha,
        alpha_mode=args.alpha_mode,
        use_adaptive_rank=args.adaptive_rank,
        energy_threshold=args.energy_threshold,
        min_rank=args.min_rank,
        max_rank=args.max_rank,
        use_dare=args.dare,
        dare_drop_rate=args.dare_drop_rate,
        use_ties=args.ties,
        ties_density=args.ties_density,
        ties_trim_single=args.ties_trim_single,
        max_delta_ratio=args.max_delta_ratio,

                # Calibration-aware SVD
        calibration_mode=args.calib_mode,
        calib_data=args.calib_data,
        calib_format=args.calib_format,
        calib_alpaca_template=args.calib_alpaca_template,
        calib_alpaca_include_output=bool(args.calib_alpaca_include_output),
        calib_max_samples=args.calib_max_samples,
        calib_max_length=args.calib_max_length,
        calib_padding=args.calib_padding,
        calib_use_attention_mask=(not args.calib_no_attention_mask),
        calib_batch_size=args.calib_batch_size,
        calib_device=args.calib_device,
        calib_dtype=args.calib_dtype,
        calib_load=args.calib_load,
        calib_save=args.calib_save,
        calib_eps=args.calib_eps,
        calib_cov_chunk_layers=args.calib_cov_chunk_layers,
        calib_cov_max_dim=args.calib_cov_max_dim,
        calib_cov_groups=args.calib_cov_groups,
        calib_cov_store_dtype=args.calib_cov_store_dtype,
        calib_collect_in=(not args.calib_no_in),
        calib_collect_out=(not args.calib_no_out),

moe_merge_method=args.moe_method,
        max_experts_to_blend=args.max_experts,
        map_schedule=args.map_schedule,
        sigmoid_k=args.sigmoid_k,
        interp_mode=args.interp_mode,

        # SVD backend
        svd_mode=args.svd_mode,
        svd_randomized_iter=args.svd_rand_iter,
        svd_randomized_oversamples=args.svd_rand_oversamples,
        svd_auto_min_dim=args.svd_auto_min_dim,
        svd_auto_full_rank_ratio=args.svd_auto_full_rank_ratio,
        svd_auto_lowrank=args.svd_auto_lowrank,
        svd_adaptive_block_size=args.svd_adaptive_block_size,
        svd_adaptive_n_test=args.svd_adaptive_n_test,
        svd_drf_steps=args.svd_drf_steps,
        svd_drf_theta=args.svd_drf_theta,
        svd_drf_resid_eps=args.svd_drf_resid_eps,
        svd_aurora_steps=args.svd_aurora_steps,
        svd_aurora_order=args.svd_aurora_order,
        svd_aurora_theta1=args.svd_aurora_theta1,
        svd_aurora_theta2=args.svd_aurora_theta2,
        svd_aurora_resid_eps=args.svd_aurora_resid_eps,
        svd_aurora_level2_keep=args.svd_aurora_level2_keep,

        # Cross-arch projection
        use_svd_projection=args.svd_projection,
        projection_rank=args.projection_rank,
        projection_min_rank=args.projection_min_rank,
        projection_adaptive_rank=args.projection_adaptive_rank,
        projection_energy_threshold=args.projection_energy_threshold,
        projection_randomized_iter=args.projection_rand_iter,
        projection_randomized_oversamples=args.projection_rand_oversamples,
        subspace_delta=args.subspace_delta,
        subspace_offdiag_threshold=args.subspace_offdiag_threshold,
        subspace_denom=args.subspace_denom,
        subspace_norm_match=(not args.no_subspace_norm_match),

        num_gpus=args.num_gpus,
        gpu_memory_fraction=args.gpu_memory_fraction,
        use_cuda_streams=not args.no_cuda_streams,
        num_cuda_streams=args.cuda_streams,
        use_mixed_precision=args.mixed_precision,
        svd_batch_size=args.svd_batch_size,
        offload_strategy=args.offload_strategy,
        offload_dir=args.offload_dir,
        prefetch_layers=args.prefetch_layers,
        use_pinned_memory=args.pinned_memory,
        max_cpu_memory_gb=args.max_cpu_memory,
        chunk_threshold_mb=args.chunk_threshold,
        chunk_size_mb=args.chunk_size,
        micro_batch_size=args.micro_batch,
        include_pattern=args.include,
        exclude_pattern=args.exclude,
        include_embed_lm_head=args.include_embed,
        auto_adjust_dare=args.auto_adjust_dare,
        seed=args.seed,
        verbose=not args.quiet,
    )
    
    seed_all(config.seed)
    
    print("=" * 80)
    print("Universal SVD-LoRA Distillation v4.2.0 (Cross-Arch Innovations)")
    print("Plug-and-Play Truncated SVD (Low-VRAM) + Low-VRAM Offload")
    print("=" * 80)
    
    if torch.cuda.is_available():
        device = torch.device("cuda:0")
        print(f"\n🖥️  GPU: {torch.cuda.get_device_name(0)}")
        print(f"   Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
        
        try:
            torch.backends.cuda.preferred_linalg_library("cusolver")
        except Exception:
            pass
    else:
        device = torch.device("cpu")
        print("\n⚠️  No GPU available, using CPU")
    
    print("\n📂 Loading model indices...")
    teacher_weight_map = read_index_map(config.teacher_path)
    student_weight_map = read_index_map(config.student_path)
    
    teacher_keys = sorted(teacher_weight_map.keys())
    student_keys = sorted(student_weight_map.keys())
    
    print(f"   Teacher: {len(teacher_keys)} tensors")
    print(f"   Student: {len(student_keys)} tensors")
    
    print("\n🔍 Detecting architectures...")
    if HAS_TRANSFORMERS:
        teacher_arch = detect_architecture_from_config(config.teacher_path)
        student_arch = detect_architecture_from_config(config.student_path)
    else:
        teacher_arch = detect_architecture_from_weights(teacher_keys)
        student_arch = detect_architecture_from_weights(student_keys)
    
    teacher_arch_w = detect_architecture_from_weights(teacher_keys)
    if teacher_arch_w.mlp_type == MLPType.MOE:
        teacher_arch.mlp_type = MLPType.MOE
        teacher_arch.num_experts = teacher_arch_w.num_experts
    
    print(f"\n📊 Teacher: {teacher_arch.attention_type.value}, {teacher_arch.mlp_type.value}")
    print(f"   Layers: {teacher_arch.num_hidden_layers}, Hidden: {teacher_arch.hidden_size}")
    if teacher_arch.num_experts > 0:
        print(f"   Experts: {teacher_arch.num_experts}")
    
    print(f"\n📊 Student: {student_arch.attention_type.value}, {student_arch.mlp_type.value}")
    print(f"   Layers: {student_arch.num_hidden_layers}, Hidden: {student_arch.hidden_size}")
    
    teacher_layers_map = scan_layers(teacher_keys)
    student_layers_map = scan_layers(student_keys)
    
    layer_mapping = create_layer_mapping(teacher_arch, student_arch)
    if layer_mapping and config.verbose:
        print(f"\n📋 Layer Name Mapping: {layer_mapping}")
    
    print(f"\n🚀 Starting GPU-accelerated distillation...")
    print(f"   CUDA Streams: {config.num_cuda_streams if config.use_cuda_streams else 'Disabled'}")
    print(f"   Mixed Precision: {config.use_mixed_precision}")
    print(f"   Offload Strategy: {config.offload_strategy}")
    print(f"   Prefetch Layers: {config.prefetch_layers}")
    print(f"   DARE: {config.use_dare} (drop_rate={config.dare_drop_rate})")
    print(f"   TIES: {config.use_ties} (density={config.ties_density}, trim_single={getattr(config, 'ties_trim_single', True)})")
    print(f"   Adaptive Rank: {config.use_adaptive_rank} (threshold={config.energy_threshold})")
    print(f"   Rank Range: {config.min_rank} - {config.max_rank}")
    print(f"   SVD Projection: {config.use_svd_projection}")

    # Calibration-aware SVD (optional)
    calib_stats: Dict[str, CalibStatsEntry] = {}
    if getattr(config, "calibration_mode", "none") != "none":
        print(f"\n🧪 Calibration-aware SVD: {config.calibration_mode}")
        if getattr(config, "calib_load", None):
            calib_stats = load_calibration_stats(config.calib_load)
            print(f"   Loaded calibration stats: {len(calib_stats)} modules")
            print(f"   From: {config.calib_load}")
        else:
            calib_stats = collect_calibration_stats(config.student_path, student_arch, config)
            print(f"   Collected calibration stats: {len(calib_stats)} modules")
            if getattr(config, "calib_save", None):
                save_calibration_stats(calib_stats, config.calib_save)
                print(f"   Saved calibration stats to: {config.calib_save}")

    start_time = time.time()

    distiller = AcceleratedDistiller(config, device, teacher_arch, student_arch)
    distiller.calib_stats = calib_stats if calib_stats else None
    lora_weights, rank_info = distiller.run(
        student_keys=student_keys,
        teacher_weight_map=teacher_weight_map,
        student_weight_map=student_weight_map,
        teacher_layers_map=teacher_layers_map,
        student_layers_map=student_layers_map,
        teacher_arch=teacher_arch,
        student_arch=student_arch,
        layer_name_mapping=layer_mapping
    )
    
    elapsed = time.time() - start_time
    
    if not lora_weights:
        print("❌ CRITICAL: No LoRA weights were generated!")
        return 1
    
    save_peft_adapter(lora_weights, config.output_path, config, student_arch, rank_info)
    
    print(f"\n✅ Generated {len(lora_weights) // 2} LoRA pairs in {elapsed:.1f}s")
    print(f"   Throughput: {len(lora_weights) // 2 / elapsed:.1f} layers/sec")
    
    if rank_info:
        ranks = list(rank_info.values())
        print(f"   Rank range: {min(ranks)}-{max(ranks)} (avg: {sum(ranks)/len(ranks):.1f})")
    
    print("\n🎉 Distillation Complete!")
    print("=" * 80)
    
    return 0


if __name__ == "__main__":
    exit(main())