upload src

src/common_lm_data.py

@@ -0,0 +1,435 @@
+#!/usr/bin/env python3
+"""Shared LM dataset helpers for fair cross-method comparisons."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Iterable, Iterator, List, Optional, Tuple
+
+import torch
+
+try:
+    from datasets import load_dataset
+    from datasets import Dataset as HFDataset
+except Exception:  # pragma: no cover - optional dependency
+    load_dataset = None
+    HFDataset = None
+
+
+def _normalize_config(config: Optional[str]) -> Optional[str]:
+    if config is None:
+        return None
+    if config.strip().lower() in {"none", "null", "-"}:
+        return None
+    return config
+
+
+def guess_text_field(dataset) -> str:
+    if hasattr(dataset, "column_names") and dataset.column_names:
+        if "text" in dataset.column_names:
+            return "text"
+        return dataset.column_names[0]
+    if hasattr(dataset, "features"):
+        names = list(dataset.features.keys())
+        if "text" in names:
+            return "text"
+        if names:
+            return names[0]
+    return "text"
+
+
+def normalize_dataset_name(name: str) -> str:
+    normalized = name.strip().lower()
+    aliases = {
+        "bookcorpus": "bookcorpus",
+        "boockcorpus": "bookcorpus",
+        "slimpajama": "slimpajama",
+        "dkyoon/slimpajama-6b": "slimpajama",
+    }
+    if normalized not in aliases:
+        raise ValueError(f"Unsupported dataset: {name}")
+    return aliases[normalized]
+
+
+def resolve_dataset_spec(
+    name: str,
+    config: Optional[str] = None,
+    split: str = "train",
+) -> Tuple[str, Optional[str], str]:
+    normalized = normalize_dataset_name(name)
+    if normalized == "bookcorpus":
+        return "bookcorpus", _normalize_config(config), split
+    if normalized == "slimpajama":
+        return "DKYoon/SlimPajama-6B", _normalize_config(config), split
+    raise ValueError(f"Unsupported dataset: {name}")
+
+
+def _sample_dataset_rows(dataset, target: int, seed: int) -> List[Dict[str, object]]:
+    if target <= 0:
+        return []
+    try:
+        dataset = dataset.shuffle(seed=seed)
+    except Exception:
+        pass
+
+    if hasattr(dataset, "__len__"):
+        limit = min(target, len(dataset))
+        dataset = dataset.select(range(limit))
+        return [row for row in dataset]
+
+    rows = []
+    for row in dataset:
+        rows.append(row)
+        if len(rows) >= target:
+            break
+    return rows
+
+
+def _iter_dataset_rows(dataset, seed: int) -> Iterator[Dict[str, object]]:
+    try:
+        dataset = dataset.shuffle(seed=seed)
+    except Exception:
+        pass
+    for row in dataset:
+        yield row
+
+
+def load_named_texts(
+    dataset_name: str,
+    *,
+    config: Optional[str] = None,
+    split: str = "train",
+    text_field: Optional[str] = None,
+    num_samples: int = 0,
+    seed: int = 0,
+) -> List[str]:
+    if load_dataset is None:
+        raise SystemExit("datasets is required for shared LM dataloaders")
+
+    hf_name, hf_config, hf_split = resolve_dataset_spec(dataset_name, config, split)
+    dataset = load_dataset(
+        hf_name,
+        hf_config,
+        split=hf_split,
+        trust_remote_code=True,
+    )
+    rows = dataset if num_samples <= 0 else _sample_dataset_rows(dataset, num_samples, seed)
+    field = text_field or guess_text_field(dataset)
+
+    texts: List[str] = []
+    for row in rows:
+        value = row.get(field, None) if isinstance(row, dict) else None
+        if isinstance(value, str) and value.strip():
+            texts.append(value)
+    return texts
+
+
+def build_token_chunks_from_rows(
+    rows: Iterable[Dict[str, object]],
+    *,
+    text_field: str,
+    tokenizer,
+    seq_len: int,
+    num_sequences: int = 0,
+    add_bos: bool = False,
+    max_rows: int = 0,
+) -> List[torch.Tensor]:
+    chunks: List[torch.Tensor] = []
+    buffer: List[int] = []
+    limit = None if num_sequences <= 0 else num_sequences
+    rows_seen = 0
+
+    for row in rows:
+        if max_rows > 0 and rows_seen >= max_rows:
+            break
+        rows_seen += 1
+
+        value = row.get(text_field, None) if isinstance(row, dict) else None
+        if not isinstance(value, str) or not value.strip():
+            continue
+
+        ids = tokenizer.encode(value, add_special_tokens=False)
+        if add_bos and tokenizer.bos_token_id is not None:
+            ids = [tokenizer.bos_token_id] + ids
+        if not ids:
+            continue
+
+        buffer.extend(ids)
+        while len(buffer) >= seq_len and (limit is None or len(chunks) < limit):
+            chunk = buffer[:seq_len]
+            buffer = buffer[seq_len:]
+            chunks.append(torch.tensor(chunk, dtype=torch.long))
+        if limit is not None and len(chunks) >= limit:
+            break
+
+    return chunks
+
+
+def collect_texts_from_rows(
+    rows: Iterable[Dict[str, object]],
+    *,
+    text_field: str,
+    tokenizer,
+    target_tokens: int = 0,
+    add_bos: bool = False,
+    max_rows: int = 0,
+) -> List[str]:
+    texts: List[str] = []
+    token_count = 0
+    rows_seen = 0
+
+    for row in rows:
+        if max_rows > 0 and rows_seen >= max_rows:
+            break
+        rows_seen += 1
+
+        value = row.get(text_field, None) if isinstance(row, dict) else None
+        if not isinstance(value, str) or not value.strip():
+            continue
+
+        texts.append(value)
+        if target_tokens > 0:
+            ids = tokenizer.encode(value, add_special_tokens=False)
+            if add_bos and tokenizer.bos_token_id is not None:
+                ids = [tokenizer.bos_token_id] + ids
+            token_count += len(ids)
+            if token_count >= target_tokens:
+                break
+
+    return texts
+
+
+def build_token_chunks(
+    texts: Iterable[str],
+    tokenizer,
+    seq_len: int,
+    num_sequences: int = 0,
+    add_bos: bool = False,
+) -> List[torch.Tensor]:
+    chunks: List[torch.Tensor] = []
+    buffer: List[int] = []
+    limit = None if num_sequences <= 0 else num_sequences
+
+    for text in texts:
+        ids = tokenizer.encode(text, add_special_tokens=False)
+        if add_bos and tokenizer.bos_token_id is not None:
+            ids = [tokenizer.bos_token_id] + ids
+        if not ids:
+            continue
+
+        buffer.extend(ids)
+        while len(buffer) >= seq_len and (limit is None or len(chunks) < limit):
+            chunk = buffer[:seq_len]
+            buffer = buffer[seq_len:]
+            chunks.append(torch.tensor(chunk, dtype=torch.long))
+        if limit is not None and len(chunks) >= limit:
+            break
+
+    return chunks
+
+
+class TokenChunkDataset(torch.utils.data.Dataset):
+    def __init__(self, chunks: List[torch.Tensor]) -> None:
+        self.chunks = chunks
+
+    def __len__(self) -> int:
+        return len(self.chunks)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        input_ids = self.chunks[idx]
+        attention_mask = torch.ones_like(input_ids)
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": input_ids.clone(),
+        }
+
+
+class TokenOnlyDataset(torch.utils.data.Dataset):
+    def __init__(self, chunks: List[torch.Tensor]) -> None:
+        self.chunks = chunks
+
+    def __len__(self) -> int:
+        return len(self.chunks)
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        return self.chunks[idx]
+
+
+class TokenInputMaskDataset(torch.utils.data.Dataset):
+    def __init__(self, chunks: List[torch.Tensor]) -> None:
+        self.chunks = chunks
+
+    def __len__(self) -> int:
+        return len(self.chunks)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        input_ids = self.chunks[idx]
+        return {
+            "input_ids": input_ids,
+            "attention_mask": torch.ones_like(input_ids),
+        }
+
+
+@dataclass
+class SharedLMDataSpec:
+    dataset: str
+    config: Optional[str] = None
+    split: str = "train"
+    text_field: Optional[str] = None
+    num_samples: int = 0
+    seq_len: int = 2048
+    num_sequences: int = 0
+    target_tokens: int = 0
+    batch_size: int = 1
+    shuffle: bool = False
+    num_workers: int = 0
+    seed: int = 0
+    add_bos: bool = False
+
+
+def build_chunks(spec: SharedLMDataSpec, tokenizer) -> List[torch.Tensor]:
+    if load_dataset is None:
+        raise SystemExit("datasets is required for shared LM dataloaders")
+
+    hf_name, hf_config, hf_split = resolve_dataset_spec(spec.dataset, spec.config, spec.split)
+    dataset = load_dataset(
+        hf_name,
+        hf_config,
+        split=hf_split,
+        trust_remote_code=True,
+    )
+
+    target_sequences = spec.num_sequences
+    if spec.target_tokens > 0:
+        token_sequences = (spec.target_tokens + spec.seq_len - 1) // spec.seq_len
+        target_sequences = max(target_sequences, token_sequences)
+    row_limit = spec.num_samples if target_sequences <= 0 else 0
+
+    rows = _iter_dataset_rows(dataset, spec.seed)
+    text_field = spec.text_field or guess_text_field(dataset)
+    chunks = build_token_chunks_from_rows(
+        rows,
+        text_field=text_field,
+        tokenizer=tokenizer,
+        seq_len=spec.seq_len,
+        num_sequences=target_sequences,
+        add_bos=spec.add_bos,
+        max_rows=row_limit,
+    )
+    return chunks
+
+
+def build_dataloader(spec: SharedLMDataSpec, tokenizer) -> torch.utils.data.DataLoader:
+    chunks = build_chunks(spec, tokenizer)
+    dataset = TokenChunkDataset(chunks)
+    return torch.utils.data.DataLoader(
+        dataset,
+        batch_size=spec.batch_size,
+        shuffle=spec.shuffle,
+        num_workers=spec.num_workers,
+    )
+
+
+def build_text_dataloader(spec: SharedLMDataSpec, tokenizer) -> torch.utils.data.DataLoader:
+    if load_dataset is None:
+        raise SystemExit("datasets is required for shared LM dataloaders")
+
+    hf_name, hf_config, hf_split = resolve_dataset_spec(spec.dataset, spec.config, spec.split)
+    dataset = load_dataset(
+        hf_name,
+        hf_config,
+        split=hf_split,
+        trust_remote_code=True,
+    )
+    rows = _iter_dataset_rows(dataset, spec.seed)
+    text_field = spec.text_field or guess_text_field(dataset)
+    row_limit = spec.num_samples
+    texts = collect_texts_from_rows(
+        rows,
+        text_field=text_field,
+        tokenizer=tokenizer,
+        target_tokens=spec.target_tokens,
+        add_bos=spec.add_bos,
+        max_rows=row_limit,
+    )
+    return torch.utils.data.DataLoader(
+        texts,
+        batch_size=spec.batch_size,
+        shuffle=spec.shuffle,
+        num_workers=spec.num_workers,
+        drop_last=True,
+    )
+
+
+def build_uidl_post_train_dataloader(
+    spec: SharedLMDataSpec,
+    tokenizer,
+) -> torch.utils.data.DataLoader:
+    dataset = TokenChunkDataset(build_chunks(spec, tokenizer))
+    return torch.utils.data.DataLoader(
+        dataset,
+        batch_size=spec.batch_size,
+        shuffle=spec.shuffle,
+        num_workers=spec.num_workers,
+    )
+
+
+def build_uidl_similarity_dataloader(
+    spec: SharedLMDataSpec,
+    tokenizer,
+) -> torch.utils.data.DataLoader:
+    dataset = TokenInputMaskDataset(build_chunks(spec, tokenizer))
+    return torch.utils.data.DataLoader(
+        dataset,
+        batch_size=spec.batch_size,
+        shuffle=spec.shuffle,
+        num_workers=spec.num_workers,
+    )
+
+
+def build_shortened_llm_dataloader(
+    spec: SharedLMDataSpec,
+    tokenizer,
+) -> torch.utils.data.DataLoader:
+    dataset = TokenOnlyDataset(build_chunks(spec, tokenizer))
+    return torch.utils.data.DataLoader(
+        dataset,
+        batch_size=spec.batch_size,
+        shuffle=spec.shuffle,
+        num_workers=spec.num_workers,
+    )
+
+
+def build_shortened_llm_examples(spec: SharedLMDataSpec, tokenizer) -> torch.Tensor:
+    chunks = build_chunks(spec, tokenizer)
+    if not chunks:
+        return torch.empty((0, spec.seq_len), dtype=torch.long)
+    return torch.stack(chunks, dim=0)
+
+
+def build_llmpruner_examples(spec: SharedLMDataSpec, tokenizer) -> torch.Tensor:
+    chunks = build_chunks(spec, tokenizer)
+    if not chunks:
+        return torch.empty((0, spec.seq_len), dtype=torch.long)
+    return torch.stack(chunks, dim=0)
+
+
+def build_replaceme_dataloader(
+    spec: SharedLMDataSpec,
+    tokenizer,
+) -> torch.utils.data.DataLoader:
+    return build_text_dataloader(spec, tokenizer)
+
+
+def build_hf_causal_dataset(spec: SharedLMDataSpec, tokenizer):
+    if HFDataset is None:
+        raise SystemExit("datasets is required for shared LM dataloaders")
+
+    chunks = build_chunks(spec, tokenizer)
+    payload = {
+        "input_ids": [chunk.tolist() for chunk in chunks],
+        "attention_mask": [torch.ones_like(chunk).tolist() for chunk in chunks],
+        "labels": [chunk.tolist() for chunk in chunks],
+    }
+    return HFDataset.from_dict(payload)

src/convert_llmpruner_checkpoint.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+import argparse
+import os
+import sys
+
+import torch
+
+
+def ensure_llmpruner_on_path() -> None:
+    repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    llmpruner_root = os.path.join(repo_root, "compare_model", "LLM-Pruner")
+    if os.path.isdir(llmpruner_root) and llmpruner_root not in sys.path:
+        sys.path.insert(0, llmpruner_root)
+
+
+def load_llmpruner_checkpoint(path: str):
+    ensure_llmpruner_on_path()
+    checkpoint = torch.load(path, map_location="cpu", weights_only=False)
+    if not isinstance(checkpoint, dict) or "model" not in checkpoint or "tokenizer" not in checkpoint:
+        raise SystemExit(
+            "Expected an LLM-Pruner checkpoint dict with `model` and `tokenizer` entries."
+        )
+    return checkpoint["model"], checkpoint["tokenizer"]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Convert an LLM-Pruner .bin checkpoint to a Hugging Face save_pretrained directory."
+    )
+    parser.add_argument("--input", required=True, help="Path to LLM-Pruner pytorch_model.bin")
+    parser.add_argument("--output_dir", required=True, help="Directory to write HF model artifacts")
+    args = parser.parse_args()
+
+    model, tokenizer = load_llmpruner_checkpoint(args.input)
+    os.makedirs(args.output_dir, exist_ok=True)
+    model.save_pretrained(args.output_dir)
+    tokenizer.save_pretrained(args.output_dir)
+    print(args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

src/eval_ppl.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+import argparse
+import csv
+import json
+import os
+import sys
+from typing import Iterable
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class IndexDataset(Dataset):
+    def __init__(self, tensors: torch.Tensor):
+        self.tensors = tensors
+
+    def __getitem__(self, index: int) -> torch.Tensor:
+        return self.tensors[index]
+
+    def __len__(self) -> int:
+        return len(self.tensors)
+
+
+def get_dataset(name: str):
+    if name == "wikitext2":
+        train_data = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
+        test_data = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+        return train_data, test_data, "text"
+    if name == "ptb":
+        train_data = load_dataset("ptb_text_only", "penn_treebank", split="train")
+        test_data = load_dataset("ptb_text_only", "penn_treebank", split="validation")
+        return train_data, test_data, "sentence"
+    raise ValueError(f"Unsupported dataset: {name}")
+
+
+def process_data(samples, tokenizer, seq_len: int, field_name: str, add_bos_to_every: bool) -> IndexDataset:
+    test_ids = tokenizer(
+        "\n\n".join(samples[field_name]),
+        return_tensors="pt",
+        add_special_tokens=False,
+    ).input_ids[0]
+
+    if not add_bos_to_every and tokenizer.bos_token_id is not None:
+        test_ids = torch.cat((torch.LongTensor([tokenizer.bos_token_id]), test_ids), dim=0)
+
+    batches = []
+    num_samples = test_ids.numel() // seq_len
+    for index in range(num_samples):
+        batch = test_ids[(index * seq_len) : ((index + 1) * seq_len)]
+        if add_bos_to_every and tokenizer.bos_token_id is not None:
+            batch = torch.cat((torch.LongTensor([tokenizer.bos_token_id]), batch), dim=0)
+        batches.append(batch)
+
+    return IndexDataset(tensors=torch.stack(batches))
+
+
+def get_loader(name: str, tokenizer, seq_len: int, batch_size: int, add_bos_to_every: bool):
+    _, test_data, field_name = get_dataset(name)
+    dataset = process_data(test_data, tokenizer, seq_len, field_name, add_bos_to_every)
+    return DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+
+@torch.no_grad()
+def evaluate_ppl(model, test_loader, device: str) -> float:
+    nlls = []
+    for batch in tqdm(test_loader, desc="Running PPL", dynamic_ncols=True):
+        batch = batch.to(device)
+        outputs = model(batch)
+        shift_logits = outputs.logits[:, :-1, :].contiguous()
+        shift_labels = batch[:, 1:].contiguous()
+        loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
+        loss = loss_fct(
+            shift_logits.reshape(-1, shift_logits.size(-1)),
+            shift_labels.view(-1),
+        )
+        nlls.append(loss.cpu())
+
+    return float(np.exp(torch.cat(nlls, dim=-1).mean().item()))
+
+
+def resolve_dtype(args) -> torch.dtype:
+    if args.use_bfloat:
+        return torch.bfloat16
+
+    dtype_name = args.dtype if args.dtype is not None else args.torch_dtype
+    if dtype_name is None:
+        dtype_name = "float16"
+
+    dtype_map = {
+        "float16": torch.float16,
+        "fp16": torch.float16,
+        "bfloat16": torch.bfloat16,
+        "bf16": torch.bfloat16,
+        "float32": torch.float32,
+        "fp32": torch.float32,
+    }
+    if dtype_name not in dtype_map:
+        raise ValueError(f"Unsupported dtype: {dtype_name}")
+    return dtype_map[dtype_name]
+
+
+def normalize_datasets(datasets: Iterable[str]) -> list[str]:
+    normalized = []
+    for dataset in datasets:
+        normalized.append("wikitext2" if dataset == "wikitext" else dataset)
+    return normalized
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Shared perplexity evaluation for abprune.")
+    parser.add_argument("--base_model", "--model-path", dest="model_path", required=True)
+    parser.add_argument("--output_dir", type=str, default=None)
+    parser.add_argument("--dataset", nargs="+", default=["wikitext2", "ptb"])
+    parser.add_argument("--max_seq_len", "--seq-len", dest="seq_len", type=int, default=1024)
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument(
+        "--dtype",
+        default=None,
+        choices=["float16", "fp16", "bfloat16", "bf16", "float32", "fp32"],
+    )
+    parser.add_argument(
+        "--torch_dtype",
+        default=None,
+        choices=["float16", "fp16", "bfloat16", "bf16", "float32", "fp32"],
+    )
+    parser.add_argument("--use_bfloat", action="store_true")
+    parser.add_argument("--add_bos_to_every", action="store_true")
+    parser.add_argument("--fix_decapoda_config", action="store_true")
+    parser.add_argument("--local_files_only", action="store_true")
+    return parser
+
+
+def maybe_fix_decapoda_config(tokenizer, enabled: bool) -> None:
+    if not enabled:
+        return
+    if tokenizer.bos_token_id is None and tokenizer.eos_token_id is not None:
+        tokenizer.bos_token = tokenizer.eos_token
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+
+def ensure_llmpruner_on_path() -> None:
+    repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    llmpruner_root = os.path.join(repo_root, "compare_model", "LLM-Pruner")
+    if os.path.isdir(llmpruner_root) and llmpruner_root not in sys.path:
+        sys.path.insert(0, llmpruner_root)
+
+
+def load_model_and_tokenizer(model_path: str, *, torch_dtype: torch.dtype, local_files_only: bool):
+    if os.path.isfile(model_path) and model_path.endswith(".bin"):
+        ensure_llmpruner_on_path()
+        checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
+        if not isinstance(checkpoint, dict) or "model" not in checkpoint or "tokenizer" not in checkpoint:
+            raise ValueError(
+                "Expected an LLM-Pruner checkpoint dict with `model` and `tokenizer` entries."
+            )
+        model = checkpoint["model"]
+        tokenizer = checkpoint["tokenizer"]
+        if torch_dtype is not None:
+            model = model.to(dtype=torch_dtype)
+        return model, tokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path,
+        local_files_only=local_files_only,
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        torch_dtype=torch_dtype,
+        local_files_only=local_files_only,
+    )
+    return model, tokenizer
+
+
+def main() -> None:
+    parser = build_arg_parser()
+    args = parser.parse_args()
+
+    datasets = normalize_datasets(args.dataset)
+    torch_dtype = resolve_dtype(args)
+
+    model, tokenizer = load_model_and_tokenizer(
+        args.model_path,
+        torch_dtype=torch_dtype,
+        local_files_only=args.local_files_only,
+    )
+    maybe_fix_decapoda_config(tokenizer, args.fix_decapoda_config)
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model.eval()
+    model.to(args.device)
+
+    metrics = {}
+    for dataset in datasets:
+        test_loader = get_loader(
+            dataset,
+            tokenizer,
+            seq_len=args.seq_len,
+            batch_size=args.batch_size,
+            add_bos_to_every=args.add_bos_to_every,
+        )
+        metrics[dataset] = evaluate_ppl(model, test_loader, args.device)
+        print(f"PPL-{dataset}: {metrics[dataset]} | add_bos_to_every: {args.add_bos_to_every} | seq_len: {args.seq_len}")
+
+    mem = None
+    if torch.cuda.is_available() and args.device.startswith("cuda"):
+        mem = torch.cuda.memory_allocated(args.device) / 1024 / 1024
+
+    result = {
+        "model_path": os.path.abspath(args.model_path),
+        "datasets": datasets,
+        "seq_len": args.seq_len,
+        "batch_size": args.batch_size,
+        "device": args.device,
+        "dtype": str(torch_dtype).replace("torch.", ""),
+        "add_bos_to_every": args.add_bos_to_every,
+        "metrics": metrics,
+        "params": int(sum(parameter.numel() for parameter in model.parameters())),
+        "mem_mib": mem,
+    }
+
+    if args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+        filename = "ppl_bos.csv" if args.add_bos_to_every else "ppl.csv"
+        csv_path = os.path.join(args.output_dir, filename)
+        with open(csv_path, "w", newline="", encoding="utf-8") as handle:
+            writer = csv.writer(handle)
+            writer.writerow([*(f"ppl_{dataset}" for dataset in datasets), "params", "mem"])
+            writer.writerow([*(metrics[dataset] for dataset in datasets), result["params"], mem])
+
+    print(json.dumps(result, ensure_ascii=True))
+
+
+if __name__ == "__main__":
+    main()

src/fbmc_metric.py

@@ -0,0 +1,519 @@
+#!/usr/bin/env python3
+"""Estimate Fisher-Barycentric Merge Cost (FBMC) for adjacent layers."""
+
+import argparse
+import csv
+import json
+import os
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+try:
+    from datasets import load_dataset
+except Exception:  # pragma: no cover - optional dependency
+    load_dataset = None
+
+try:
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+except Exception as exc:  # pragma: no cover - fail early with clear error
+    raise SystemExit("transformers is required: pip install transformers") from exc
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Compute FBMC for adjacent layers of a Hugging Face causal LM."
+    )
+    parser.add_argument("--model", required=True, help="HF model id or local path")
+    parser.add_argument(
+        "--dataset",
+        action="append",
+        default=[],
+        help=(
+            "HF dataset name (repeatable). Optional if using --text or --text_file."
+        ),
+    )
+    parser.add_argument(
+        "--dataset_config",
+        action="append",
+        default=[],
+        help="Optional dataset config (repeatable or single shared config).",
+    )
+    parser.add_argument(
+        "--dataset_split",
+        default="train",
+        help="Dataset split to use (default: train)",
+    )
+    parser.add_argument(
+        "--dataset_text_field",
+        default=None,
+        help="Text field in dataset (default: auto-detect, applies to all datasets)",
+    )
+    parser.add_argument(
+        "--text",
+        action="append",
+        default=[],
+        help="Inline text samples (can pass multiple)",
+    )
+    parser.add_argument(
+        "--text_file",
+        default=None,
+        help="Path to a text file for calibration data",
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=128,
+        help="Number of token sequences to use",
+    )
+    parser.add_argument(
+        "--seq_len", type=int, default=256, help="Sequence length"
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=2, help="Batch size"
+    )
+    parser.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device for model + compute",
+    )
+    parser.add_argument(
+        "--dtype",
+        default="auto",
+        choices=["auto", "float32", "float16", "bfloat16"],
+        help="Model dtype",
+    )
+    parser.add_argument(
+        "--layer_path",
+        default=None,
+        help="Override layer attribute path (e.g., model.layers)",
+    )
+    parser.add_argument(
+        "--fisher_mode",
+        default="tensor",
+        choices=["tensor", "param"],
+        help="Fisher approximation granularity",
+    )
+    parser.add_argument("--eps", type=float, default=1e-8, help="Stability epsilon")
+    parser.add_argument(
+        "--output",
+        default=None,
+        help="Optional JSON output path",
+    )
+    parser.add_argument(
+        "--output_csv",
+        default=None,
+        help="Optional CSV output path",
+    )
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument(
+        "--trust_remote_code",
+        action="store_true",
+        help="Allow custom model code from hub",
+    )
+    return parser.parse_args()
+
+
+def resolve_attr(root: object, path: str) -> Optional[object]:
+    cur = root
+    for part in path.split("."):
+        if not hasattr(cur, part):
+            return None
+        cur = getattr(cur, part)
+    return cur
+
+
+def find_layers(model, layer_path: Optional[str]) -> List[torch.nn.Module]:
+    if layer_path:
+        layers = resolve_attr(model, layer_path)
+        if layers is None:
+            raise ValueError(f"layer_path '{layer_path}' not found on model")
+        return list(layers)
+
+    # Common decoder-only layer containers. Add more if needed.
+    candidate_paths = [
+        "model.layers",  # LLaMA, Mistral, Qwen2, Gemma
+        "model.decoder.layers",  # OPT
+        "transformer.h",  # GPT-2, GPT-J, Bloom, Falcon
+        "transformer.blocks",  # MPT
+        "gpt_neox.layers",  # GPT-NeoX
+        "layers",  # fallback
+    ]
+    for path in candidate_paths:
+        layers = resolve_attr(model, path)
+        if layers is not None:
+            try:
+                return list(layers)
+            except TypeError:
+                continue
+    raise ValueError(
+        "Could not locate transformer layers. Pass --layer_path explicitly."
+    )
+
+
+def guess_text_field(dataset) -> str:
+    if hasattr(dataset, "column_names") and dataset.column_names:
+        if "text" in dataset.column_names:
+            return "text"
+        return dataset.column_names[0]
+    if hasattr(dataset, "features"):
+        names = list(dataset.features.keys())
+        if "text" in names:
+            return "text"
+        if names:
+            return names[0]
+    return "text"
+
+
+def _normalize_config(config: Optional[str]) -> Optional[str]:
+    if config is None:
+        return None
+    if config.strip().lower() in {"none", "null", "-"}:
+        return None
+    return config
+
+
+def _expand_dataset_configs(
+    datasets: List[str], configs: List[str]
+) -> List[Optional[str]]:
+    if not configs:
+        return [None] * len(datasets)
+    if len(configs) == 1 and len(datasets) > 1:
+        return [_normalize_config(configs[0])] * len(datasets)
+    if len(configs) != len(datasets):
+        raise SystemExit(
+            "Provide zero, one, or matching-count --dataset_config values."
+        )
+    return [_normalize_config(cfg) for cfg in configs]
+
+
+def _sample_dataset_rows(
+    dataset, target: int, seed: int
+) -> List[Dict[str, object]]:
+    if target <= 0:
+        return []
+    try:
+        dataset = dataset.shuffle(seed=seed)
+    except Exception:
+        pass
+
+    if hasattr(dataset, "__len__"):
+        limit = min(target, len(dataset))
+        dataset = dataset.select(range(limit))
+        return [row for row in dataset]
+
+    # IterableDataset fallback.
+    rows = []
+    for row in dataset:
+        rows.append(row)
+        if len(rows) >= target:
+            break
+    return rows
+
+
+def load_texts(args: argparse.Namespace) -> List[str]:
+    texts: List[str] = []
+    if args.text_file:
+        with open(args.text_file, "r", encoding="utf-8") as handle:
+            texts.extend([line.strip() for line in handle if line.strip()])
+    if args.text:
+        texts.extend([t for t in args.text if t])
+
+    if args.dataset:
+        if load_dataset is None:
+            raise SystemExit("datasets is required for --dataset")
+
+        datasets = list(args.dataset)
+        configs = _expand_dataset_configs(datasets, list(args.dataset_config))
+        num_datasets = len(datasets)
+        base = args.num_samples // num_datasets
+        remainder = args.num_samples % num_datasets
+
+        for idx, (dataset_name, config) in enumerate(zip(datasets, configs)):
+            target = base + (1 if idx < remainder else 0)
+            dataset = load_dataset(
+                dataset_name,
+                config,
+                split=args.dataset_split,
+                trust_remote_code=True,
+            )
+            rows = _sample_dataset_rows(dataset, target, args.seed + idx)
+            text_field = args.dataset_text_field or guess_text_field(dataset)
+            for row in rows:
+                value = row.get(text_field, None) if isinstance(row, dict) else None
+                if isinstance(value, str) and value.strip():
+                    texts.append(value)
+
+    return texts
+
+
+def build_token_chunks(
+    texts: List[str], tokenizer, seq_len: int, num_samples: int
+) -> List[torch.Tensor]:
+    chunks: List[torch.Tensor] = []
+    buffer: List[int] = []
+    for text in texts:
+        ids = tokenizer.encode(text, add_special_tokens=False)
+        if not ids:
+            continue
+        buffer.extend(ids)
+        while len(buffer) >= seq_len and len(chunks) < num_samples:
+            chunk = buffer[:seq_len]
+            buffer = buffer[seq_len:]
+            chunks.append(torch.tensor(chunk, dtype=torch.long))
+        if len(chunks) >= num_samples:
+            break
+    return chunks
+
+
+def get_dtype(dtype: str):
+    if dtype == "auto":
+        return None
+    if dtype == "float16":
+        return torch.float16
+    if dtype == "bfloat16":
+        return torch.bfloat16
+    return torch.float32
+
+
+def compute_fisher(
+    model,
+    layers: List[torch.nn.Module],
+    dataloader,
+    fisher_mode: str,
+    device: str,
+) -> Tuple[List[Dict[str, object]], int, List[Dict[str, int]]]:
+    # Only compute grads for layer params.
+    for param in model.parameters():
+        param.requires_grad_(False)
+    for layer in layers:
+        for param in layer.parameters():
+            param.requires_grad_(True)
+
+    fisher_sums: List[Dict[str, object]] = []
+    param_numels: List[Dict[str, int]] = []
+    for layer in layers:
+        layer_sums: Dict[str, object] = {}
+        layer_numels: Dict[str, int] = {}
+        for name, param in layer.named_parameters():
+            if not param.requires_grad:
+                continue
+            if fisher_mode == "param":
+                layer_sums[name] = torch.zeros_like(
+                    param, dtype=torch.float32, device="cpu"
+                )
+            else:
+                layer_sums[name] = 0.0
+                layer_numels[name] = param.numel()
+        fisher_sums.append(layer_sums)
+        param_numels.append(layer_numels)
+
+    num_batches = 0
+    model.eval()
+    for batch in dataloader:
+        input_ids = batch[0].to(device)
+        outputs = model(input_ids=input_ids, labels=input_ids)
+        loss = outputs.loss
+        loss.backward()
+        for layer_idx, layer in enumerate(layers):
+            layer_sums = fisher_sums[layer_idx]
+            for name, param in layer.named_parameters():
+                if not param.requires_grad:
+                    continue
+                if param.grad is None:
+                    continue
+                grad_sq = param.grad.detach().float().pow(2)
+                if fisher_mode == "param":
+                    layer_sums[name] += grad_sq.cpu()
+                else:
+                    layer_sums[name] += float(grad_sq.sum().item())
+        model.zero_grad(set_to_none=True)
+        num_batches += 1
+
+    if num_batches == 0:
+        raise RuntimeError("No batches processed; check dataset or text inputs.")
+
+    return fisher_sums, num_batches, param_numels
+
+
+def compute_fbmc_costs(
+    layers: List[torch.nn.Module],
+    fisher_sums: List[Dict[str, object]],
+    num_batches: int,
+    param_numels: List[Dict[str, int]],
+    fisher_mode: str,
+    eps: float,
+) -> List[Dict[str, object]]:
+    layer_params: List[Dict[str, torch.nn.Parameter]] = []
+    for layer in layers:
+        layer_params.append({name: param for name, param in layer.named_parameters()})
+
+    results: List[Dict[str, object]] = []
+    for idx in range(len(layers) - 1):
+        cost = 0.0
+        matched = 0
+        skipped = 0
+        params_i = layer_params[idx]
+        params_j = layer_params[idx + 1]
+        for name, param_i in params_i.items():
+            param_j = params_j.get(name)
+            if param_j is None or param_j.shape != param_i.shape:
+                skipped += 1
+                continue
+            matched += 1
+            if fisher_mode == "param":
+                fisher_i = fisher_sums[idx][name] / num_batches
+                fisher_j = fisher_sums[idx + 1][name] / num_batches
+                diff = (param_i.detach().float().cpu() - param_j.detach().float().cpu())
+                denom = fisher_i + fisher_j + eps
+                term = (fisher_i * fisher_j / denom) * diff * diff
+                cost += 0.5 * float(term.sum().item())
+            else:
+                fisher_i = fisher_sums[idx][name] / (
+                    num_batches * param_numels[idx][name]
+                )
+                fisher_j = fisher_sums[idx + 1][name] / (
+                    num_batches * param_numels[idx + 1][name]
+                )
+                denom = fisher_i + fisher_j + eps
+                if denom == 0:
+                    continue
+                diff_sq = (
+                    param_i.detach().float() - param_j.detach().float()
+                ).pow(2)
+                cost += 0.5 * (fisher_i * fisher_j / denom) * float(
+                    diff_sq.sum().item()
+                )
+        results.append(
+            {
+                "layer_i": idx,
+                "layer_j": idx + 1,
+                "fbmc": cost,
+                "matched_params": matched,
+                "skipped_params": skipped,
+            }
+        )
+    return results
+
+
+def main() -> None:
+    args = parse_args()
+    torch.manual_seed(args.seed)
+
+    dtype = get_dtype(args.dtype)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model,
+        torch_dtype=dtype,
+        trust_remote_code=args.trust_remote_code,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model, trust_remote_code=args.trust_remote_code
+    )
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    layers = find_layers(model, args.layer_path)
+    if len(layers) < 2:
+        raise SystemExit("Model has fewer than 2 layers; cannot compute FBMC.")
+
+    texts = load_texts(args)
+    if not texts:
+        raise SystemExit(
+            "No calibration text found. Provide --dataset, --text, or --text_file."
+        )
+
+    chunks = build_token_chunks(texts, tokenizer, args.seq_len, args.num_samples)
+    if not chunks:
+        raise SystemExit("Not enough text to build token sequences.")
+
+    dataset = torch.utils.data.TensorDataset(torch.stack(chunks))
+    dataloader = torch.utils.data.DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=False
+    )
+
+    model.to(args.device)
+
+    fisher_sums, num_batches, param_numels = compute_fisher(
+        model,
+        layers,
+        dataloader,
+        fisher_mode=args.fisher_mode,
+        device=args.device,
+    )
+
+    costs = compute_fbmc_costs(
+        layers,
+        fisher_sums,
+        num_batches,
+        param_numels,
+        fisher_mode=args.fisher_mode,
+        eps=args.eps,
+    )
+
+    costs_sorted = sorted(costs, key=lambda x: x["fbmc"])
+    best = costs_sorted[0]
+
+    print("FBMC results (layer order):")
+    for item in costs:
+        print(
+            f"layers {item['layer_i']} & {item['layer_j']} -> "
+            f"fbmc={item['fbmc']:.6e} "
+            f"(matched={item['matched_params']}, skipped={item['skipped_params']})"
+        )
+    print("\nFBMC results (lowest cost first):")
+    for item in costs_sorted:
+        print(
+            f"layers {item['layer_i']} & {item['layer_j']} -> "
+            f"fbmc={item['fbmc']:.6e} "
+            f"(matched={item['matched_params']}, skipped={item['skipped_params']})"
+        )
+    print(
+        f"\nBest pair: layers {best['layer_i']} & {best['layer_j']} "
+        f"(fbmc={best['fbmc']:.6e})"
+    )
+
+    if args.output:
+        payload = {
+            "model": args.model,
+            "num_layers": len(layers),
+            "fisher_mode": args.fisher_mode,
+            "num_batches": num_batches,
+            "num_sequences": len(chunks),
+            "seq_len": args.seq_len,
+            "best_pair": best,
+            "pairs": costs_sorted,
+        }
+        os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+        with open(args.output, "w", encoding="utf-8") as handle:
+            json.dump(payload, handle, indent=2)
+        print(f"\nWrote results to {args.output}")
+
+    if args.output_csv:
+        os.makedirs(os.path.dirname(args.output_csv) or ".", exist_ok=True)
+        with open(args.output_csv, "w", encoding="utf-8", newline="") as handle:
+            writer = csv.DictWriter(
+                handle,
+                fieldnames=[
+                    "layer_i",
+                    "layer_j",
+                    "fbmc",
+                    "matched_params",
+                    "skipped_params",
+                ],
+            )
+            writer.writeheader()
+            for item in costs_sorted:
+                writer.writerow(
+                    {
+                        "layer_i": item["layer_i"],
+                        "layer_j": item["layer_j"],
+                        "fbmc": item["fbmc"],
+                        "matched_params": item["matched_params"],
+                        "skipped_params": item["skipped_params"],
+                    }
+                )
+        print(f"Wrote CSV results to {args.output_csv}")
+
+
+if __name__ == "__main__":
+    main()

src/fuse_layers.py

@@ -0,0 +1,2416 @@
+#!/usr/bin/env python3
+"""Fuse adjacent layers via attention head alignment + Fisher-barycentric merge."""
+
+import argparse
+import copy
+import gc
+import json
+import os
+import random
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Set, Tuple
+
+import torch
+
+try:
+    import numpy as np
+except Exception:  # pragma: no cover - optional dependency
+    np = None
+
+try:
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+except Exception as exc:  # pragma: no cover - fail early with clear error
+    raise SystemExit("transformers is required: pip install transformers") from exc
+
+try:
+    import ppl_eval
+except Exception as exc:  # pragma: no cover - optional dependency
+    raise SystemExit("ppl_eval.py is required (missing or invalid)") from exc
+
+from fuse_layers_data import (
+    FixedSeqDataset,
+    build_token_chunks,
+    expand_dataset_configs,
+    load_instruction_records,
+    load_texts,
+    load_texts_from_datasets,
+)
+from common_lm_data import SharedLMDataSpec, build_chunks, build_dataloader
+from fuse_layers_distill import (
+    commutator_precondition,
+    compute_fisher_gate_priors,
+    distill_reparam_merge,
+    lora_ce_finetune,
+)
+from fuse_layers_model import (
+    apply_norm_policy,
+    build_head_permutation,
+    clone_state_dict,
+    compute_fisher,
+    compute_head_means,
+    decrement_config,
+    drop_layer,
+    find_attention_module,
+    find_colon_modules,
+    find_layer_container,
+    get_dtype,
+    get_norm_pair,
+    merge_layers,
+    permute_attention_heads,
+)
+from fuse_layers_select import select_layer_auto
+from progressive_loader import load_causal_lm, load_progressive_model
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Fuse layer i and i+1 using head alignment + Fisher barycenter."
+    )
+    parser.add_argument("--model", required=True, help="HF model id or local path")
+    parser.add_argument(
+        "--model_cache_dir",
+        default=None,
+        help="Optional cache dir for model/tokenizer downloads",
+    )
+    parser.add_argument(
+        "--layer",
+        type=str,
+        default="auto",
+        help="Layer index i (int) or 'auto' to select via auto metric",
+    )
+    parser.add_argument(
+        "--selection_method",
+        choices=["dwce", "sequential"],
+        default="dwce",
+        help=(
+            "Pair selection policy for progressive pruning. "
+            "'dwce' uses downstream-weighted composition error; "
+            "'sequential' always takes the next available pair."
+        ),
+    )
+    parser.add_argument(
+        "--exclude_pairs",
+        "--exclude_layers",
+        nargs="*",
+        default=None,
+        dest="exclude_pairs",
+        help=(
+            "Exclude pair indices from consideration for any fusion. Indices refer to "
+            "pair start positions in [0..N-2]. Negative indices count from the end "
+            "(-1 = last pair, -2 = second last). Accepts space- or comma-separated ints. "
+            "Alias: --exclude_layers (deprecated)."
+        ),
+    )
+    parser.add_argument(
+        "--output_dir", required=True, help="Directory to write fused model"
+    )
+    parser.add_argument(
+        "--dataset",
+        action="append",
+        default=[],
+        help=(
+            "HF dataset name (repeatable). Optional if using --text or --text_file."
+        ),
+    )
+    parser.add_argument(
+        "--dataset_config",
+        action="append",
+        default=[],
+        help="Optional dataset config (repeatable or single shared config).",
+    )
+    parser.add_argument(
+        "--dataset_split",
+        default="train",
+        help="Dataset split to use (default: train)",
+    )
+    parser.add_argument(
+        "--dataset_text_field",
+        default=None,
+        help="Text field in dataset (default: auto-detect, applies to all datasets)",
+    )
+    parser.add_argument(
+        "--text",
+        action="append",
+        default=[],
+        help="Inline text samples (can pass multiple)",
+    )
+    parser.add_argument(
+        "--text_file",
+        default=None,
+        help="Path to a text file for calibration data",
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=128,
+        help="Number of token sequences to use",
+    )
+    parser.add_argument(
+        "--target_tokens",
+        type=int,
+        default=0,
+        help="Target token budget for common_lm_data-backed calibration/distillation (0 = disabled)",
+    )
+    parser.add_argument("--seq_len", type=int, default=256, help="Sequence length")
+    parser.add_argument("--batch_size", type=int, default=2, help="Batch size")
+    parser.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device for model + compute",
+    )
+    parser.add_argument(
+        "--dtype",
+        default="auto",
+        choices=["auto", "float32", "float16", "bfloat16"],
+        help="Model dtype",
+    )
+    parser.add_argument(
+        "--layer_path",
+        default=None,
+        help="Override layer attribute path (e.g., model.layers)",
+    )
+    parser.add_argument(
+        "--fisher_mode",
+        default="tensor",
+        choices=["tensor", "param"],
+        help="Fisher approximation granularity",
+    )
+    parser.add_argument(
+        "--no_head_permute",
+        action="store_true",
+        help=(
+            "Deprecated alias for --no_head_permute_merge. "
+            "Disables merge-stage head permutation only."
+        ),
+    )
+    parser.add_argument(
+        "--no_head_permute_merge",
+        action="store_true",
+        help="Disable attention head permutation alignment before merge",
+    )
+    parser.add_argument(
+        "--no_head_permute_select",
+        action="store_true",
+        help="Disable attention head permutation alignment during auto selection",
+    )
+    parser.add_argument("--eps", type=float, default=1e-8, help="Stability epsilon")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument(
+        "--trust_remote_code",
+        action="store_true",
+        help="Allow custom model code from hub",
+    )
+    parser.add_argument(
+        "--save_metadata",
+        action="store_true",
+        help="Backward-compatible no-op; metadata is always written.",
+    )
+    parser.add_argument(
+        "--skip_eval",
+        action="store_true",
+        help="Skip pre/post perplexity evaluation",
+    )
+    parser.add_argument(
+        "--eval_dataset",
+        action="append",
+        default=[],
+        help="Evaluation dataset name (repeatable). Defaults to wikitext.",
+    )
+    parser.add_argument(
+        "--eval_dataset_config",
+        action="append",
+        default=[],
+        help="Evaluation dataset config (repeatable or single shared config).",
+    )
+    parser.add_argument(
+        "--eval_split",
+        default="test",
+        help="Evaluation dataset split (default: test)",
+    )
+    parser.add_argument(
+        "--eval_text_field",
+        default=None,
+        help="Evaluation text field override (default: auto-detect)",
+    )
+    parser.add_argument(
+        "--eval_model_family",
+        type=str,
+        choices=["auto", "llama", "qwen"],
+        default="auto",
+        help="Model family for BOS handling during eval",
+    )
+    parser.add_argument(
+        "--eval_add_bos",
+        type=str,
+        choices=["auto", "always", "never"],
+        default="auto",
+        help="Whether to prepend BOS to each eval sample",
+    )
+    parser.add_argument(
+        "--eval_num_samples",
+        type=int,
+        default=0,
+        help="Number of token sequences per eval dataset (0 = all)",
+    )
+    parser.add_argument(
+        "--eval_seq_len",
+        type=int,
+        default=2048,
+        help="Sequence length for eval",
+    )
+    parser.add_argument(
+        "--eval_batch_size",
+        type=int,
+        default=None,
+        help="Batch size for eval (defaults to --batch_size)",
+    )
+    parser.add_argument(
+        "--eval_max_batches",
+        type=int,
+        default=None,
+        help="Optional max number of eval batches per dataset",
+    )
+    parser.add_argument(
+        "--eval_cache_dir",
+        default=None,
+        help="Optional datasets cache dir for eval",
+    )
+    parser.add_argument(
+        "--eval_num_workers",
+        type=int,
+        default=0,
+        help="Eval DataLoader workers",
+    )
+    parser.add_argument(
+        "--eval_device",
+        default=None,
+        help="Device for eval (defaults to --device)",
+    )
+    parser.add_argument(
+        "--skip_distill",
+        action="store_true",
+        help="Skip reparameterized distillation after head alignment/Fisher setup",
+    )
+    parser.add_argument(
+        "--distill_calib_samples",
+        type=int,
+        default=256,
+        help="Number of distillation sequences from calibration datasets",
+    )
+    parser.add_argument(
+        "--distill_inst_samples",
+        type=int,
+        default=0,
+        help="Number of distillation sequences from instruction dataset (0 = all)",
+    )
+    parser.add_argument(
+        "--distill_seq_len",
+        type=int,
+        default=512,
+        help="Sequence length for distillation",
+    )
+    parser.add_argument(
+        "--distill_batch_size",
+        type=int,
+        default=2,
+        help="Batch size for distillation",
+    )
+    parser.add_argument(
+        "--distill_epochs",
+        type=float,
+        default=1.0,
+        help="Number of distillation epochs (float allowed, e.g. 0.5)",
+    )
+    parser.add_argument(
+        "--distill_lr",
+        type=float,
+        default=1e-4,
+        help="Learning rate for distillation",
+    )
+    parser.add_argument(
+        "--distill_method",
+        choices=["reparam"],
+        default="reparam",
+        help="Distillation strategy (reparam only).",
+    )
+    parser.add_argument(
+        "--distill_kl_weight",
+        type=float,
+        default=1e-2,
+        help="Weight for KL loss on logits",
+    )
+    parser.add_argument(
+        "--distill_kl_temp",
+        type=float,
+        default=4.0,
+        help="Temperature for KL distillation on logits",
+    )
+    parser.add_argument(
+        "--distill_hidden_mse_weight",
+        type=float,
+        default=1.0,
+        help="Weight for hidden-state MSE in reparam distillation (0 disables it)",
+    )
+    parser.add_argument(
+        "--distill_attn_mse_weight",
+        type=float,
+        default=0.0,
+        help="Weight for auxiliary attention-output MSE in reparam distillation",
+    )
+    parser.add_argument(
+        "--distill_mlp_mse_weight",
+        type=float,
+        default=0.0,
+        help="Weight for auxiliary MLP-output MSE in reparam distillation",
+    )
+    parser.add_argument(
+        "--reparam_eta",
+        type=float,
+        default=1e-2,
+        help="Eta: ||lambda - lambda_gate||^2 regularizer weight for --distill_method reparam",
+    )
+    parser.add_argument(
+        "--reparam_gamma",
+        type=float,
+        default=1e-4,
+        help="Gamma: ||U - U0||^2 regularizer weight for --distill_method reparam",
+    )
+    parser.add_argument(
+        "--reparam_attn_reg_scale",
+        type=float,
+        default=1.0,
+        help="Relative scale applied to attention-parameter reparam regularizers",
+    )
+    parser.add_argument(
+        "--reparam_mlp_reg_scale",
+        type=float,
+        default=1.0,
+        help="Relative scale applied to MLP-parameter reparam regularizers",
+    )
+    parser.add_argument(
+        "--reparam_param_subset",
+        type=str,
+        choices=["all", "mlp", "attn"],
+        default="all",
+        help="Restrict reparam merge/recovery capacity to only this parameter family",
+    )
+    parser.add_argument(
+        "--norm_policy",
+        type=str,
+        choices=["hybrid", "merge_all", "copy_n1", "copy_n1_n2"],
+        default="hybrid",
+        help="Norm merge policy (default: hybrid)",
+    )
+    parser.add_argument(
+        "--distill_weight_decay",
+        type=float,
+        default=0.0,
+        help="Weight decay for distillation",
+    )
+    parser.add_argument(
+        "--distill_max_grad_norm",
+        type=float,
+        default=1.0,
+        help="Max grad norm for distillation",
+    )
+    parser.add_argument(
+        "--distill_grad_accum_steps",
+        type=int,
+        default=1,
+        help="Gradient accumulation steps for distillation",
+    )
+    parser.add_argument(
+        "--distill_log_steps",
+        type=int,
+        default=100,
+        help="Log distillation loss every N steps",
+    )
+    parser.add_argument(
+        "--distill_eval_every",
+        type=int,
+        default=0,
+        help="Evaluate PPL every N distill steps (0 = disable)",
+    )
+    parser.add_argument(
+        "--distill_eval_max_batches",
+        type=int,
+        default=None,
+        help="Max eval batches per dataset during distill (default: all)",
+    )
+    parser.add_argument(
+        "--distill_teacher_device",
+        default=None,
+        help="Device for teacher model during distillation (defaults to --device)",
+    )
+    parser.add_argument(
+        "--comm_enabled",
+        action="store_true",
+        help=(
+            "Enable commutator-style preconditioning before each progressive "
+            "cycle's fusion."
+        ),
+    )
+    parser.add_argument(
+        "--comm_include_cycle1",
+        action="store_true",
+        help="Run commutator preconditioning for cycle 1 as well (default: skip cycle 1).",
+    )
+    parser.add_argument(
+        "--comm_topk",
+        type=int,
+        default=1,
+        help="Top-K lowest-score pairs used as the commutator candidate set",
+    )
+    parser.add_argument(
+        "--comm_sample_eta",
+        type=float,
+        default=0.5,
+        help="Mixture weight between uniform and score-biased candidate sampling",
+    )
+    parser.add_argument(
+        "--comm_sample_dwce_scale",
+        type=float,
+        default=1.0,
+        help="Scale c in softmax(-c * score(i)) for commutator pair sampling",
+    )
+    parser.add_argument(
+        "--comm_temp",
+        type=float,
+        default=2.0,
+        help="Temperature for teacher-anchor KL in commutator preconditioning",
+    )
+    parser.add_argument(
+        "--comm_steps_ratio",
+        type=float,
+        default=0.1,
+        help="Run this fraction of distillation optimizer steps for commutator phase",
+    )
+    parser.add_argument(
+        "--comm_lr_scale",
+        type=float,
+        default=0.1,
+        help="Commutator LR = --distill_lr * this scale",
+    )
+    parser.add_argument(
+        "--comm_train_mode",
+        choices=["lora", "full"],
+        default="lora",
+        help=(
+            "Commutator trainable parameter mode: "
+            "'lora' updates LoRA adapters on sampled receiver layers; "
+            "'full' updates full receiver-layer weights."
+        ),
+    )
+    parser.add_argument(
+        "--comm_interaction_mode",
+        choices=["mse", "relative"],
+        default="relative",
+        help="Interaction loss form: plain MSE or relative MSE",
+    )
+    parser.add_argument(
+        "--comm_interaction_eps",
+        type=float,
+        default=1e-8,
+        help="Epsilon for relative commutator interaction normalization",
+    )
+    parser.add_argument(
+        "--comm_mu",
+        type=float,
+        default=None,
+        help=(
+            "Weight for interaction loss. Defaults to 0.1 for --comm_interaction_mode=mse "
+            "and 0.5 for --comm_interaction_mode=relative."
+        ),
+    )
+    parser.add_argument(
+        "--comm_mu_auto",
+        action="store_true",
+        help="Enable automatic mu scaling via gradient-norm balancing",
+    )
+    parser.add_argument(
+        "--comm_mu_auto_rho",
+        type=float,
+        default=0.1,
+        help="Target anchor-to-interaction gradient ratio constant for auto-mu",
+    )
+    parser.add_argument(
+        "--comm_mu_auto_eps",
+        type=float,
+        default=1e-8,
+        help="Numerical epsilon in auto-mu denominator",
+    )
+    parser.add_argument(
+        "--comm_log_steps",
+        type=int,
+        default=50,
+        help="Log commutator preconditioning loss every N optimizer steps",
+    )
+    parser.add_argument(
+        "--comm_skip_post_reselect",
+        action="store_true",
+        help=(
+            "Keep the pre-comm selected fusion pair and skip recomputing "
+            "selection after commutator preconditioning."
+        ),
+    )
+    parser.add_argument(
+        "--redistrib_teacher_source",
+        type=str,
+        choices=["base_model", "previous_cycle"],
+        default="base_model",
+        help=(
+            "Teacher source for commutator preconditioning teacher loading. "
+            "'base_model' uses --model for all cycles; "
+            "'previous_cycle' uses cycle-1 checkpoint (cycle 1 falls back to base_model)."
+        ),
+    )
+    parser.add_argument(
+        "--lora_epochs",
+        type=float,
+        default=1.0,
+        help="LoRA CE finetuning epochs after distill (0 = disable)",
+    )
+    parser.add_argument(
+        "--lora_rank",
+        type=int,
+        default=8,
+        help="LoRA rank (r)",
+    )
+    parser.add_argument(
+        "--lora_alpha",
+        type=float,
+        default=16.0,
+        help="LoRA alpha",
+    )
+    parser.add_argument(
+        "--lora_dropout",
+        type=float,
+        default=0.0,
+        help="LoRA dropout",
+    )
+    parser.add_argument(
+        "--lora_kl_enabled",
+        action="store_true",
+        help="Add KL regularization between pre/post LoRA logits",
+    )
+    parser.add_argument(
+        "--lora_kl_weight",
+        type=float,
+        default=1e-1,
+        help="KL weight for LoRA regularization",
+    )
+    parser.add_argument(
+        "--lora_kl_temp",
+        type=float,
+        default=4.0,
+        help="Temperature for LoRA KL regularization",
+    )
+    parser.add_argument(
+        "--lora_target_modules",
+        nargs="*",
+        default=[
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "down_proj",
+            "up_proj",
+        ],
+        help="Module name suffixes to LoRA-wrap",
+    )
+    parser.add_argument(
+        "--lora_respect_exclude_pairs",
+        action="store_true",
+        help=(
+            "When attaching LoRA adapters, skip linear modules under layers touched by "
+            "--exclude_pairs (i and i+1 for each excluded pair)."
+        ),
+    )
+    parser.add_argument(
+        "--lora_lr",
+        type=float,
+        default=1e-4,
+        help="Learning rate for LoRA finetuning",
+    )
+    parser.add_argument(
+        "--lora_weight_decay",
+        type=float,
+        default=0.0,
+        help="Weight decay for LoRA finetuning",
+    )
+    parser.add_argument(
+        "--lora_max_grad_norm",
+        type=float,
+        default=1.0,
+        help="Max grad norm for LoRA finetuning",
+    )
+    parser.add_argument(
+        "--lora_grad_accum_steps",
+        type=int,
+        default=1,
+        help="Gradient accumulation steps for LoRA finetuning",
+    )
+    parser.add_argument(
+        "--lora_log_steps",
+        type=int,
+        default=100,
+        help="Log LoRA loss every N steps",
+    )
+    parser.add_argument(
+        "--lora_eval_every",
+        type=int,
+        default=0,
+        help="Evaluate PPL every N LoRA steps (0 = disable)",
+    )
+    parser.add_argument(
+        "--lora_eval_max_batches",
+        type=int,
+        default=None,
+        help="Max eval batches per dataset during LoRA (default: all)",
+    )
+    parser.add_argument(
+        "--instruction_dataset",
+        default=None,
+        help="HF dataset name for alpaca-style instruction data",
+    )
+    parser.add_argument(
+        "--instruction_config",
+        default=None,
+        help="Optional instruction dataset config",
+    )
+    parser.add_argument(
+        "--instruction_split",
+        default="train",
+        help="Instruction dataset split",
+    )
+    parser.add_argument(
+        "--instruction_field_instruction",
+        default="instruction",
+        help="Instruction field name",
+    )
+    parser.add_argument(
+        "--instruction_field_input",
+        default="input",
+        help="Optional input field name",
+    )
+    parser.add_argument(
+        "--instruction_field_output",
+        default="output",
+        help="Response/output field name",
+    )
+    parser.add_argument(
+        "--auto_max_batches",
+        type=int,
+        default=0,
+        help="Max calibration batches for auto selection scoring (0 = all)",
+    )
+    parser.add_argument(
+        "--auto_metric",
+        type=str,
+        choices=[
+            "dwce",
+            "cosine",
+            "hybrid",
+            "hybrid_cosine",
+            "hybrid_global_rel",
+        ],
+        default="dwce",
+        help=(
+            "Auto pair scoring metric. 'dwce' uses downstream-weighted composition error; "
+            "'cosine' uses average token-level cosine distance between adjacent layer outputs; "
+            "'hybrid'/'hybrid_cosine' use DWCE to shortlist then adjacent cosine for final scoring; "
+            "'hybrid_global_rel' uses DWCE to shortlist then reranks by the change in "
+            "pair-to-final-layer cosine relation after surrogate fusion."
+        ),
+    )
+    parser.add_argument(
+        "--auto_cosine_topk",
+        type=int,
+        default=3,
+        help="Top-K DWCE candidates to rescore with cosine in --auto_metric=hybrid",
+    )
+    parser.add_argument(
+        "--auto_norm",
+        type=str,
+        choices=["relative", "none"],
+        default="relative",
+        help="Normalization mode for DWCE scoring (ignored for cosine)",
+    )
+    parser.add_argument(
+        "--auto_dwce_mode",
+        type=str,
+        choices=["separate", "shared"],
+        default="separate",
+        help=(
+            "DWCE implementation for auto scoring. "
+            "'separate' runs distinct Fisher and DWCE backward passes; "
+            "'shared' reuses one backward pass and replays DWCE with cached gradients."
+        ),
+    )
+    parser.add_argument(
+        "--num_progressive",
+        type=int,
+        default=0,
+        help="Number of progressive fusions (>0 required)",
+    )
+    parser.add_argument(
+        "--resume_from_cycle",
+        type=int,
+        default=0,
+        help=(
+            "Resume from this completed cycle index. When > 0, --model should point "
+            "to the saved full model directory for that cycle."
+        ),
+    )
+    parser.add_argument(
+        "--save_full_model_cycles",
+        nargs="*",
+        default=[],
+        help=(
+            "Cycle indices whose full models should be saved. Requesting cycle c "
+            "also saves cycle c-1 automatically (c=1 saves only cycle 1)."
+        ),
+    )
+    return parser.parse_args()
+
+
+def parse_exclude_pairs(exclude_raw: Optional[List[str]], num_pairs: int) -> List[int]:
+    """Parse --exclude_pairs into normalized pair indices for the current model.
+
+    Indices refer to the start of an adjacent pair (i, i+1) and must be in [0..N-2].
+    Negative indices count from the end (-1 = last pair).
+    """
+    if not exclude_raw:
+        return []
+    exclude: List[int] = []
+    for item in exclude_raw:
+        if item is None:
+            continue
+        for part in str(item).split(","):
+            part = part.strip()
+            if not part:
+                continue
+            try:
+                idx = int(part)
+            except ValueError as exc:
+                raise SystemExit("--exclude_pairs must contain integers.") from exc
+            if idx < 0:
+                idx = num_pairs + idx
+            if 0 <= idx < num_pairs:
+                exclude.append(idx)
+    return sorted(set(exclude))
+
+
+def parse_cycle_list(raw_values: Optional[List[str]]) -> List[int]:
+    if not raw_values:
+        return []
+    cycles: List[int] = []
+    for item in raw_values:
+        if item is None:
+            continue
+        for part in str(item).split(","):
+            part = part.strip()
+            if not part:
+                continue
+            try:
+                cycles.append(int(part))
+            except ValueError as exc:
+                raise SystemExit(
+                    "--save_full_model_cycles must contain integers."
+                ) from exc
+    return cycles
+
+
+def resolve_full_model_save_cycles(
+    requested_cycles: List[int], num_progressive: int
+) -> Set[int]:
+    resolved: Set[int] = set()
+    for cycle in requested_cycles:
+        if cycle <= 0 or cycle > num_progressive:
+            raise SystemExit(
+                "--save_full_model_cycles entries must be within [1, --num_progressive]."
+            )
+        resolved.add(cycle)
+        if cycle > 1:
+            resolved.add(cycle - 1)
+    return resolved
+
+
+def load_resume_metadata(model_path: str) -> Optional[Dict[str, object]]:
+    resume_meta_path = os.path.join(model_path, "resume_info.json")
+    if not os.path.exists(resume_meta_path):
+        return None
+    with open(resume_meta_path, "r", encoding="utf-8") as handle:
+        loaded = json.load(handle)
+    return loaded if isinstance(loaded, dict) else None
+
+
+def build_generator(seed: int) -> torch.Generator:
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(int(seed))
+    return generator
+
+
+def capture_rng_state() -> Dict[str, object]:
+    state: Dict[str, object] = {
+        "python_random_state": random.getstate(),
+        "torch_cpu_rng_state": torch.get_rng_state(),
+    }
+    if np is not None:
+        state["numpy_random_state"] = np.random.get_state()
+    if torch.cuda.is_available():
+        state["torch_cuda_rng_state_all"] = torch.cuda.get_rng_state_all()
+    return state
+
+
+def restore_rng_state(state: Dict[str, object]) -> None:
+    python_state = state.get("python_random_state")
+    if python_state is not None:
+        random.setstate(python_state)
+
+    numpy_state = state.get("numpy_random_state")
+    if numpy_state is not None and np is not None:
+        np.random.set_state(numpy_state)
+
+    torch_cpu_state = state.get("torch_cpu_rng_state")
+    if torch_cpu_state is not None:
+        torch.set_rng_state(torch_cpu_state)
+
+    torch_cuda_state = state.get("torch_cuda_rng_state_all")
+    if torch_cuda_state is not None and torch.cuda.is_available():
+        torch.cuda.set_rng_state_all(torch_cuda_state)
+
+
+def save_rng_state(path: str) -> None:
+    torch.save(capture_rng_state(), path)
+
+
+def load_rng_state(path: str) -> Optional[Dict[str, object]]:
+    if not os.path.exists(path):
+        return None
+    loaded = torch.load(path, map_location="cpu", weights_only=False)
+    return loaded if isinstance(loaded, dict) else None
+
+
+def configure_reproducibility(seed: int) -> None:
+    random.seed(seed)
+    if np is not None:
+        np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        if hasattr(torch.backends, "cudnn"):
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+    if hasattr(torch, "use_deterministic_algorithms"):
+        torch.use_deterministic_algorithms(True, warn_only=True)
+
+
+def save_loader_generator_state(
+    base_dir: str,
+    *,
+    distill_generator: Optional[torch.Generator] = None,
+    lora_generator: Optional[torch.Generator] = None,
+) -> None:
+    state: Dict[str, object] = {}
+    if distill_generator is not None:
+        state["distill_generator_state"] = distill_generator.get_state()
+    if lora_generator is not None:
+        state["lora_generator_state"] = lora_generator.get_state()
+    if state:
+        torch.save(state, os.path.join(base_dir, "loader_generators.pt"))
+
+
+def load_loader_generator_state(base_dir: str) -> Optional[Dict[str, object]]:
+    path = os.path.join(base_dir, "loader_generators.pt")
+    if not os.path.exists(path):
+        return None
+    loaded = torch.load(path, map_location="cpu")
+    return loaded if isinstance(loaded, dict) else None
+
+
+def resolve_layer_idx(
+    args: argparse.Namespace,
+    model,
+    layers: List[torch.nn.Module],
+    dataloader,
+    previous_scores,
+    start_index: int,
+    exclude_pairs: Set[int],
+):
+    layer_arg = str(getattr(args, "layer", "auto")).strip().lower()
+    selection_method = str(getattr(args, "selection_method", "dwce")).strip().lower()
+
+    if layer_arg != "auto":
+        try:
+            layer_idx = int(layer_arg)
+        except ValueError as exc:
+            raise SystemExit("--layer must be 'auto' or an integer index") from exc
+        num_pairs = max(len(layers) - 1, 0)
+        if layer_idx < 0:
+            layer_idx += num_pairs
+        if layer_idx in exclude_pairs:
+            raise SystemExit(f"--layer resolved to excluded pair index {layer_idx}")
+        return layer_idx, previous_scores, {"method": "manual", "exclude_pairs": sorted(exclude_pairs)}
+
+    if selection_method == "sequential":
+        num_pairs = len(layers) - 1
+        for layer_idx in range(max(start_index, 0), num_pairs):
+            if layer_idx not in exclude_pairs:
+                return layer_idx, previous_scores, {
+                    "method": "sequential",
+                    "start_index": max(start_index, 0),
+                    "exclude_pairs": sorted(exclude_pairs),
+                }
+        raise SystemExit("No eligible layer pairs remain after exclusions")
+
+    layer_idx, dwce_scores, dwce_meta = select_layer_auto(
+        model,
+        layers,
+        dataloader,
+        args,
+        previous_scores=previous_scores,
+        start_index=start_index,
+        exclude_pairs=exclude_pairs,
+    )
+    return layer_idx, dwce_scores, dwce_meta
+
+
+@dataclass
+class PreparedData:
+    calib_loader: torch.utils.data.DataLoader
+    calib_num_sequences: int
+    distill_loader: Optional[torch.utils.data.DataLoader]
+    distill_generator: Optional[torch.Generator]
+    distill_meta: Dict[str, object]
+    lora_loader: Optional[torch.utils.data.DataLoader]
+    lora_generator: Optional[torch.Generator]
+    lora_meta: Dict[str, object]
+    eval_datasets: List[str]
+    eval_configs: List[Optional[str]]
+    eval_dataloaders: Optional[Dict[str, torch.utils.data.DataLoader]]
+
+
+def resolve_eval_datasets(args: argparse.Namespace) -> Tuple[List[str], List[Optional[str]]]:
+    eval_datasets = args.eval_dataset or ["wikitext"]
+    eval_configs = args.eval_dataset_config or ["wikitext-2-raw-v1"]
+    eval_configs = ppl_eval._expand_dataset_configs(eval_datasets, eval_configs)
+    return eval_datasets, eval_configs
+
+
+def run_ppl_eval(
+    model_id_or_path: str,
+    eval_datasets: List[str],
+    eval_configs: List[Optional[str]],
+    args: argparse.Namespace,
+    prepared_eval_dataloaders: Optional[Dict[str, torch.utils.data.DataLoader]] = None,
+) -> Dict[str, float]:
+    eval_device = args.eval_device or args.device
+    dtype = get_dtype(args.dtype)
+
+    eval_model = load_causal_lm(
+        model_id_or_path,
+        torch_dtype=dtype,
+        trust_remote_code=args.trust_remote_code,
+    )
+
+    eval_model.to(eval_device)
+    if prepared_eval_dataloaders is not None:
+        results = ppl_eval.evaluate_ppl_dataloaders(
+            eval_model,
+            prepared_eval_dataloaders,
+            eval_device,
+            max_batches=args.eval_max_batches,
+        )
+    else:
+        eval_batch_size = args.eval_batch_size or args.batch_size
+        eval_tokenizer = AutoTokenizer.from_pretrained(
+            model_id_or_path, trust_remote_code=args.trust_remote_code
+        )
+        if eval_tokenizer.pad_token is None and eval_tokenizer.eos_token is not None:
+            eval_tokenizer.pad_token = eval_tokenizer.eos_token
+
+        results = ppl_eval.evaluate_ppl_datasets(
+            eval_model,
+            eval_tokenizer,
+            datasets=eval_datasets,
+            configs=eval_configs,
+            split=args.eval_split,
+            text_field=args.eval_text_field,
+            num_samples=args.eval_num_samples,
+            seq_len=args.eval_seq_len,
+            batch_size=eval_batch_size,
+            device=eval_device,
+            seed=args.seed,
+            shuffle=False,
+            model_family=args.eval_model_family,
+            add_bos=args.eval_add_bos,
+            max_batches=args.eval_max_batches,
+            cache_dir=args.eval_cache_dir,
+            num_workers=args.eval_num_workers,
+        )
+
+    del eval_model
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    return results
+
+
+def build_calibration_dataloader(
+    args: argparse.Namespace, tokenizer
+) -> Tuple[List[str], List[torch.Tensor], torch.utils.data.DataLoader]:
+    if args.dataset:
+        datasets = list(args.dataset)
+        configs = expand_dataset_configs(datasets, list(args.dataset_config))
+        chunks: List[torch.Tensor] = []
+        for idx, (dataset_name, config) in enumerate(zip(datasets, configs)):
+            spec = SharedLMDataSpec(
+                dataset=dataset_name,
+                config=config,
+                split=args.dataset_split,
+                text_field=args.dataset_text_field,
+                seq_len=args.seq_len,
+                num_sequences=args.num_samples,
+                seed=args.seed + idx,
+            )
+            chunks.extend(build_chunks(spec, tokenizer))
+        if not chunks:
+            raise SystemExit("Not enough text to build token sequences.")
+        input_ids = torch.stack(chunks)
+        attention_mask = torch.ones_like(input_ids)
+        dataset = torch.utils.data.TensorDataset(input_ids, attention_mask)
+        dataloader = torch.utils.data.DataLoader(
+            dataset, batch_size=args.batch_size, shuffle=False
+        )
+        return [], chunks, dataloader
+
+    texts = load_texts(args)
+    if not texts:
+        raise SystemExit(
+            "No calibration text found. Provide --dataset, --text, or --text_file."
+        )
+
+    chunks = build_token_chunks(texts, tokenizer, args.seq_len, args.num_samples)
+    if not chunks:
+        raise SystemExit("Not enough text to build token sequences.")
+
+    input_ids = torch.stack(chunks)
+    attention_mask = torch.ones_like(input_ids)
+    dataset = torch.utils.data.TensorDataset(input_ids, attention_mask)
+    dataloader = torch.utils.data.DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=False
+    )
+    return texts, chunks, dataloader
+
+
+def prepare_distillation_data(
+    args: argparse.Namespace, tokenizer, include_instruction: bool = True
+) -> Tuple[Optional[torch.utils.data.DataLoader], Optional[torch.Generator], Dict[str, object]]:
+    if (
+        include_instruction
+        and args.distill_inst_samples != 0
+        and not args.instruction_dataset
+    ):
+        print(
+            "Warning: --distill_inst_samples > 0 but no --instruction_dataset "
+            "provided; instruction distillation will be skipped."
+        )
+
+    calib_texts: List[str] = []
+    calib_dataset = None
+    if args.target_tokens > 0 and args.dataset:
+        datasets = list(args.dataset)
+        configs = expand_dataset_configs(datasets, list(args.dataset_config))
+        per_dataset = args.target_tokens // len(datasets)
+        remainder = args.target_tokens % len(datasets)
+        calib_chunks: List[torch.Tensor] = []
+        for idx, (dataset_name, config) in enumerate(zip(datasets, configs)):
+            dataset_tokens = per_dataset + (remainder if idx == 0 else 0)
+            spec = SharedLMDataSpec(
+                dataset=dataset_name,
+                config=config,
+                split=args.dataset_split,
+                text_field=args.dataset_text_field,
+                seq_len=args.distill_seq_len,
+                target_tokens=dataset_tokens,
+                seed=args.seed + 17 + idx,
+            )
+            calib_chunks.extend(build_chunks(spec, tokenizer))
+        if calib_chunks:
+            input_ids = torch.stack(calib_chunks)
+            attention_mask = torch.ones_like(input_ids)
+            calib_dataset = torch.utils.data.TensorDataset(input_ids, attention_mask)
+    else:
+        calib_texts = load_texts_from_datasets(
+            datasets=list(args.dataset),
+            configs=expand_dataset_configs(list(args.dataset), list(args.dataset_config)),
+            split=args.dataset_split,
+            text_field=args.dataset_text_field,
+            num_samples=args.distill_calib_samples,
+            seed=args.seed + 17,
+        )
+    inst_records = []
+    if include_instruction:
+        inst_records = load_instruction_records(args, args.distill_inst_samples)
+
+    distill_datasets = []
+    if calib_dataset is not None:
+        distill_datasets.append(calib_dataset)
+    elif calib_texts:
+        calib_records = [{"text": text} for text in calib_texts]
+        distill_datasets.append(
+            FixedSeqDataset(calib_records, tokenizer, args.distill_seq_len)
+        )
+    if inst_records:
+        distill_datasets.append(
+            FixedSeqDataset(inst_records, tokenizer, args.distill_seq_len)
+        )
+
+    distill_meta: Dict[str, object] = {
+        "calib_texts": len(calib_texts),
+        "calib_sequences": len(calib_dataset) if calib_dataset is not None else len(calib_texts),
+        "inst_sequences": len(inst_records),
+        "total_sequences": 0,
+    }
+
+    if not distill_datasets:
+        return None, None, distill_meta
+    if len(distill_datasets) == 1:
+        distill_dataset = distill_datasets[0]
+    else:
+        distill_dataset = torch.utils.data.ConcatDataset(distill_datasets)
+
+    distill_meta["total_sequences"] = len(distill_dataset)
+    distill_generator = build_generator(
+        args.seed + 1000 + (1000000 if include_instruction else 0)
+    )
+    distill_loader = torch.utils.data.DataLoader(
+        distill_dataset,
+        batch_size=args.distill_batch_size,
+        shuffle=True,
+        generator=distill_generator,
+    )
+    return distill_loader, distill_generator, distill_meta
+
+
+def prepare_eval_dataloaders(
+    args: argparse.Namespace,
+    tokenizer,
+    model: torch.nn.Module,
+    eval_datasets: List[str],
+    eval_configs: List[Optional[str]],
+) -> Optional[Dict[str, torch.utils.data.DataLoader]]:
+    needs_eval = (not args.skip_eval) or (
+        (not args.skip_distill and args.distill_eval_every)
+        or (args.lora_epochs > 0 and args.lora_eval_every)
+    )
+    if not needs_eval:
+        return None
+
+    eval_batch_size = args.eval_batch_size or args.batch_size
+    resolved_family = args.eval_model_family
+    if resolved_family == "auto":
+        resolved_family = ppl_eval._infer_model_family(model)
+
+    return ppl_eval.prepare_ppl_dataloaders(
+        tokenizer=tokenizer,
+        datasets=eval_datasets,
+        configs=eval_configs,
+        split=args.eval_split,
+        text_field=args.eval_text_field,
+        num_samples=args.eval_num_samples,
+        seq_len=args.eval_seq_len,
+        batch_size=eval_batch_size,
+        seed=args.seed,
+        shuffle=False,
+        model_family=resolved_family,
+        add_bos=args.eval_add_bos,
+        cache_dir=args.eval_cache_dir,
+        num_workers=args.eval_num_workers,
+        model=model,
+    )
+
+
+def prepare_all_data(
+    args: argparse.Namespace,
+    tokenizer,
+    model: torch.nn.Module,
+    eval_datasets: List[str],
+    eval_configs: List[Optional[str]],
+    loader_generator_state: Optional[Dict[str, object]] = None,
+) -> PreparedData:
+    texts, chunks, calib_loader = build_calibration_dataloader(args, tokenizer)
+    calib_num_sequences = len(chunks)
+    del texts
+    del chunks
+
+    distill_loader = None
+    distill_generator = None
+    distill_meta: Dict[str, object] = {
+        "calib_texts": 0,
+        "calib_sequences": 0,
+        "inst_sequences": 0,
+        "total_sequences": 0,
+    }
+    lora_loader = None
+    lora_generator = None
+    lora_meta: Dict[str, object] = {
+        "calib_texts": 0,
+        "calib_sequences": 0,
+        "inst_sequences": 0,
+        "total_sequences": 0,
+    }
+    if (not args.skip_distill) or bool(getattr(args, "comm_enabled", False)):
+        distill_loader, distill_generator, distill_meta = prepare_distillation_data(
+            args, tokenizer, include_instruction=False
+        )
+        if (
+            distill_generator is not None
+            and loader_generator_state is not None
+            and loader_generator_state.get("distill_generator_state") is not None
+        ):
+            distill_generator.set_state(loader_generator_state["distill_generator_state"])
+    if args.lora_epochs > 0:
+        lora_loader, lora_generator, lora_meta = prepare_distillation_data(
+            args, tokenizer, include_instruction=True
+        )
+        if (
+            lora_generator is not None
+            and loader_generator_state is not None
+            and loader_generator_state.get("lora_generator_state") is not None
+        ):
+            lora_generator.set_state(loader_generator_state["lora_generator_state"])
+
+    eval_dataloaders = prepare_eval_dataloaders(
+        args, tokenizer, model, eval_datasets, eval_configs
+    )
+
+    return PreparedData(
+        calib_loader=calib_loader,
+        calib_num_sequences=calib_num_sequences,
+        distill_loader=distill_loader,
+        distill_generator=distill_generator,
+        distill_meta=distill_meta,
+        lora_loader=lora_loader,
+        lora_generator=lora_generator,
+        lora_meta=lora_meta,
+        eval_datasets=eval_datasets,
+        eval_configs=eval_configs,
+        eval_dataloaders=eval_dataloaders,
+    )
+
+
+def evaluate_ppl_model(
+    model: torch.nn.Module,
+    tokenizer,
+    eval_datasets: List[str],
+    eval_configs: List[Optional[str]],
+    args: argparse.Namespace,
+    max_batches: Optional[int] = None,
+    prepared_eval_dataloaders: Optional[Dict[str, torch.utils.data.DataLoader]] = None,
+) -> Dict[str, float]:
+    eval_device = args.eval_device or args.device
+    prev_mode = model.training
+    try:
+        prev_device = next(model.parameters()).device
+    except StopIteration:
+        prev_device = torch.device(eval_device)
+
+    model.eval()
+    if str(prev_device) != eval_device:
+        model.to(eval_device)
+
+    if prepared_eval_dataloaders is not None:
+        results = ppl_eval.evaluate_ppl_dataloaders(
+            model,
+            prepared_eval_dataloaders,
+            eval_device,
+            max_batches=max_batches if max_batches is not None else args.eval_max_batches,
+        )
+    else:
+        eval_batch_size = args.eval_batch_size or args.batch_size
+        results = ppl_eval.evaluate_ppl_datasets(
+            model,
+            tokenizer,
+            datasets=eval_datasets,
+            configs=eval_configs,
+            split=args.eval_split,
+            text_field=args.eval_text_field,
+            num_samples=args.eval_num_samples,
+            seq_len=args.eval_seq_len,
+            batch_size=eval_batch_size,
+            device=eval_device,
+            seed=args.seed,
+            shuffle=False,
+            model_family=args.eval_model_family,
+            add_bos=args.eval_add_bos,
+            max_batches=max_batches if max_batches is not None else args.eval_max_batches,
+            cache_dir=args.eval_cache_dir,
+            num_workers=args.eval_num_workers,
+        )
+
+    if prev_mode:
+        model.train()
+    if str(prev_device) != eval_device:
+        model.to(prev_device)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    return results
+
+
+def has_post_fusion_data(
+    distill_loader: Optional[torch.utils.data.DataLoader],
+    distill_meta: Optional[Dict[str, object]],
+) -> bool:
+    if distill_loader is None or distill_meta is None:
+        return False
+    return distill_meta.get("total_sequences", 0) > 0
+
+
+def summarize_gate_lambdas(gates: Dict[str, torch.Tensor]) -> Dict[str, object]:
+    if not gates:
+        return {"num_tensors": 0, "num_elements": 0}
+
+    total_sum = 0.0
+    total_elems = 0
+    per_tensor_mean: Dict[str, Optional[float]] = {}
+    for name, gate in gates.items():
+        g = gate.detach().float()
+        if g.numel() == 0:
+            per_tensor_mean[name] = None
+            continue
+        per_tensor_mean[name] = float(g.mean().item())
+        total_sum += float(g.sum().item())
+        total_elems += int(g.numel())
+
+    global_mean = None if total_elems == 0 else total_sum / float(total_elems)
+    return {
+        "num_tensors": len(gates),
+        "num_elements": total_elems,
+        "global_mean": global_mean,
+        "per_tensor_mean": per_tensor_mean,
+    }
+
+
+def compute_path_bytes(path: str) -> int:
+    if os.path.isfile(path):
+        return os.path.getsize(path)
+
+    total = 0
+    for root, _, files in os.walk(path):
+        for name in files:
+            file_path = os.path.join(root, name)
+            if os.path.islink(file_path):
+                continue
+            try:
+                total += os.path.getsize(file_path)
+            except OSError:
+                continue
+    return total
+
+
+def save_stage_checkpoint(
+    model: torch.nn.Module,
+    tokenizer,
+    stage_dir: str,
+    stage_name: str,
+    ppl_results: Optional[Dict[str, float]],
+) -> Dict[str, object]:
+    os.makedirs(stage_dir, exist_ok=True)
+    colon_modules = find_colon_modules(model)
+    if colon_modules:
+        raise RuntimeError(
+            "Unexpected module names with ':' detected before save: "
+            f"{', '.join(colon_modules)}."
+        )
+
+    model.save_pretrained(stage_dir)
+    tokenizer.save_pretrained(stage_dir)
+
+    stage_meta = {
+        "stage": stage_name,
+        "path": stage_dir,
+        "weight_bytes": compute_path_bytes(stage_dir),
+        "post_ppl": ppl_results,
+    }
+    with open(
+        os.path.join(stage_dir, "stage_metrics.json"),
+        "w",
+        encoding="utf-8",
+    ) as handle:
+        json.dump(stage_meta, handle, indent=2)
+    return stage_meta
+
+
+def save_cycle_full_model(
+    model: torch.nn.Module,
+    tokenizer,
+    cycle_dir: str,
+    cycle_idx: int,
+    args: argparse.Namespace,
+    ppl_results: Optional[Dict[str, float]],
+) -> Dict[str, object]:
+    full_model_dir = os.path.join(cycle_dir, "full_model")
+    stage_meta = save_stage_checkpoint(
+        model=model,
+        tokenizer=tokenizer,
+        stage_dir=full_model_dir,
+        stage_name=f"cycle_{cycle_idx}_full_model",
+        ppl_results=ppl_results,
+    )
+    resume_meta = {
+        "base_model": getattr(args, "base_model_id", args.model),
+        "cycle": cycle_idx,
+        "output_dir": args.output_dir,
+        "layer_path": args.layer_path,
+        "rng_state": "rng_state.pt",
+        "loader_generators": "loader_generators.pt",
+    }
+    with open(
+        os.path.join(full_model_dir, "resume_info.json"),
+        "w",
+        encoding="utf-8",
+    ) as handle:
+        json.dump(resume_meta, handle, indent=2)
+    stage_meta["resume_info"] = "resume_info.json"
+    return stage_meta
+
+
+def run_lora_phase(
+    model: torch.nn.Module,
+    tokenizer,
+    eval_datasets: List[str],
+    eval_configs: List[Optional[str]],
+    args: argparse.Namespace,
+    lora_loader: Optional[torch.utils.data.DataLoader] = None,
+    lora_meta: Optional[Dict[str, object]] = None,
+    eval_dataloaders: Optional[Dict[str, torch.utils.data.DataLoader]] = None,
+    cycle_idx: Optional[int] = None,
+    num_cycles: Optional[int] = None,
+) -> List[Dict[str, object]]:
+    lora_eval_history: List[Dict[str, object]] = []
+    if args.lora_epochs <= 0:
+        return lora_eval_history
+    if not has_post_fusion_data(lora_loader, lora_meta):
+        print("No post-fusion sequences built; skipping LoRA finetuning.")
+        return lora_eval_history
+
+    lora_ce_finetune(
+        model=model,
+        dataloader=lora_loader,
+        eval_tokenizer=tokenizer,
+        eval_datasets=eval_datasets,
+        eval_configs=eval_configs,
+        eval_history=lora_eval_history,
+        args=args,
+        eval_dataloaders=eval_dataloaders,
+        progressive_cycle=cycle_idx,
+        progressive_total=num_cycles,
+    )
+    return lora_eval_history
+
+
+def run_progressive(
+    args: argparse.Namespace,
+    model: torch.nn.Module,
+    tokenizer,
+    prepared: PreparedData,
+) -> None:
+    eval_datasets = prepared.eval_datasets
+    eval_configs = prepared.eval_configs
+
+    dataloader = prepared.calib_loader
+    num_sequences = prepared.calib_num_sequences
+
+    model.to(args.device)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    progressive_meta_path = os.path.join(args.output_dir, "progressive_metadata.json")
+    existing_meta: Dict[str, object] = {}
+    if args.resume_from_cycle > 0 and os.path.exists(progressive_meta_path):
+        with open(progressive_meta_path, "r", encoding="utf-8") as handle:
+            loaded_meta = json.load(handle)
+        if isinstance(loaded_meta, dict):
+            existing_meta = loaded_meta
+
+    bootstrap_meta = {
+        "base_model": getattr(args, "base_model_id", args.model),
+        "num_progressive": args.num_progressive,
+        "layer_path": args.layer_path,
+        "resume_from_cycle": args.resume_from_cycle,
+        "save_full_model_cycles": sorted(args.full_model_save_cycles),
+        "cycles": (
+            existing_meta.get("cycles", [])
+            if isinstance(existing_meta.get("cycles"), list)
+            else []
+        ),
+    }
+    with open(
+        progressive_meta_path,
+        "w",
+        encoding="utf-8",
+    ) as handle:
+        json.dump(bootstrap_meta, handle, indent=2)
+
+    pre_eval = None
+    if not args.skip_eval:
+        pre_eval = evaluate_ppl_model(
+            model,
+            tokenizer,
+            eval_datasets,
+            eval_configs,
+            args,
+            prepared_eval_dataloaders=prepared.eval_dataloaders,
+        )
+        print("Pre-pruning perplexity:")
+        for dataset_name, ppl in pre_eval.items():
+            print(f"{dataset_name}: {ppl:.4f}")
+
+    parent, name, container = find_layer_container(model, args.layer_path)
+    layers = list(container)
+    if args.num_progressive > (len(layers) - 1 + args.resume_from_cycle):
+        raise SystemExit(
+            f"--num_progressive ({args.num_progressive}) exceeds available pairs "
+            f"after resume offset ({len(layers) - 1 + args.resume_from_cycle})"
+        )
+
+    dwce_scores = None
+    dwce_meta = None
+    last_fused_idx = 0
+    cycle_summaries: List[Dict[str, object]] = []
+    existing_cycles = existing_meta.get("cycles", [])
+    if isinstance(existing_cycles, list):
+        for entry in existing_cycles:
+            if not isinstance(entry, dict):
+                continue
+            cycle_value = entry.get("cycle")
+            if isinstance(cycle_value, int) and cycle_value <= args.resume_from_cycle:
+                cycle_summaries.append(entry)
+
+    comm_enabled = bool(getattr(args, "comm_enabled", False))
+    comm_teacher_model = None
+    comm_teacher_cycle: Optional[int] = None
+    teacher_device = args.distill_teacher_device or args.device
+    previous_cycle_teacher_model = None
+    previous_cycle_teacher_cycle: Optional[int] = None
+
+    def _release_comm_teacher() -> None:
+        nonlocal comm_teacher_model, comm_teacher_cycle
+        if comm_teacher_model is not None:
+            del comm_teacher_model
+            comm_teacher_model = None
+            comm_teacher_cycle = None
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+    def _release_previous_cycle_teacher() -> None:
+        nonlocal previous_cycle_teacher_model, previous_cycle_teacher_cycle
+        if previous_cycle_teacher_model is not None:
+            del previous_cycle_teacher_model
+            previous_cycle_teacher_model = None
+            previous_cycle_teacher_cycle = None
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+    def _snapshot_previous_cycle_teacher(cycle_idx: int) -> None:
+        nonlocal previous_cycle_teacher_model, previous_cycle_teacher_cycle
+        _release_previous_cycle_teacher()
+        previous_cycle_teacher_model = copy.deepcopy(model)
+        previous_cycle_teacher_model.to(teacher_device)
+        previous_cycle_teacher_model.eval()
+        previous_cycle_teacher_cycle = cycle_idx
+
+    def _get_previous_cycle_teacher(
+        cycle_idx: int,
+    ) -> Tuple[Optional[torch.nn.Module], str, Optional[int]]:
+        prev_cycle = cycle_idx - 1
+        if prev_cycle <= 0:
+            return None, "base_model", 0
+        if (
+            previous_cycle_teacher_model is not None
+            and previous_cycle_teacher_cycle == prev_cycle
+        ):
+            return previous_cycle_teacher_model, "previous_cycle_memory", prev_cycle
+        teacher_model = load_progressive_model(
+            getattr(args, "base_model_id", args.model),
+            args.output_dir,
+            cycle=prev_cycle,
+            device=teacher_device,
+            dtype=args.dtype,
+            trust_remote_code=args.trust_remote_code,
+            layer_path=args.layer_path,
+        )
+        teacher_model.eval()
+        return teacher_model, "previous_cycle_disk", prev_cycle
+
+    def _get_comm_teacher(cycle_idx: int) -> Tuple[Optional[torch.nn.Module], str, Optional[int]]:
+        nonlocal comm_teacher_model, comm_teacher_cycle
+        if not comm_enabled:
+            return None, "disabled", None
+
+        source = str(getattr(args, "redistrib_teacher_source", "base_model"))
+        if source == "base_model":
+            if comm_teacher_model is None:
+                print(
+                    "[comm] Loading fixed base teacher for anchor loss "
+                    f"(device={teacher_device})."
+                )
+                comm_teacher_model = AutoModelForCausalLM.from_pretrained(
+                    getattr(args, "base_model_id", args.model),
+                    torch_dtype=get_dtype(args.dtype),
+                    trust_remote_code=args.trust_remote_code,
+                )
+                comm_teacher_model.to(teacher_device)
+                comm_teacher_model.eval()
+                comm_teacher_cycle = 0
+            return comm_teacher_model, "base_model", 0
+
+        prev_cycle = cycle_idx - 1
+        if prev_cycle <= 0:
+            if comm_teacher_model is None or comm_teacher_cycle != 0:
+                _release_comm_teacher()
+                print(
+                    "[comm] --redistrib_teacher_source=previous_cycle but cycle 1 "
+                    "has no prior checkpoint; using base teacher."
+                )
+                comm_teacher_model = AutoModelForCausalLM.from_pretrained(
+                    getattr(args, "base_model_id", args.model),
+                    torch_dtype=get_dtype(args.dtype),
+                    trust_remote_code=args.trust_remote_code,
+                )
+                comm_teacher_model.to(teacher_device)
+                comm_teacher_model.eval()
+                comm_teacher_cycle = 0
+            return comm_teacher_model, "base_model", 0
+
+        if (
+            previous_cycle_teacher_model is not None
+            and previous_cycle_teacher_cycle == prev_cycle
+        ):
+            if comm_teacher_model is not previous_cycle_teacher_model:
+                _release_comm_teacher()
+                comm_teacher_model = previous_cycle_teacher_model
+                comm_teacher_cycle = prev_cycle
+            return comm_teacher_model, "previous_cycle_memory", prev_cycle
+
+        if comm_teacher_model is None or comm_teacher_cycle != prev_cycle:
+            _release_comm_teacher()
+            print(
+                "[comm] Loading teacher from previous cycle "
+                f"{prev_cycle} (device={teacher_device})."
+            )
+            comm_teacher_model = load_progressive_model(
+                getattr(args, "base_model_id", args.model),
+                args.output_dir,
+                cycle=prev_cycle,
+                device=teacher_device,
+                dtype=args.dtype,
+                trust_remote_code=args.trust_remote_code,
+                layer_path=args.layer_path,
+            )
+            comm_teacher_model.eval()
+            comm_teacher_cycle = prev_cycle
+        return comm_teacher_model, "previous_cycle_disk", prev_cycle
+
+    if args.resume_from_cycle > 0:
+        _snapshot_previous_cycle_teacher(args.resume_from_cycle)
+
+    start_cycle = args.resume_from_cycle + 1
+    for cycle_idx in range(start_cycle, args.num_progressive + 1):
+        print(f"[progressive] Cycle {cycle_idx}/{args.num_progressive}")
+        run_comm = comm_enabled and (
+            cycle_idx > 1 or bool(getattr(args, "comm_include_cycle1", False))
+        )
+        comm_stats: Dict[str, object] = {"enabled": False}
+        comm_post_eval = None
+        if run_comm:
+            # Preconditioning updates model weights, so DWCE reuse is unreliable.
+            start_index = 0
+            reuse_scores = None
+        else:
+            start_index = last_fused_idx if cycle_idx > 1 else 0
+            reuse_scores = dwce_scores
+        exclude_pairs = set(parse_exclude_pairs(args.exclude_pairs, max(len(layers) - 1, 0)))
+        layer_idx, dwce_scores, dwce_meta = resolve_layer_idx(
+            args,
+            model,
+            layers,
+            dataloader,
+            reuse_scores,
+            start_index,
+            exclude_pairs,
+        )
+
+        if run_comm:
+            dwce_scores_pre_comm = dwce_scores
+            if prepared.calib_loader is None:
+                print(
+                    "[comm] Enabled but no calibration sequences were built; skipping."
+                )
+            else:
+                (
+                    comm_teacher_model_loaded,
+                    comm_teacher_source,
+                    comm_teacher_cycle_idx,
+                ) = _get_comm_teacher(cycle_idx)
+                if comm_teacher_model_loaded is None:
+                    raise RuntimeError("comm_enabled but teacher model was not loaded.")
+                comm_stats = commutator_precondition(
+                    student_model=model,
+                    student_layers=layers,
+                    teacher_model=comm_teacher_model_loaded,
+                    dataloader=prepared.calib_loader,
+                    dwce_scores=dwce_scores_pre_comm,
+                    exclude_pairs=exclude_pairs,
+                    args=args,
+                    progressive_cycle=cycle_idx,
+                    progressive_total=args.num_progressive,
+                )
+                if comm_stats.get("enabled"):
+                    comm_stats["teacher_source"] = comm_teacher_source
+                    comm_stats["teacher_cycle"] = comm_teacher_cycle_idx
+                    comm_stats["dwce_scores_pre"] = dwce_scores_pre_comm
+                    print(
+                        "[comm] Done:"
+                        f" opt_steps={comm_stats.get('opt_steps')}"
+                        f" lr={comm_stats.get('lr')}"
+                    )
+                    if not args.skip_eval:
+                        comm_post_eval = evaluate_ppl_model(
+                            model,
+                            tokenizer,
+                            eval_datasets,
+                            eval_configs,
+                            args,
+                            prepared_eval_dataloaders=prepared.eval_dataloaders,
+                        )
+                        comm_stats["post_ppl"] = comm_post_eval
+                        print(f"[progressive] Cycle {cycle_idx} post-comm perplexity:")
+                        for dataset_name, ppl in comm_post_eval.items():
+                            print(f"{dataset_name}: {ppl:.4f}")
+
+                    if bool(getattr(args, "comm_skip_post_reselect", False)):
+                        comm_stats["post_selection_recomputed"] = False
+                        comm_stats["selected_layer_post"] = int(layer_idx)
+                        print(
+                            "[comm] Keeping pre-comm DWCE pair selection for fusion."
+                        )
+                    else:
+                        print(
+                            "[comm] Recomputing DWCE after preconditioning for fusion selection."
+                        )
+                        layer_idx, dwce_scores, dwce_meta = resolve_layer_idx(
+                            args,
+                            model,
+                            layers,
+                            dataloader,
+                            None,
+                            0,
+                            exclude_pairs,
+                        )
+                        comm_stats["post_selection_recomputed"] = True
+                        comm_stats["selected_layer_post"] = int(layer_idx)
+
+        if layer_idx < 0 or layer_idx >= len(layers) - 1:
+            raise SystemExit("--layer must be in [0, num_layers-2]")
+
+        num_layers_before = len(layers)
+        layer_a = layers[layer_idx]
+        layer_b = layers[layer_idx + 1]
+
+        norm1_state = None
+        norm2_state = None
+        norm1, norm2, norm_names = get_norm_pair(layer_a)
+        if norm1 is not None:
+            norm1_state = clone_state_dict(norm1)
+        if norm2 is not None:
+            norm2_state = clone_state_dict(norm2)
+
+        attn_a = find_attention_module(layer_a)
+        attn_b = find_attention_module(layer_b)
+        hidden_size = getattr(model.config, "hidden_size", None)
+        if hidden_size is None:
+            hidden_size = getattr(model.config, "n_embd", None)
+        if hidden_size is None:
+            raise SystemExit("Model config missing hidden_size/n_embd")
+
+        no_head_permute_merge = bool(
+            getattr(args, "no_head_permute_merge", False)
+            or getattr(args, "no_head_permute", False)
+        )
+        if no_head_permute_merge:
+            print("[fuse] Head permutation disabled; merging with original head order.")
+        else:
+            mean_a, mean_b, num_heads, num_kv_heads, head_dim = compute_head_means(
+                model,
+                attn_a,
+                attn_b,
+                dataloader,
+                args.device,
+                hidden_size,
+            )
+
+            perm = build_head_permutation(
+                mean_a,
+                mean_b,
+                num_heads=num_heads,
+                num_kv_heads=num_kv_heads,
+                eps=args.eps,
+            )
+            permute_attention_heads(
+                attn_b, perm, num_heads, num_kv_heads, head_dim=head_dim
+            )
+
+        fisher_sums, num_batches, param_numels = compute_fisher(
+            model,
+            layer_a,
+            layer_b,
+            dataloader,
+            fisher_mode=args.fisher_mode,
+            device=args.device,
+        )
+
+        distill_ready = has_post_fusion_data(
+            prepared.distill_loader, prepared.distill_meta
+        )
+        teacher_cycle = cycle_idx - 1
+        teacher_source = "previous_cycle" if teacher_cycle > 0 else "base_model"
+
+        merge_method = "fisher"
+        distill_method = str(getattr(args, "distill_method", "reparam"))
+        reparam_stats: Optional[Dict[str, object]] = None
+        reparam_gate_summary: Optional[Dict[str, object]] = None
+
+        needs_teacher_for_reparam = (
+            (not args.skip_distill)
+            and distill_ready
+            and float(args.distill_epochs) > 0.0
+        )
+        teacher_model = None
+        teacher_parent = None
+        teacher_layer_attr = None
+        teacher_layers: Optional[List[torch.nn.Module]] = None
+        teacher_from_cache = False
+        if needs_teacher_for_reparam:
+            teacher_model, teacher_source, teacher_cycle = _get_previous_cycle_teacher(
+                cycle_idx
+            )
+            teacher_from_cache = (
+                teacher_source == "previous_cycle_memory"
+                and teacher_model is previous_cycle_teacher_model
+            )
+            if teacher_model is None:
+                teacher_model = load_causal_lm(
+                    getattr(args, "base_model_id", args.model),
+                    torch_dtype=get_dtype(args.dtype),
+                    trust_remote_code=args.trust_remote_code,
+                    cache_dir=args.model_cache_dir,
+                )
+                teacher_model.to(teacher_device)
+                teacher_model.eval()
+                teacher_source = "base_model"
+                teacher_cycle = 0
+            teacher_parent, teacher_layer_attr, teacher_container = find_layer_container(
+                teacher_model, args.layer_path
+            )
+            teacher_layers = list(teacher_container)
+
+        do_reparam = (
+            (not args.skip_distill)
+            and distill_ready
+            and prepared.distill_loader is not None
+        )
+        if (not args.skip_distill) and not do_reparam:
+            print("[reparam] No distillation sequences built; skipping reparam distill.")
+        distill_post_eval = None
+        if do_reparam:
+            lambda_source = "fisher_prior"
+            reparam_gate_targets: Dict[str, object] = compute_fisher_gate_priors(
+                layer_a=layer_a,
+                layer_b=layer_b,
+                fisher_a=fisher_sums[0],
+                fisher_b=fisher_sums[1],
+                num_batches=num_batches,
+                numels_a=param_numels[0],
+                numels_b=param_numels[1],
+                fisher_mode=args.fisher_mode,
+                eps=float(args.eps),
+            )
+            if not reparam_gate_targets:
+                raise SystemExit("[reparam] No mergeable parameters found; cannot continue.")
+            if float(args.distill_epochs) > 0.0 and (
+                teacher_model is None or teacher_layers is None
+            ):
+                raise SystemExit("--distill_method reparam requires a teacher model.")
+            print(
+                f"[reparam] Cycle {cycle_idx}: training U + gates for pair "
+                f"{layer_idx}-{layer_idx + 1} (epochs={args.distill_epochs}, "
+                f"hidden_mse_w={args.distill_hidden_mse_weight}, "
+                f"attn_mse_w={args.distill_attn_mse_weight}, "
+                f"mlp_mse_w={args.distill_mlp_mse_weight}, "
+                f"eta={args.reparam_eta}, gamma={args.reparam_gamma}, "
+                f"attn_reg_scale={args.reparam_attn_reg_scale}, "
+                f"mlp_reg_scale={args.reparam_mlp_reg_scale}, "
+                f"param_subset={args.reparam_param_subset}, "
+                f"lambda_init={lambda_source})."
+            )
+            merged, final_gates, reparam_stats = distill_reparam_merge(
+                student_model=model,
+                student_parent=parent,
+                student_layer_attr=name,
+                student_layers=layers,
+                teacher_model=teacher_model,
+                teacher_parent=teacher_parent,
+                teacher_layer_attr=teacher_layer_attr,
+                teacher_layers=teacher_layers,
+                layer_idx=layer_idx,
+                gate_lambdas=reparam_gate_targets,
+                dataloader=prepared.distill_loader,
+                args=args,
+                progressive_cycle=cycle_idx,
+                progressive_total=args.num_progressive,
+            )
+            reparam_gate_summary = summarize_gate_lambdas(final_gates)
+            merge_method = "reparam"
+            if reparam_stats is not None:
+                reparam_stats["lambda_init"] = lambda_source
+        else:
+            merged = merge_layers(
+                layer_a,
+                layer_b,
+                fisher_sums[0],
+                fisher_sums[1],
+                num_batches,
+                param_numels[0],
+                param_numels[1],
+                fisher_mode=args.fisher_mode,
+                eps=args.eps,
+            )
+
+            apply_norm_policy(
+                layer_a,
+                args.norm_policy,
+                norm1_state,
+                norm2_state,
+                norm_names,
+            )
+        if teacher_model is not None and not teacher_from_cache:
+            del teacher_model
+            teacher_model = None
+            teacher_parent = None
+            teacher_layer_attr = None
+            teacher_layers = None
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        new_container = drop_layer(container, layer_idx + 1)
+        setattr(parent, name, new_container)
+        decrement_config(model.config)
+        layers = list(new_container)
+
+        lora_post_eval = None
+        if (not args.skip_eval) and (not args.skip_distill) and do_reparam:
+            distill_post_eval = evaluate_ppl_model(
+                model,
+                tokenizer,
+                eval_datasets,
+                eval_configs,
+                args,
+                prepared_eval_dataloaders=prepared.eval_dataloaders,
+            )
+            print(f"[progressive] Cycle {cycle_idx} post-distill perplexity:")
+            for dataset_name, ppl in distill_post_eval.items():
+                print(f"{dataset_name}: {ppl:.4f}")
+
+        post_eval = None
+        if not args.skip_eval:
+            if distill_post_eval is not None:
+                post_eval = distill_post_eval
+            else:
+                post_eval = evaluate_ppl_model(
+                    model,
+                    tokenizer,
+                    eval_datasets,
+                    eval_configs,
+                    args,
+                    prepared_eval_dataloaders=prepared.eval_dataloaders,
+                )
+            print(f"[progressive] Cycle {cycle_idx} perplexity:")
+            for dataset_name, ppl in post_eval.items():
+                print(f"{dataset_name}: {ppl:.4f}")
+
+        cycle_dir = os.path.join(args.output_dir, f"cycle_{cycle_idx}")
+        os.makedirs(cycle_dir, exist_ok=True)
+        fused_layer_file = "fused_layer.pt"
+        fused_layer_path = os.path.join(cycle_dir, fused_layer_file)
+        torch.save(layers[layer_idx].state_dict(), fused_layer_path)
+
+        cycle_meta: Dict[str, object] = {
+            "cycle": cycle_idx,
+            "layer_merged": layer_idx,
+            "num_layers_before": num_layers_before,
+            "num_layers_after": num_layers_before - 1,
+            "fused_layer_state": fused_layer_file,
+            "dwce_score": dwce_scores[layer_idx] if dwce_scores else None,
+            "dwce_scores": dwce_scores,
+            "dwce_meta": dwce_meta,
+            "fisher_num_batches": num_batches,
+            "merge_method": merge_method,
+            "merged_params": merged,
+            "num_sequences": num_sequences,
+            "teacher_source": teacher_source,
+            "teacher_cycle": teacher_cycle,
+            "eval": {
+                "datasets": eval_datasets,
+                "configs": eval_configs,
+                "split": args.eval_split,
+                "num_samples": args.eval_num_samples,
+                "seq_len": args.eval_seq_len,
+                "post_ppl": post_eval,
+            },
+            "comm": comm_stats,
+            "distill": {
+                "enabled": not args.skip_distill,
+                "method": distill_method,
+                "calib_samples": args.distill_calib_samples,
+                "inst_samples": args.distill_inst_samples,
+                "seq_len": args.distill_seq_len,
+                "batch_size": args.distill_batch_size,
+                "epochs": args.distill_epochs,
+                "lr": args.distill_lr,
+                "kl_weight": args.distill_kl_weight,
+                "kl_temp": args.distill_kl_temp,
+                "hidden_mse_weight": args.distill_hidden_mse_weight,
+                "attn_mse_weight": args.distill_attn_mse_weight,
+                "mlp_mse_weight": args.distill_mlp_mse_weight,
+                "reparam_eta": args.reparam_eta,
+                "reparam_gamma": args.reparam_gamma,
+                "reparam_attn_reg_scale": args.reparam_attn_reg_scale,
+                "reparam_mlp_reg_scale": args.reparam_mlp_reg_scale,
+                "reparam_param_subset": args.reparam_param_subset,
+                "reparam_stats": reparam_stats,
+                "reparam_gate_summary": reparam_gate_summary,
+                "post_ppl": distill_post_eval,
+                "weight_decay": args.distill_weight_decay,
+                "max_grad_norm": args.distill_max_grad_norm,
+                "grad_accum_steps": args.distill_grad_accum_steps,
+                "instruction_dataset": args.instruction_dataset,
+                "instruction_config": args.instruction_config,
+                "instruction_split": args.instruction_split,
+            },
+            "lora": {
+                "enabled": args.lora_epochs > 0,
+                "seq_len": args.distill_seq_len,
+                "batch_size": args.distill_batch_size,
+                "epochs": args.lora_epochs,
+                "rank": args.lora_rank,
+                "alpha": args.lora_alpha,
+                "dropout": args.lora_dropout,
+                "target_modules": args.lora_target_modules,
+                "respect_exclude_pairs": args.lora_respect_exclude_pairs,
+                "kl_enabled": args.lora_kl_enabled,
+                "kl_weight": args.lora_kl_weight,
+                "kl_temp": args.lora_kl_temp,
+                "post_ppl": lora_post_eval,
+                "lr": args.lora_lr,
+                "weight_decay": args.lora_weight_decay,
+                "max_grad_norm": args.lora_max_grad_norm,
+                "grad_accum_steps": args.lora_grad_accum_steps,
+                "log_steps": args.lora_log_steps,
+                "eval_every": args.lora_eval_every,
+                "eval_max_batches": args.lora_eval_max_batches,
+            },
+            "norm_policy": args.norm_policy,
+        }
+
+        saved_full_model_dir = None
+        if cycle_idx in args.full_model_save_cycles:
+            cycle_meta["full_model_saved"] = True
+            cycle_meta["full_model"] = save_cycle_full_model(
+                model=model,
+                tokenizer=tokenizer,
+                cycle_dir=cycle_dir,
+                cycle_idx=cycle_idx,
+                args=args,
+                ppl_results=post_eval,
+            )
+            saved_full_model_dir = os.path.join(cycle_dir, "full_model")
+        else:
+            cycle_meta["full_model_saved"] = False
+
+        with open(
+            os.path.join(cycle_dir, "cycle_metadata.json"),
+            "w",
+            encoding="utf-8",
+        ) as handle:
+            json.dump(cycle_meta, handle, indent=2)
+
+        cycle_summaries.append(
+            {
+                "cycle": cycle_idx,
+                "layer_merged": layer_idx,
+                "dwce_score": dwce_scores[layer_idx] if dwce_scores else None,
+                "comm_post_ppl": comm_post_eval,
+                "distill_post_ppl": distill_post_eval,
+                "lora_post_ppl": lora_post_eval,
+                "post_ppl": post_eval,
+                "cycle_dir": f"cycle_{cycle_idx}",
+            }
+        )
+
+        last_fused_idx = layer_idx
+        _snapshot_previous_cycle_teacher(cycle_idx)
+
+        parent, name, container = find_layer_container(model, args.layer_path)
+        layers = list(container)
+        if dwce_scores:
+            dwce_scores = dwce_scores[: max(len(layers) - 1, 0)]
+
+        # Encourage allocator to release cached blocks between cycles.
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        gc.collect()
+
+        if saved_full_model_dir is not None:
+            save_rng_state(os.path.join(saved_full_model_dir, "rng_state.pt"))
+            save_loader_generator_state(
+                saved_full_model_dir,
+                distill_generator=prepared.distill_generator,
+                lora_generator=prepared.lora_generator,
+            )
+
+    _release_comm_teacher()
+    _release_previous_cycle_teacher()
+
+    final_pre_lora_eval = cycle_summaries[-1]["post_ppl"] if cycle_summaries else None
+    final_pre_lora_dir = f"{os.path.abspath(args.output_dir.rstrip(os.sep))}_final_pre_lora_hf"
+    final_pre_lora_meta = save_stage_checkpoint(
+        model=model,
+        tokenizer=tokenizer,
+        stage_dir=final_pre_lora_dir,
+        stage_name="final_pre_lora",
+        ppl_results=final_pre_lora_eval,
+    )
+
+    # Optional final LoRA finetune after all pruning cycles.
+    lora_eval_history: List[Dict[str, object]] = []
+    lora_post_eval = None
+    lora_ready = has_post_fusion_data(prepared.lora_loader, prepared.lora_meta)
+    if args.lora_epochs > 0:
+        if not lora_ready:
+            print("No post-fusion sequences built; skipping LoRA finetuning.")
+        else:
+            print(
+                f"[progressive] Running final LoRA finetuning (epochs={args.lora_epochs})."
+            )
+            lora_eval_history = run_lora_phase(
+                model=model,
+                tokenizer=tokenizer,
+                eval_datasets=eval_datasets,
+                eval_configs=eval_configs,
+                args=args,
+                lora_loader=prepared.lora_loader,
+                lora_meta=prepared.lora_meta,
+                eval_dataloaders=prepared.eval_dataloaders,
+                cycle_idx=args.num_progressive,
+                num_cycles=args.num_progressive,
+            )
+
+            if not args.skip_eval:
+                lora_post_eval = evaluate_ppl_model(
+                    model,
+                    tokenizer,
+                    eval_datasets,
+                    eval_configs,
+                    args,
+                    prepared_eval_dataloaders=prepared.eval_dataloaders,
+                )
+                print("[progressive] Post-LoRA perplexity:")
+                for dataset_name, ppl in lora_post_eval.items():
+                    print(f"{dataset_name}: {ppl:.4f}")
+
+            # Update final cycle metadata and summary with the post-LoRA PPL.
+            if cycle_summaries:
+                cycle_summaries[-1]["lora_post_ppl"] = lora_post_eval
+                if lora_post_eval is not None:
+                    cycle_summaries[-1]["post_ppl"] = lora_post_eval
+
+            final_cycle_dir = os.path.join(
+                args.output_dir, f"cycle_{args.num_progressive}"
+            )
+            final_cycle_meta_path = os.path.join(final_cycle_dir, "cycle_metadata.json")
+            if os.path.exists(final_cycle_meta_path):
+                with open(final_cycle_meta_path, "r", encoding="utf-8") as handle:
+                    final_cycle_meta = json.load(handle)
+
+                lora_meta_entry = final_cycle_meta.get("lora")
+                if not isinstance(lora_meta_entry, dict):
+                    lora_meta_entry = {}
+                    final_cycle_meta["lora"] = lora_meta_entry
+                lora_meta_entry["ran"] = True
+                lora_meta_entry["post_ppl"] = lora_post_eval
+                if lora_post_eval is not None and isinstance(
+                    final_cycle_meta.get("eval"), dict
+                ):
+                    final_cycle_meta["eval"]["post_ppl"] = lora_post_eval
+
+                if lora_eval_history:
+                    lora_path = os.path.join(final_cycle_dir, "ppl_over_lora.json")
+                    with open(lora_path, "w", encoding="utf-8") as handle:
+                        json.dump(lora_eval_history, handle, indent=2)
+                    lora_meta_entry["ppl_over_lora"] = "ppl_over_lora.json"
+
+                with open(final_cycle_meta_path, "w", encoding="utf-8") as handle:
+                    json.dump(final_cycle_meta, handle, indent=2)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    final_post_lora_meta = save_stage_checkpoint(
+        model=model,
+        tokenizer=tokenizer,
+        stage_dir=args.output_dir,
+        stage_name="final_post_lora" if lora_post_eval is not None else "final_model",
+        ppl_results=lora_post_eval,
+    )
+
+    progressive_meta = {
+        "base_model": getattr(args, "base_model_id", args.model),
+        "num_progressive": args.num_progressive,
+        "layer_path": args.layer_path,
+        "resume_from_cycle": args.resume_from_cycle,
+        "save_full_model_cycles": sorted(args.full_model_save_cycles),
+        "num_sequences": num_sequences,
+        "seq_len": args.seq_len,
+        "lora": {
+            "enabled": args.lora_epochs > 0,
+            "ran": args.lora_epochs > 0 and lora_ready,
+            "seq_len": args.distill_seq_len,
+            "batch_size": args.distill_batch_size,
+            "epochs": args.lora_epochs,
+            "rank": args.lora_rank,
+            "alpha": args.lora_alpha,
+            "dropout": args.lora_dropout,
+            "target_modules": args.lora_target_modules,
+            "respect_exclude_pairs": args.lora_respect_exclude_pairs,
+            "kl_enabled": args.lora_kl_enabled,
+            "kl_weight": args.lora_kl_weight,
+            "kl_temp": args.lora_kl_temp,
+            "post_ppl": lora_post_eval,
+            "ppl_over_lora": (
+                f"cycle_{args.num_progressive}/ppl_over_lora.json"
+                if lora_eval_history
+                else None
+            ),
+            "lr": args.lora_lr,
+            "weight_decay": args.lora_weight_decay,
+            "max_grad_norm": args.lora_max_grad_norm,
+            "grad_accum_steps": args.lora_grad_accum_steps,
+            "log_steps": args.lora_log_steps,
+            "eval_every": args.lora_eval_every,
+            "eval_max_batches": args.lora_eval_max_batches,
+        },
+        "artifacts": {
+            "final_pre_lora": final_pre_lora_meta,
+            "final_post_lora": final_post_lora_meta,
+        },
+        "eval": {
+            "datasets": eval_datasets,
+            "configs": eval_configs,
+            "split": args.eval_split,
+            "num_samples": args.eval_num_samples,
+            "seq_len": args.eval_seq_len,
+            "pre_ppl": pre_eval,
+            "post_ppl": cycle_summaries[-1]["post_ppl"] if cycle_summaries else None,
+        },
+        "cycles": cycle_summaries,
+        "final_num_layers": len(layers),
+    }
+
+    with open(
+        os.path.join(args.output_dir, "progressive_metadata.json"),
+        "w",
+        encoding="utf-8",
+    ) as handle:
+        json.dump(progressive_meta, handle, indent=2)
+
+    print(
+        f"[progressive] Completed {args.num_progressive} cycles. "
+        f"Final model saved to {args.output_dir}."
+    )
+
+
+def main() -> None:
+    args = parse_args()
+    if args.num_progressive <= 0:
+        raise SystemExit(
+            "Single-cycle mode has been removed. Pass --num_progressive > 0."
+        )
+    if args.resume_from_cycle < 0:
+        raise SystemExit("--resume_from_cycle must be >= 0.")
+    if args.resume_from_cycle >= args.num_progressive:
+        raise SystemExit("--resume_from_cycle must be smaller than --num_progressive.")
+    args.full_model_save_cycles = resolve_full_model_save_cycles(
+        parse_cycle_list(args.save_full_model_cycles),
+        args.num_progressive,
+    )
+    args.base_model_id = args.model
+    if args.resume_from_cycle > 0:
+        resume_meta = load_resume_metadata(args.model)
+        if resume_meta is None:
+            raise SystemExit(
+                "--resume_from_cycle requires --model to point to a saved cycle full model "
+                "directory containing resume_info.json."
+            )
+        resume_cycle = resume_meta.get("cycle")
+        if resume_cycle is not None and int(resume_cycle) != args.resume_from_cycle:
+            raise SystemExit(
+                "resume_info.json cycle does not match --resume_from_cycle."
+            )
+        base_model = resume_meta.get("base_model")
+        if isinstance(base_model, str) and base_model:
+            args.base_model_id = base_model
+    configure_reproducibility(args.seed)
+
+    eval_datasets, eval_configs = resolve_eval_datasets(args)
+    dtype = get_dtype(args.dtype)
+    model = load_causal_lm(
+        args.model,
+        torch_dtype=dtype,
+        trust_remote_code=args.trust_remote_code,
+        cache_dir=args.model_cache_dir,
+    )
+    loader_generator_state = None
+    if args.resume_from_cycle > 0:
+        rng_state_path = os.path.join(args.model, "rng_state.pt")
+        rng_state = load_rng_state(rng_state_path)
+        if rng_state is not None:
+            restore_rng_state(rng_state)
+        loader_generator_state = load_loader_generator_state(args.model)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model,
+        trust_remote_code=args.trust_remote_code,
+        cache_dir=args.model_cache_dir,
+    )
+    print(model)
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+        
+    # for llama?
+    model.config.use_cache = False
+
+    prepared = prepare_all_data(
+        args,
+        tokenizer,
+        model,
+        eval_datasets,
+        eval_configs,
+        loader_generator_state=loader_generator_state,
+    )
+    run_progressive(args, model, tokenizer, prepared)
+
+
+if __name__ == "__main__":
+    main()

src/fuse_layers_data.py

@@ -0,0 +1,280 @@
+#!/usr/bin/env python3
+"""Dataset and text helpers for fuse_layers."""
+
+import argparse
+from typing import Dict, List, Optional
+
+import torch
+
+try:
+    from datasets import load_dataset
+except Exception:  # pragma: no cover - optional dependency
+    load_dataset = None
+
+
+def guess_text_field(dataset) -> str:
+    if hasattr(dataset, "column_names") and dataset.column_names:
+        if "text" in dataset.column_names:
+            return "text"
+        return dataset.column_names[0]
+    if hasattr(dataset, "features"):
+        names = list(dataset.features.keys())
+        if "text" in names:
+            return "text"
+        if names:
+            return names[0]
+    return "text"
+
+
+def _normalize_config(config: Optional[str]) -> Optional[str]:
+    if config is None:
+        return None
+    if config.strip().lower() in {"none", "null", "-"}:
+        return None
+    return config
+
+
+def expand_dataset_configs(
+    datasets: List[str], configs: List[str]
+) -> List[Optional[str]]:
+    if not configs:
+        return [None] * len(datasets)
+    if len(configs) == 1 and len(datasets) > 1:
+        return [_normalize_config(configs[0])] * len(datasets)
+    if len(configs) != len(datasets):
+        raise SystemExit(
+            "Provide zero, one, or matching-count --dataset_config values."
+        )
+    return [_normalize_config(cfg) for cfg in configs]
+
+
+def _sample_dataset_rows(
+    dataset, target: int, seed: int
+) -> List[Dict[str, object]]:
+    if target <= 0:
+        return []
+    try:
+        dataset = dataset.shuffle(seed=seed)
+    except Exception:
+        pass
+
+    if hasattr(dataset, "__len__"):
+        limit = min(target, len(dataset))
+        dataset = dataset.select(range(limit))
+        return [row for row in dataset]
+
+    rows = []
+    for row in dataset:
+        rows.append(row)
+        if len(rows) >= target:
+            break
+    return rows
+
+
+def load_texts(args: argparse.Namespace) -> List[str]:
+    texts: List[str] = []
+    if args.text_file:
+        with open(args.text_file, "r", encoding="utf-8") as handle:
+            texts.extend([line.strip() for line in handle if line.strip()])
+    if args.text:
+        texts.extend([t for t in args.text if t])
+
+    if args.dataset:
+        if load_dataset is None:
+            raise SystemExit("datasets is required for --dataset")
+
+        datasets = list(args.dataset)
+        configs = expand_dataset_configs(datasets, list(args.dataset_config))
+        num_datasets = len(datasets)
+        base = args.num_samples // num_datasets
+        remainder = args.num_samples % num_datasets
+
+        for idx, (dataset_name, config) in enumerate(zip(datasets, configs)):
+            target = base + (1 if idx < remainder else 0)
+            dataset = load_dataset(
+                dataset_name,
+                config,
+                split=args.dataset_split,
+                trust_remote_code=True,
+            )
+            rows = _sample_dataset_rows(dataset, target, args.seed + idx)
+            text_field = args.dataset_text_field or guess_text_field(dataset)
+            for row in rows:
+                value = row.get(text_field, None) if isinstance(row, dict) else None
+                if isinstance(value, str) and value.strip():
+                    texts.append(value)
+
+    return texts
+
+
+def load_texts_from_datasets(
+    datasets: List[str],
+    configs: List[Optional[str]],
+    split: str,
+    text_field: Optional[str],
+    num_samples: int,
+    seed: int,
+) -> List[str]:
+    if not datasets:
+        return []
+    if load_dataset is None:
+        raise SystemExit("datasets is required for --dataset")
+
+    texts: List[str] = []
+    num_datasets = len(datasets)
+    base = num_samples // num_datasets
+    remainder = num_samples % num_datasets
+
+    for idx, (dataset_name, config) in enumerate(zip(datasets, configs)):
+        target = base + (1 if idx < remainder else 0)
+        dataset = load_dataset(
+            dataset_name,
+            config,
+            split=split,
+            trust_remote_code=True,
+        )
+        rows = _sample_dataset_rows(dataset, target, seed + idx)
+        field = text_field or guess_text_field(dataset)
+        for row in rows:
+            value = row.get(field, None) if isinstance(row, dict) else None
+            if isinstance(value, str) and value.strip():
+                texts.append(value)
+    return texts
+
+
+def format_alpaca_example(instruction: str, inp: str, output: str) -> str:
+    if inp:
+        return (
+            "### Instruction:\n"
+            f"{instruction}\n\n"
+            "### Input:\n"
+            f"{inp}\n\n"
+            "### Response:\n"
+            f"{output}"
+        )
+    return (
+        "### Instruction:\n"
+        f"{instruction}\n\n"
+        "### Response:\n"
+        f"{output}"
+    )
+
+
+def build_alpaca_messages(
+    instruction: str, inp: str, output: str
+) -> List[Dict[str, str]]:
+    if inp:
+        user_content = f"{instruction}\n\nInput:\n{inp}"
+    else:
+        user_content = instruction
+    return [
+        {"role": "user", "content": user_content},
+        {"role": "assistant", "content": output},
+    ]
+
+
+class FixedSeqDataset(torch.utils.data.Dataset):
+    def __init__(self, records: List[Dict[str, object]], tokenizer, seq_len: int) -> None:
+        self.records = records
+        self.tokenizer = tokenizer
+        self.seq_len = seq_len
+        self.pad_id = tokenizer.pad_token_id
+        if self.pad_id is None:
+            self.pad_id = tokenizer.eos_token_id or 0
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __getitem__(self, idx: int):
+        record = self.records[idx]
+        chat_template = getattr(self.tokenizer, "chat_template", None)
+        if (
+            "messages" in record
+            and hasattr(self.tokenizer, "apply_chat_template")
+            and chat_template
+        ):
+            ids = self.tokenizer.apply_chat_template(
+                record["messages"],
+                tokenize=True,
+                add_generation_prompt=False,
+            )
+        else:
+            text = record.get("text", "")
+            ids = self.tokenizer.encode(text, add_special_tokens=False)
+
+        # Transformers may return a BatchEncoding here instead of a plain list.
+        if hasattr(ids, "input_ids"):
+            ids = ids.input_ids
+        if isinstance(ids, torch.Tensor):
+            ids = ids.tolist()
+        elif not isinstance(ids, list):
+            ids = list(ids)
+
+        if len(ids) > self.seq_len:
+            ids = ids[: self.seq_len]
+        attn = [1] * len(ids)
+        if len(ids) < self.seq_len:
+            pad_len = self.seq_len - len(ids)
+            ids = ids + [self.pad_id] * pad_len
+            attn = attn + [0] * pad_len
+
+        return (
+            torch.tensor(ids, dtype=torch.long),
+            torch.tensor(attn, dtype=torch.long),
+        )
+
+
+def load_instruction_records(
+    args: argparse.Namespace, num_samples: int
+) -> List[Dict[str, object]]:
+    if not args.instruction_dataset:
+        return []
+    if load_dataset is None:
+        raise SystemExit("datasets is required for instruction dataset")
+
+    dataset = load_dataset(
+        args.instruction_dataset,
+        _normalize_config(args.instruction_config),
+        split=args.instruction_split,
+        trust_remote_code=True,
+    )
+    if num_samples > 0:
+        rows = _sample_dataset_rows(dataset, num_samples, args.seed)
+    else:
+        rows = dataset
+    records: List[Dict[str, object]] = []
+    for row in rows:
+        if not isinstance(row, dict):
+            continue
+        instruction = str(row.get(args.instruction_field_instruction, "")).strip()
+        inp = str(row.get(args.instruction_field_input, "")).strip()
+        output = str(row.get(args.instruction_field_output, "")).strip()
+        if not instruction or not output:
+            continue
+        records.append(
+            {
+                "messages": build_alpaca_messages(instruction, inp, output),
+                "text": format_alpaca_example(instruction, inp, output),
+            }
+        )
+    return records
+
+
+def build_token_chunks(
+    texts: List[str], tokenizer, seq_len: int, num_samples: int
+) -> List[torch.Tensor]:
+    chunks: List[torch.Tensor] = []
+    buffer: List[int] = []
+    limit = None if num_samples <= 0 else num_samples
+    for text in texts:
+        ids = tokenizer.encode(text, add_special_tokens=False)
+        if not ids:
+            continue
+        buffer.extend(ids)
+        while len(buffer) >= seq_len and (limit is None or len(chunks) < limit):
+            chunk = buffer[:seq_len]
+            buffer = buffer[seq_len:]
+            chunks.append(torch.tensor(chunk, dtype=torch.long))
+        if limit is not None and len(chunks) >= limit:
+            break
+    return chunks

src/fuse_layers_distill.py

@@ -0,0 +1,2018 @@
+#!/usr/bin/env python3
+"""Distillation helpers for fuse_layers."""
+
+import argparse
+import itertools
+import math
+import os
+from contextlib import contextmanager, nullcontext
+from typing import Dict, List, Optional, Set, Tuple
+
+import torch
+import torch.nn.functional as F
+
+try:
+    import ppl_eval
+except Exception as exc:  # pragma: no cover - optional dependency
+    raise SystemExit("ppl_eval.py is required (missing or invalid)") from exc
+try:
+    from tqdm import tqdm
+except Exception:  # pragma: no cover - optional dependency
+    tqdm = None
+
+try:
+    from torch.func import functional_call as _functional_call
+except Exception:  # pragma: no cover - depends on torch version
+    try:
+        from torch.nn.utils.stateless import functional_call as _functional_call
+    except Exception:  # pragma: no cover - depends on torch version
+        _functional_call = None
+
+from fuse_layers_model import find_attention_module, find_mlp_module
+
+
+def _tqdm_enabled() -> bool:
+    value = os.environ.get("DISABLE_TQDM", os.environ.get("TQDM_DISABLE", "0"))
+    return value.strip().lower() not in {"1", "true", "yes", "on"}
+
+
+@contextmanager
+def temporary_layers(parent: object, name: str, new_layers: torch.nn.Module):
+    original = getattr(parent, name)
+    setattr(parent, name, new_layers)
+    try:
+        yield
+    finally:
+        setattr(parent, name, original)
+
+
+@contextmanager
+def temporary_norm(parent: object):
+    if hasattr(parent, "norm"):
+        original = getattr(parent, "norm")
+        setattr(parent, "norm", torch.nn.Identity())
+        try:
+            yield
+        finally:
+            setattr(parent, "norm", original)
+    else:
+        yield
+
+
+def forward_truncated(
+    parent: torch.nn.Module,
+    layer_attr: str,
+    layers: List[torch.nn.Module],
+    upto: int,
+    input_ids: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    truncated = torch.nn.ModuleList(layers[:upto])
+    with temporary_layers(parent, layer_attr, truncated), temporary_norm(parent):
+        outputs = parent(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            use_cache=False,
+        )
+    if hasattr(outputs, "last_hidden_state"):
+        return outputs.last_hidden_state
+    return outputs[0]
+
+
+def _masked_hidden_mse(diff: torch.Tensor, attention_mask: torch.Tensor) -> Optional[torch.Tensor]:
+    diff_f = diff.float()
+    mask = attention_mask.to(device=diff.device, dtype=torch.float32)
+    denom = mask.sum() * diff_f.size(-1)
+    if denom.item() == 0:
+        return None
+    return (diff_f.pow(2) * mask.unsqueeze(-1)).sum() / denom
+
+
+def _extract_hidden_like(output) -> Optional[torch.Tensor]:
+    if torch.is_tensor(output):
+        return output
+    if isinstance(output, (tuple, list)) and output:
+        first = output[0]
+        if torch.is_tensor(first):
+            return first
+    if hasattr(output, "last_hidden_state"):
+        hidden = getattr(output, "last_hidden_state")
+        if torch.is_tensor(hidden):
+            return hidden
+    return None
+
+
+@contextmanager
+def capture_module_output(module: torch.nn.Module):
+    cache: Dict[str, Optional[torch.Tensor]] = {"output": None}
+
+    def hook(_module, _inputs, output):
+        cache["output"] = _extract_hidden_like(output)
+
+    handle = module.register_forward_hook(hook)
+    try:
+        yield cache
+    finally:
+        handle.remove()
+
+
+_ATTN_NAME_FRAGMENTS = (
+    "self_attn.",
+    "attn.",
+    "attention.",
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "q_norm",
+    "k_norm",
+)
+_MLP_NAME_FRAGMENTS = (
+    "mlp.",
+    "ffn.",
+    "feed_forward",
+    "feedforward",
+    "gate_proj",
+    "up_proj",
+    "down_proj",
+    "fc1",
+    "fc2",
+    "dense_h_to_4h",
+    "dense_4h_to_h",
+    "w1",
+    "w2",
+    "w3",
+)
+
+
+def _classify_param_family(name: str) -> str:
+    lowered = name.lower()
+    if any(fragment in lowered for fragment in _MLP_NAME_FRAGMENTS):
+        return "mlp"
+    if any(fragment in lowered for fragment in _ATTN_NAME_FRAGMENTS):
+        return "attn"
+    return "other"
+
+
+def _family_reg_scale(family: str, attn_scale: float, mlp_scale: float) -> float:
+    if family == "attn":
+        return attn_scale
+    if family == "mlp":
+        return mlp_scale
+    return 1.0
+
+
+def _subset_allows_param(name: str, subset: str) -> bool:
+    if subset == "all":
+        return True
+    return _classify_param_family(name) == subset
+
+
+def _gate_logit_from_prior(prior: torch.Tensor) -> torch.Tensor:
+    # Stable logit: log(p) - log(1 - p).
+    return torch.log(prior) - torch.log1p(-prior)
+
+
+def _build_gate_priors(
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    fisher_a: Dict[str, object],
+    fisher_b: Dict[str, object],
+    num_batches: int,
+    numels_a: Dict[str, int],
+    numels_b: Dict[str, int],
+    fisher_mode: str,
+    eps: float,
+    clamp_eps: float,
+) -> Dict[str, torch.Tensor]:
+    """Return lambda priors for parameters that can be merged."""
+    priors: Dict[str, torch.Tensor] = {}
+    params_b = {name: param for name, param in layer_b.named_parameters()}
+    for name, param_a in layer_a.named_parameters():
+        param_b = params_b.get(name)
+        if param_b is None or param_b.shape != param_a.shape:
+            continue
+        if fisher_mode == "param":
+            fa = fisher_a[name] / max(num_batches, 1)
+            fb = fisher_b[name] / max(num_batches, 1)
+            denom = fa + fb
+            if not isinstance(denom, torch.Tensor):
+                denom = torch.tensor(float(denom))
+            # If Fisher is uninformative, default to symmetric init.
+            prior = torch.where(
+                denom > eps,
+                fa / (denom + eps),
+                torch.full_like(denom, 0.5),
+            )
+            prior = prior.clamp(clamp_eps, 1.0 - clamp_eps)
+            priors[name] = prior
+        else:
+            fa = fisher_a[name] / (max(num_batches, 1) * numels_a[name])
+            fb = fisher_b[name] / (max(num_batches, 1) * numels_b[name])
+            denom = fa + fb
+            if denom <= eps:
+                prior_val = 0.5
+            else:
+                prior_val = float(fa / (denom + eps))
+            prior_val = min(max(prior_val, clamp_eps), 1.0 - clamp_eps)
+            priors[name] = torch.tensor(prior_val, dtype=torch.float32)
+    return priors
+
+
+def compute_fisher_gate_priors(
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    fisher_a: Dict[str, object],
+    fisher_b: Dict[str, object],
+    num_batches: int,
+    numels_a: Dict[str, int],
+    numels_b: Dict[str, int],
+    fisher_mode: str,
+    eps: float,
+    clamp_eps: float = 1e-4,
+) -> Dict[str, torch.Tensor]:
+    """Compute Fisher prior gate lambdas (lambda_prior) for mergeable parameters."""
+    return _build_gate_priors(
+        layer_a=layer_a,
+        layer_b=layer_b,
+        fisher_a=fisher_a,
+        fisher_b=fisher_b,
+        num_batches=num_batches,
+        numels_a=numels_a,
+        numels_b=numels_b,
+        fisher_mode=fisher_mode,
+        eps=eps,
+        clamp_eps=clamp_eps,
+    )
+
+
+class ReparamMergedLayer(torch.nn.Module):
+    """Virtual layer that merges parameters via W0/U reparameterization.
+
+    Parameters of layer_a/layer_b are treated as frozen (detached). We train:
+      - gate logits s (lambda = sigmoid(s))
+      - U (initialized as U0 = (W_a - W_b) / 2)
+
+    Forward uses:
+      W_merge = W0 + (2 * lambda - 1) * U
+      where W0 = (W_a + W_b) / 2
+    """
+
+    def __init__(
+        self,
+        layer_a: torch.nn.Module,
+        layer_b: torch.nn.Module,
+        gate_targets: Dict[str, object],
+        param_subset: str = "all",
+        clamp_eps: float = 1e-4,
+    ) -> None:
+        super().__init__()
+        self.layer_a = layer_a
+        self.layer_b = layer_b
+        self.param_subset = param_subset
+        self._name_map: Dict[str, str] = {}
+
+        self.gates = torch.nn.ParameterDict()
+        self.u = torch.nn.ParameterDict()
+
+        params_b = {name: param for name, param in layer_b.named_parameters()}
+        try:
+            device = next(layer_a.parameters()).device
+        except StopIteration:
+            device = torch.device("cpu")
+
+        for name, param_a in layer_a.named_parameters():
+            param_b = params_b.get(name)
+            if param_b is None or param_b.shape != param_a.shape:
+                continue
+            if not _subset_allows_param(name, self.param_subset):
+                continue
+
+            target = gate_targets.get(name)
+            if target is None:
+                target_t = torch.tensor(0.5, device=device, dtype=torch.float32)
+            elif isinstance(target, torch.Tensor):
+                target_t = target.detach().to(device=device, dtype=torch.float32)
+            else:
+                target_t = torch.tensor(float(target), device=device, dtype=torch.float32)
+
+            target_t = target_t.clamp(clamp_eps, 1.0 - clamp_eps)
+            s0 = _gate_logit_from_prior(target_t)
+            u0 = 0.5 * (param_a.detach().float() - param_b.detach().float())
+
+            safe = name.replace(".", "__")
+            if safe in self.gates:
+                safe = f"{safe}_{len(self.gates)}"
+            self._name_map[name] = safe
+            self.gates[safe] = torch.nn.Parameter(s0)
+            self.u[safe] = torch.nn.Parameter(u0)
+
+    def __getattr__(self, name: str):
+        # Delegate model-specific attributes (e.g. Qwen's `attention_type`) to
+        # the underlying layer so the parent forward doesn't break.
+        try:
+            return super().__getattr__(name)
+        except AttributeError as exc:
+            try:
+                layer_a = super().__getattr__("layer_a")
+                if hasattr(layer_a, name):
+                    return getattr(layer_a, name)
+            except AttributeError:
+                pass
+            try:
+                layer_b = super().__getattr__("layer_b")
+                if hasattr(layer_b, name):
+                    return getattr(layer_b, name)
+            except AttributeError:
+                pass
+            raise exc
+
+    def _safe_for(self, orig: str) -> Optional[str]:
+        return self._name_map.get(orig)
+
+    def gate_lambdas(self) -> Dict[str, torch.Tensor]:
+        out: Dict[str, torch.Tensor] = {}
+        for orig, safe in self._name_map.items():
+            out[orig] = torch.sigmoid(self.gates[safe]).detach()
+        return out
+
+    def _merged_params(self) -> Dict[str, torch.Tensor]:
+        params_a = {name: p for name, p in self.layer_a.named_parameters()}
+        params_b = {name: p for name, p in self.layer_b.named_parameters()}
+        merged_params: Dict[str, torch.Tensor] = {}
+
+        for name, param_a in params_a.items():
+            param_b = params_b.get(name)
+            safe = self._safe_for(name)
+            if safe is None or param_b is None or param_b.shape != param_a.shape:
+                merged_params[name] = param_a.detach()
+                continue
+
+            lam = torch.sigmoid(self.gates[safe]).to(dtype=torch.float32)
+            u = self.u[safe].to(dtype=torch.float32)
+            w0 = 0.5 * (param_a.detach().float() + param_b.detach().float())
+            merged = w0 + (2.0 * lam - 1.0) * u
+            merged_params[name] = merged.to(dtype=param_a.dtype)
+        return merged_params
+
+    def forward(self, *args, **kwargs):
+        if _functional_call is None:
+            raise RuntimeError(
+                "Reparam distillation requires torch.func.functional_call"
+            )
+
+        merged_params = self._merged_params()
+        return _functional_call(self.layer_a, merged_params, args, kwargs)
+
+    def materialize_into_layer_a(self) -> int:
+        merged = 0
+        params_a = {name: p for name, p in self.layer_a.named_parameters()}
+        params_b = {name: p for name, p in self.layer_b.named_parameters()}
+        with torch.no_grad():
+            for orig, safe in self._name_map.items():
+                param_a = params_a.get(orig)
+                param_b = params_b.get(orig)
+                if param_a is None or param_b is None or param_b.shape != param_a.shape:
+                    continue
+                lam = torch.sigmoid(self.gates[safe]).to(device=param_a.device, dtype=torch.float32)
+                u = self.u[safe].to(device=param_a.device, dtype=torch.float32)
+                w0 = 0.5 * (param_a.detach().float() + param_b.detach().float())
+                merged_param = w0 + (2.0 * lam - 1.0) * u
+                param_a.copy_(merged_param.to(dtype=param_a.dtype))
+                merged += 1
+        return merged
+
+
+def distill_reparam_merge(
+    student_model: torch.nn.Module,
+    student_parent: object,
+    student_layer_attr: str,
+    student_layers: List[torch.nn.Module],
+    teacher_model: torch.nn.Module,
+    teacher_parent: object,
+    teacher_layer_attr: str,
+    teacher_layers: List[torch.nn.Module],
+    layer_idx: int,
+    gate_lambdas: Dict[str, object],
+    dataloader,
+    args: argparse.Namespace,
+    progressive_cycle: Optional[int] = None,
+    progressive_total: Optional[int] = None,
+) -> Tuple[int, Dict[str, torch.Tensor], Dict[str, object]]:
+    """Reparameterized distillation that materializes a fused layer into layer_a.
+
+    Trains U and gate logits s (lambda = sigmoid(s)) using:
+      - composition MSE + distill-KL
+      - eta * ||lambda - lambda_gate||^2 + gamma * ||U - U0||^2
+    """
+    total_epochs = float(args.distill_epochs)
+
+    hidden_mse_weight = float(getattr(args, "distill_hidden_mse_weight", 1.0))
+    if hidden_mse_weight < 0.0:
+        raise SystemExit("--distill_hidden_mse_weight must be >= 0")
+    attn_mse_weight = float(getattr(args, "distill_attn_mse_weight", 0.0))
+    if attn_mse_weight < 0.0:
+        raise SystemExit("--distill_attn_mse_weight must be >= 0")
+    mlp_mse_weight = float(getattr(args, "distill_mlp_mse_weight", 0.0))
+    if mlp_mse_weight < 0.0:
+        raise SystemExit("--distill_mlp_mse_weight must be >= 0")
+    param_subset = str(getattr(args, "reparam_param_subset", "all"))
+    if param_subset not in {"all", "mlp", "attn"}:
+        raise SystemExit("--reparam_param_subset must be one of: all, mlp, attn")
+
+    kl_weight = float(args.distill_kl_weight)
+    kl_temp = float(args.distill_kl_temp)
+    if kl_weight < 0.0:
+        raise SystemExit("--distill_kl_weight must be >= 0")
+    if kl_temp <= 0.0:
+        raise SystemExit("--distill_kl_temp must be > 0")
+
+    eta = float(getattr(args, "reparam_eta", 0.0))
+    gamma = float(getattr(args, "reparam_gamma", 0.0))
+    if eta < 0.0:
+        raise SystemExit("--reparam_eta must be >= 0")
+    if gamma < 0.0:
+        raise SystemExit("--reparam_gamma must be >= 0")
+    attn_reg_scale = float(getattr(args, "reparam_attn_reg_scale", 1.0))
+    mlp_reg_scale = float(getattr(args, "reparam_mlp_reg_scale", 1.0))
+    if attn_reg_scale < 0.0:
+        raise SystemExit("--reparam_attn_reg_scale must be >= 0")
+    if mlp_reg_scale < 0.0:
+        raise SystemExit("--reparam_mlp_reg_scale must be >= 0")
+    if (
+        total_epochs > 0.0
+        and hidden_mse_weight == 0.0
+        and attn_mse_weight == 0.0
+        and mlp_mse_weight == 0.0
+        and kl_weight == 0.0
+        and eta == 0.0
+        and gamma == 0.0
+    ):
+        raise SystemExit(
+            "Reparam distillation has no active loss terms. "
+            "Enable hidden/attention/MLP MSE, KL, or at least one reparam regularizer."
+        )
+
+    if not gate_lambdas:
+        raise SystemExit("Reparam distillation requires non-empty gate lambdas.")
+
+    layer_a = student_layers[layer_idx]
+    layer_b = student_layers[layer_idx + 1]
+
+    reparam_layer = ReparamMergedLayer(
+        layer_a,
+        layer_b,
+        gate_lambdas,
+        param_subset=param_subset,
+        clamp_eps=1e-4,
+    )
+    if not reparam_layer._name_map:
+        raise RuntimeError(
+            "No mergeable parameters found for reparam distillation under "
+            f"--reparam_param_subset={param_subset!r}."
+        )
+
+    teacher_attn = None
+    student_attn = None
+    if attn_mse_weight > 0.0:
+        try:
+            teacher_attn = find_attention_module(teacher_layers[layer_idx + 1])
+            student_attn = find_attention_module(reparam_layer.layer_a)
+        except ValueError as exc:
+            raise SystemExit(
+                "Attention-output preservation was requested but an attention module "
+                f"could not be resolved: {exc}"
+            ) from exc
+
+    teacher_mlp = None
+    student_mlp = None
+    if mlp_mse_weight > 0.0:
+        try:
+            teacher_mlp = find_mlp_module(teacher_layers[layer_idx + 1])
+            student_mlp = find_mlp_module(reparam_layer.layer_a)
+        except ValueError as exc:
+            raise SystemExit(
+                "MLP-output preservation was requested but an MLP module could not be "
+                f"resolved: {exc}"
+            ) from exc
+
+    # Virtual layer list: replace layer_a with reparam layer and remove layer_b.
+    virtual_layers = list(student_layers)
+    virtual_layers[layer_idx] = reparam_layer
+    del virtual_layers[layer_idx + 1]
+
+    # Only (U, s) are trainable.
+    for param in student_model.parameters():
+        param.requires_grad_(False)
+    for param in reparam_layer.gates.parameters():
+        param.requires_grad_(True)
+    for param in reparam_layer.u.parameters():
+        param.requires_grad_(True)
+
+    do_train = total_epochs > 0.0
+    if do_train:
+        teacher_model.eval()
+        student_model.train()
+
+    # Rough memory heads-up (esp. when --fisher_mode param makes per-element gates).
+    total_gate_elems = sum(int(p.numel()) for p in reparam_layer.gates.parameters())
+    total_u_elems = sum(int(p.numel()) for p in reparam_layer.u.parameters())
+    gate_mib = total_gate_elems * 4.0 / (1024.0 * 1024.0)
+    u_mib = total_u_elems * 4.0 / (1024.0 * 1024.0)
+    family_counts: Dict[str, int] = {"attn": 0, "mlp": 0, "other": 0}
+    for orig in reparam_layer._name_map:
+        family_counts[_classify_param_family(orig)] += 1
+    print(
+        f"[reparam] subset={param_subset} gates={len(reparam_layer.gates)} "
+        f"(attn={family_counts['attn']}, mlp={family_counts['mlp']}, other={family_counts['other']}) "
+        f"elems={total_gate_elems} (~{gate_mib:.1f} MiB), "
+        f"U_elems={total_u_elems} (~{u_mib:.1f} MiB; +optimizer state)"
+    )
+
+    optimizer = None
+    if do_train:
+        optimizer = torch.optim.AdamW(
+            [*reparam_layer.gates.parameters(), *reparam_layer.u.parameters()],
+            lr=float(args.distill_lr),
+            weight_decay=float(args.distill_weight_decay),
+        )
+
+    device_type = torch.device(args.device).type
+    amp_dtype = None
+    if args.dtype == "float16":
+        amp_dtype = torch.float16
+    elif args.dtype == "bfloat16":
+        amp_dtype = torch.bfloat16
+    use_amp = do_train and amp_dtype is not None and device_type == "cuda"
+    use_scaler = use_amp and amp_dtype == torch.float16
+    scaler = torch.cuda.amp.GradScaler() if use_scaler else None
+
+    full_epochs = int(total_epochs) if do_train else 0
+    fractional = (total_epochs - full_epochs) if do_train else 0.0
+    if fractional < 1e-8:
+        fractional = 0.0
+
+    epoch_plan = [(epoch_idx, None) for epoch_idx in range(full_epochs)]
+    if fractional > 0:
+        try:
+            batches_per_epoch = len(dataloader)
+        except TypeError as exc:
+            raise SystemExit(
+                "Fractional distill epochs require a dataloader with finite length."
+            ) from exc
+        if batches_per_epoch > 0:
+            frac_batches = int(round(fractional * batches_per_epoch))
+            if frac_batches <= 0:
+                frac_batches = 1
+            epoch_plan.append((full_epochs, frac_batches))
+
+    grad_accum = int(getattr(args, "distill_grad_accum_steps", 1))
+    if grad_accum <= 0:
+        raise SystemExit("--distill_grad_accum_steps must be >= 1")
+
+    log_steps = int(getattr(args, "distill_log_steps", 100))
+    max_grad_norm = getattr(args, "distill_max_grad_norm", 1.0)
+
+    params_a = {name: p for name, p in layer_a.named_parameters()}
+    params_b = {name: p for name, p in layer_b.named_parameters()}
+
+    step = 0
+    for epoch_idx, max_batches in epoch_plan:
+        if max_batches is None:
+            epoch_iter = dataloader
+        else:
+            epoch_iter = itertools.islice(dataloader, max_batches)
+        iterator = epoch_iter
+        if tqdm is not None and _tqdm_enabled():
+            if progressive_cycle is not None:
+                if progressive_total is not None:
+                    desc = (
+                        f"Reparam (cycle {progressive_cycle}/{progressive_total}, "
+                        f"epoch {epoch_idx+1})"
+                    )
+                else:
+                    desc = f"Reparam (cycle {progressive_cycle}, epoch {epoch_idx+1})"
+            else:
+                desc = f"Reparam (epoch {epoch_idx+1})"
+            iterator = tqdm(epoch_iter, desc=desc, unit="batch", total=max_batches)
+
+        for batch in iterator:
+            input_ids = batch[0].to(args.device)
+            attention_mask = batch[1].to(args.device)
+            teacher_ids = input_ids.to(args.distill_teacher_device or args.device)
+            teacher_mask = attention_mask.to(args.distill_teacher_device or args.device)
+
+            teacher_depth = layer_idx + 2
+            student_depth = layer_idx + 1
+
+            autocast_ctx = (
+                torch.autocast(device_type=device_type, dtype=amp_dtype)
+                if use_amp
+                else nullcontext()
+            )
+            with autocast_ctx:
+                teacher_attn_ctx = (
+                    capture_module_output(teacher_attn)
+                    if teacher_attn is not None
+                    else nullcontext({"output": None})
+                )
+                teacher_mlp_ctx = (
+                    capture_module_output(teacher_mlp)
+                    if teacher_mlp is not None
+                    else nullcontext({"output": None})
+                )
+                with torch.no_grad():
+                    with teacher_attn_ctx as teacher_attn_cache, teacher_mlp_ctx as teacher_mlp_cache:
+                        teacher_hidden = forward_truncated(
+                            teacher_parent,
+                            teacher_layer_attr,
+                            teacher_layers,
+                            teacher_depth,
+                            teacher_ids,
+                            attention_mask=teacher_mask,
+                        )
+
+                student_attn_ctx = (
+                    capture_module_output(student_attn)
+                    if student_attn is not None
+                    else nullcontext({"output": None})
+                )
+                student_mlp_ctx = (
+                    capture_module_output(student_mlp)
+                    if student_mlp is not None
+                    else nullcontext({"output": None})
+                )
+                with student_attn_ctx as student_attn_cache, student_mlp_ctx as student_mlp_cache:
+                    student_hidden = forward_truncated(
+                        student_parent,
+                        student_layer_attr,
+                        virtual_layers,
+                        student_depth,
+                        input_ids,
+                        attention_mask=attention_mask,
+                    )
+
+                if teacher_hidden.device != student_hidden.device:
+                    teacher_hidden = teacher_hidden.to(student_hidden.device)
+
+                mse_loss = None
+                if hidden_mse_weight > 0.0:
+                    diff = student_hidden - teacher_hidden
+                    mse_loss = _masked_hidden_mse(diff, attention_mask)
+                    if mse_loss is None:
+                        continue
+
+                attn_aux_loss = None
+                if attn_mse_weight > 0.0:
+                    teacher_attn_hidden = teacher_attn_cache.get("output")
+                    student_attn_hidden = student_attn_cache.get("output")
+                    if teacher_attn_hidden is None or student_attn_hidden is None:
+                        raise RuntimeError(
+                            "Attention-output preservation is enabled, but the forward "
+                            "hook did not capture attention outputs."
+                        )
+                    if teacher_attn_hidden.device != student_attn_hidden.device:
+                        teacher_attn_hidden = teacher_attn_hidden.to(student_attn_hidden.device)
+                    attn_aux_loss = _masked_hidden_mse(
+                        student_attn_hidden - teacher_attn_hidden,
+                        attention_mask,
+                    )
+                    if attn_aux_loss is None:
+                        continue
+
+                mlp_aux_loss = None
+                if mlp_mse_weight > 0.0:
+                    teacher_mlp_hidden = teacher_mlp_cache.get("output")
+                    student_mlp_hidden = student_mlp_cache.get("output")
+                    if teacher_mlp_hidden is None or student_mlp_hidden is None:
+                        raise RuntimeError(
+                            "MLP-output preservation is enabled, but the forward hook "
+                            "did not capture MLP outputs."
+                        )
+                    if teacher_mlp_hidden.device != student_mlp_hidden.device:
+                        teacher_mlp_hidden = teacher_mlp_hidden.to(student_mlp_hidden.device)
+                    mlp_aux_loss = _masked_hidden_mse(
+                        student_mlp_hidden - teacher_mlp_hidden,
+                        attention_mask,
+                    )
+                    if mlp_aux_loss is None:
+                        continue
+
+                kl_loss = None
+                if kl_weight > 0.0:
+                    with torch.no_grad():
+                        teacher_outputs = teacher_model(
+                            input_ids=teacher_ids,
+                            attention_mask=teacher_mask,
+                            use_cache=False,
+                        )
+                        teacher_logits = teacher_outputs.logits
+
+                    virtual_container = torch.nn.ModuleList(virtual_layers)
+                    with temporary_layers(
+                        student_parent, student_layer_attr, virtual_container
+                    ):
+                        student_outputs = student_model(
+                            input_ids=input_ids,
+                            attention_mask=attention_mask,
+                            use_cache=False,
+                        )
+                    student_logits = student_outputs.logits
+                    if teacher_logits.device != student_logits.device:
+                        teacher_logits = teacher_logits.to(student_logits.device)
+
+                    shift_teacher_logits = teacher_logits[:, :-1, :].contiguous()
+                    shift_student_logits = student_logits[:, :-1, :].contiguous()
+                    shift_mask = attention_mask[:, 1:].contiguous()
+                    log_p_t = F.log_softmax(shift_teacher_logits / kl_temp, dim=-1)
+                    log_p_s = F.log_softmax(shift_student_logits / kl_temp, dim=-1)
+                    p_t = log_p_t.exp()
+                    kl_flat = (p_t * (log_p_t - log_p_s)).sum(dim=-1)
+                    kl_denom = shift_mask.sum()
+                    if kl_denom.item() == 0:
+                        continue
+                    kl_loss = (
+                        kl_flat * shift_mask.to(kl_flat.dtype)
+                    ).sum() / kl_denom
+
+                lambda_reg = None
+                if eta > 0.0:
+                    reg_sum: Optional[torch.Tensor] = None
+                    reg_elems = 0
+                    for orig, safe in reparam_layer._name_map.items():
+                        lam = torch.sigmoid(reparam_layer.gates[safe]).float()
+                        target = gate_lambdas.get(orig)
+                        if target is None:
+                            target_t = 0.5
+                        elif isinstance(target, torch.Tensor):
+                            target_t = target.to(device=lam.device, dtype=lam.dtype)
+                        else:
+                            target_t = float(target)
+                        diff_lam = lam - target_t
+                        family = _classify_param_family(orig)
+                        scale = _family_reg_scale(
+                            family,
+                            attn_scale=attn_reg_scale,
+                            mlp_scale=mlp_reg_scale,
+                        )
+                        if scale <= 0.0:
+                            continue
+                        part = diff_lam.pow(2).sum() * scale
+                        reg_sum = part if reg_sum is None else reg_sum + part
+                        reg_elems += int(float(diff_lam.numel()) * scale)
+                    if reg_elems > 0 and reg_sum is not None:
+                        lambda_reg = reg_sum / float(reg_elems)
+
+                u_reg = None
+                if gamma > 0.0:
+                    reg_sum: Optional[torch.Tensor] = None
+                    reg_elems = 0
+                    for orig, safe in reparam_layer._name_map.items():
+                        u = reparam_layer.u[safe].float()
+                        param_a = params_a.get(orig)
+                        param_b = params_b.get(orig)
+                        if param_a is None or param_b is None or param_b.shape != param_a.shape:
+                            continue
+                        u0 = 0.5 * (param_a.detach().float() - param_b.detach().float())
+                        diff_u = u - u0
+                        family = _classify_param_family(orig)
+                        scale = _family_reg_scale(
+                            family,
+                            attn_scale=attn_reg_scale,
+                            mlp_scale=mlp_reg_scale,
+                        )
+                        if scale <= 0.0:
+                            continue
+                        part = diff_u.pow(2).sum() * scale
+                        reg_sum = part if reg_sum is None else reg_sum + part
+                        reg_elems += int(float(diff_u.numel()) * scale)
+                    if reg_elems > 0 and reg_sum is not None:
+                        u_reg = reg_sum / float(reg_elems)
+
+            total_loss = None
+            if mse_loss is not None:
+                total_loss = hidden_mse_weight * mse_loss
+            if attn_aux_loss is not None:
+                total_loss = attn_mse_weight * attn_aux_loss if total_loss is None else total_loss + (attn_mse_weight * attn_aux_loss)
+            if mlp_aux_loss is not None:
+                total_loss = mlp_mse_weight * mlp_aux_loss if total_loss is None else total_loss + (mlp_mse_weight * mlp_aux_loss)
+            if kl_loss is not None:
+                total_loss = kl_weight * (kl_temp ** 2) * kl_loss if total_loss is None else total_loss + (kl_weight * (kl_temp ** 2) * kl_loss)
+            if lambda_reg is not None:
+                total_loss = eta * lambda_reg if total_loss is None else total_loss + (eta * lambda_reg)
+            if u_reg is not None:
+                total_loss = gamma * u_reg if total_loss is None else total_loss + (gamma * u_reg)
+            if total_loss is None:
+                continue
+
+            if grad_accum > 1:
+                total_loss = total_loss / grad_accum
+            if use_scaler:
+                scaler.scale(total_loss).backward()
+            else:
+                total_loss.backward()
+
+            if (step + 1) % grad_accum == 0:
+                if max_grad_norm is not None:
+                    if use_scaler:
+                        scaler.unscale_(optimizer)
+                    torch.nn.utils.clip_grad_norm_(
+                        [*reparam_layer.gates.parameters(), *reparam_layer.u.parameters()],
+                        float(max_grad_norm),
+                    )
+                if use_scaler:
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+            if log_steps and (step == 0 or (step + 1) % log_steps == 0):
+                log_parts = [f"loss={total_loss.item():.6e}"]
+                if mse_loss is not None:
+                    log_parts.append(f"mse={mse_loss.item():.6e}")
+                else:
+                    log_parts.append("mse=disabled")
+                if attn_aux_loss is not None:
+                    log_parts.append(f"attn_mse={attn_aux_loss.item():.6e}")
+                elif attn_mse_weight > 0.0:
+                    log_parts.append("attn_mse=nan")
+                if mlp_aux_loss is not None:
+                    log_parts.append(f"mlp_mse={mlp_aux_loss.item():.6e}")
+                elif mlp_mse_weight > 0.0:
+                    log_parts.append("mlp_mse=nan")
+                if kl_loss is not None:
+                    log_parts.append(f"kl={kl_loss.item():.6e}")
+                if lambda_reg is not None:
+                    log_parts.append(f"lam_reg={lambda_reg.item():.6e}")
+                if u_reg is not None:
+                    log_parts.append(f"u_reg={u_reg.item():.6e}")
+                print(
+                    f"[reparam] epoch={epoch_idx+1} step={step+1} " + " ".join(log_parts)
+                )
+            step += 1
+
+    merged = reparam_layer.materialize_into_layer_a()
+    final_lambdas = reparam_layer.gate_lambdas()
+    stats: Dict[str, object] = {
+        "enabled": True,
+        "epochs": total_epochs,
+        "lr": float(args.distill_lr),
+        "hidden_mse_weight": hidden_mse_weight,
+        "attn_mse_weight": attn_mse_weight,
+        "mlp_mse_weight": mlp_mse_weight,
+        "eta": eta,
+        "gamma": gamma,
+        "attn_reg_scale": attn_reg_scale,
+        "mlp_reg_scale": mlp_reg_scale,
+        "param_subset": param_subset,
+        "num_gates": len(final_lambdas),
+        "num_attn_gates": family_counts["attn"],
+        "num_mlp_gates": family_counts["mlp"],
+        "num_other_gates": family_counts["other"],
+    }
+    return merged, final_lambdas, stats
+
+
+class LoRALinear(torch.nn.Module):
+    def __init__(
+        self,
+        base: torch.nn.Linear,
+        rank: int,
+        alpha: float,
+        dropout: float,
+    ) -> None:
+        super().__init__()
+        if rank <= 0:
+            raise ValueError("LoRA rank must be positive")
+        self.base = base
+        self.rank = int(rank)
+        self.alpha = float(alpha)
+        self.scaling = self.alpha / float(self.rank)
+        self.enabled = True
+        if dropout > 0:
+            self.dropout = torch.nn.Dropout(dropout)
+        else:
+            self.dropout = torch.nn.Identity()
+
+        self.lora_A = torch.nn.Linear(base.in_features, self.rank, bias=False)
+        self.lora_B = torch.nn.Linear(self.rank, base.out_features, bias=False)
+        torch.nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.lora_B.weight)
+
+        self.lora_A.to(device=base.weight.device, dtype=base.weight.dtype)
+        self.lora_B.to(device=base.weight.device, dtype=base.weight.dtype)
+        self.merged = False
+
+    def lora_parameters(self) -> List[torch.nn.Parameter]:
+        return [*self.lora_A.parameters(), *self.lora_B.parameters()]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        result = self.base(x)
+        if self.merged or not self.enabled:
+            return result
+        lora_out = self.lora_B(self.lora_A(self.dropout(x)))
+        return result + lora_out * self.scaling
+
+    def merge(self) -> None:
+        if self.merged:
+            return
+        delta = torch.matmul(self.lora_B.weight, self.lora_A.weight)
+        delta = delta.to(dtype=self.base.weight.dtype) * self.scaling
+        self.base.weight.data.add_(delta)
+        self.merged = True
+
+
+def _get_child_module(parent: torch.nn.Module, part: str) -> torch.nn.Module:
+    if isinstance(parent, (torch.nn.ModuleList, torch.nn.Sequential)) and part.isdigit():
+        return parent[int(part)]
+    if isinstance(parent, torch.nn.ModuleDict):
+        return parent[part]
+    return getattr(parent, part)
+
+
+def _set_child_module(parent: torch.nn.Module, part: str, module: torch.nn.Module) -> None:
+    if isinstance(parent, (torch.nn.ModuleList, torch.nn.Sequential)) and part.isdigit():
+        parent[int(part)] = module
+        return
+    if isinstance(parent, torch.nn.ModuleDict):
+        parent[part] = module
+        return
+    setattr(parent, part, module)
+
+
+def _resolve_parent_module(
+    root: torch.nn.Module, module_name: str
+) -> Optional[tuple]:
+    if not module_name:
+        return None
+    parts = module_name.split(".")
+    parent = root
+    for part in parts[:-1]:
+        parent = _get_child_module(parent, part)
+    return parent, parts[-1]
+
+
+def _resolve_module_by_path(root: torch.nn.Module, module_path: str) -> Optional[torch.nn.Module]:
+    if not module_path:
+        return None
+    parts = [part for part in module_path.split(".") if part]
+    node = root
+    for part in parts:
+        try:
+            node = _get_child_module(node, part)
+        except Exception:
+            return None
+    return node
+
+
+def _resolve_layer_container_for_lora(
+    model: torch.nn.Module, layer_path: Optional[str]
+) -> Tuple[Optional[str], Optional[object]]:
+    """Resolve transformer layer container with optional auto-detection.
+
+    Mirrors the candidate path strategy used elsewhere, so LoRA filtering can work
+    even when --layer_path is not provided.
+    """
+    if isinstance(layer_path, str) and layer_path and layer_path.lower() != "none":
+        container = _resolve_module_by_path(model, layer_path)
+        if container is not None:
+            try:
+                list(container)
+                return layer_path, container
+            except TypeError:
+                pass
+
+    candidate_paths = [
+        "model.layers",  # LLaMA, Mistral, Qwen2, Gemma
+        "model.decoder.layers",  # OPT
+        "transformer.h",  # GPT-2, GPT-J, Bloom, Falcon
+        "transformer.blocks",  # MPT
+        "gpt_neox.layers",  # GPT-NeoX
+        "layers",  # fallback
+    ]
+    for path in candidate_paths:
+        container = _resolve_module_by_path(model, path)
+        if container is None:
+            continue
+        try:
+            list(container)
+        except TypeError:
+            continue
+        return path, container
+
+    return None, None
+
+
+def _parse_exclude_pairs_local(raw_values, num_pairs: int) -> Set[int]:
+    if not raw_values or num_pairs <= 0:
+        return set()
+    exclude: Set[int] = set()
+    for item in raw_values:
+        if item is None:
+            continue
+        for part in str(item).split(","):
+            part = part.strip()
+            if not part:
+                continue
+            try:
+                idx = int(part)
+            except ValueError as exc:
+                raise SystemExit("--exclude_pairs must contain integers.") from exc
+            if idx < 0:
+                idx = num_pairs + idx
+            if 0 <= idx < num_pairs:
+                exclude.add(idx)
+    return exclude
+
+
+def _extract_layer_index_from_module_name(
+    module_name: str, layer_path: str
+) -> Optional[int]:
+    if not layer_path:
+        return None
+    prefix = f"{layer_path}."
+    if not module_name.startswith(prefix):
+        return None
+    rest = module_name[len(prefix) :]
+    if not rest:
+        return None
+    idx_text = rest.split(".", 1)[0]
+    if not idx_text.isdigit():
+        return None
+    return int(idx_text)
+
+
+def _select_linear_modules_for_lora_targets(
+    model: torch.nn.Module,
+    args: argparse.Namespace,
+    *,
+    log_tag: str,
+) -> Tuple[List[Tuple[str, torch.nn.Linear]], Optional[Set[str]], Set[int], Optional[str]]:
+    raw_targets = getattr(args, "lora_target_modules", None)
+    target_modules: Optional[Set[str]] = None
+    if raw_targets:
+        target_modules = {str(item) for item in raw_targets if str(item)}
+
+    exclude_layer_indices: Set[int] = set()
+    resolved_layer_path: Optional[str] = None
+    if bool(getattr(args, "lora_respect_exclude_pairs", False)):
+        requested_layer_path = getattr(args, "layer_path", None)
+        resolved_layer_path, layer_container = _resolve_layer_container_for_lora(
+            model, requested_layer_path
+        )
+        if isinstance(layer_container, (torch.nn.ModuleList, list, tuple)):
+            num_pairs = max(len(layer_container) - 1, 0)
+            exclude_pairs = _parse_exclude_pairs_local(
+                getattr(args, "exclude_pairs", None), num_pairs
+            )
+            for pair_idx in exclude_pairs:
+                exclude_layer_indices.add(pair_idx)
+                exclude_layer_indices.add(pair_idx + 1)
+        else:
+            print(
+                f"[{log_tag}] Warning: --lora_respect_exclude_pairs enabled, but "
+                f"could not resolve layer path '{requested_layer_path}'."
+            )
+
+    linear_modules = [
+        (name, module)
+        for name, module in model.named_modules()
+        if isinstance(module, torch.nn.Linear)
+        and (target_modules is None or name.split(".")[-1] in target_modules)
+        and (
+            not exclude_layer_indices
+            or _extract_layer_index_from_module_name(name, resolved_layer_path or "")
+            not in exclude_layer_indices
+        )
+    ]
+    return linear_modules, target_modules, exclude_layer_indices, resolved_layer_path
+
+
+def apply_lora_adapters(
+    model: torch.nn.Module, args: argparse.Namespace
+) -> List[LoRALinear]:
+    if args.lora_rank <= 0:
+        raise SystemExit("--lora_rank must be > 0 when --lora_epochs > 0")
+    linear_modules, target_modules, exclude_layer_indices, _ = (
+        _select_linear_modules_for_lora_targets(model, args, log_tag="lora")
+    )
+    if not linear_modules:
+        raise SystemExit(
+            "No Linear modules found for LoRA adapters "
+            "(check --lora_target_modules / --exclude_pairs / --lora_respect_exclude_pairs)."
+        )
+
+    lora_modules: List[LoRALinear] = []
+    for name, module in linear_modules:
+        resolved = _resolve_parent_module(model, name)
+        if resolved is None:
+            continue
+        parent, attr = resolved
+        wrapped = LoRALinear(
+            base=module,
+            rank=args.lora_rank,
+            alpha=args.lora_alpha,
+            dropout=args.lora_dropout,
+        )
+        _set_child_module(parent, attr, wrapped)
+        lora_modules.append(wrapped)
+
+    for param in model.parameters():
+        param.requires_grad_(False)
+    for lora_module in lora_modules:
+        for param in lora_module.lora_parameters():
+            param.requires_grad_(True)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    percent = 100.0 * trainable_params / max(total_params, 1)
+    target_note = ""
+    if target_modules is not None:
+        target_note = f" target={sorted(target_modules)}"
+    exclude_note = ""
+    if exclude_layer_indices:
+        exclude_note = f" excluded_layers={sorted(exclude_layer_indices)}"
+    print(
+        "[lora] Applied adapters to "
+        f"{len(lora_modules)} linear modules "
+        f"({trainable_params}/{total_params} trainable, {percent:.4f}%)."
+        f"{target_note}{exclude_note}"
+    )
+    return lora_modules
+
+
+def merge_lora_adapters(model: torch.nn.Module) -> None:
+    lora_entries = [
+        (name, module)
+        for name, module in model.named_modules()
+        if isinstance(module, LoRALinear)
+    ]
+    for name, module in lora_entries:
+        module.merge()
+        resolved = _resolve_parent_module(model, name)
+        if resolved is None:
+            continue
+        parent, attr = resolved
+        _set_child_module(parent, attr, module.base)
+
+
+def set_lora_enabled(lora_modules: List[LoRALinear], enabled: bool) -> None:
+    for module in lora_modules:
+        module.enabled = enabled
+
+
+def lora_ce_finetune(
+    model: torch.nn.Module,
+    dataloader,
+    eval_tokenizer,
+    eval_datasets: List[str],
+    eval_configs: List[Optional[str]],
+    eval_history: List[Dict[str, object]],
+    args: argparse.Namespace,
+    eval_dataloaders: Optional[Dict[str, object]] = None,
+    progressive_cycle: Optional[int] = None,
+    progressive_total: Optional[int] = None,
+) -> None:
+    total_epochs = float(args.lora_epochs)
+    if total_epochs <= 0:
+        return
+
+    use_kl = bool(getattr(args, "lora_kl_enabled", False))
+    kl_weight = float(getattr(args, "lora_kl_weight", 0.0))
+    kl_temp = float(getattr(args, "lora_kl_temp", 1.0))
+    if use_kl:
+        if kl_weight < 0.0:
+            raise SystemExit("--lora_kl_weight must be >= 0")
+        if kl_temp <= 0.0:
+            raise SystemExit("--lora_kl_temp must be > 0")
+        if kl_weight == 0.0:
+            use_kl = False
+
+    lora_modules = apply_lora_adapters(model, args)
+    if not lora_modules:
+        return
+
+    model.train()
+
+    lora_params = []
+    for module in lora_modules:
+        lora_params.extend(module.lora_parameters())
+
+    optimizer = torch.optim.AdamW(
+        lora_params,
+        lr=args.lora_lr,
+        weight_decay=args.lora_weight_decay,
+    )
+
+    device_type = torch.device(args.device).type
+    amp_dtype = None
+    if args.dtype == "float16":
+        amp_dtype = torch.float16
+    elif args.dtype == "bfloat16":
+        amp_dtype = torch.bfloat16
+    use_amp = amp_dtype is not None and device_type == "cuda"
+    use_scaler = use_amp and amp_dtype == torch.float16
+    scaler = torch.cuda.amp.GradScaler() if use_scaler else None
+
+    full_epochs = int(total_epochs)
+    fractional = total_epochs - full_epochs
+    if fractional < 1e-8:
+        fractional = 0.0
+
+    epoch_plan = [(epoch_idx, None) for epoch_idx in range(full_epochs)]
+    if fractional > 0:
+        try:
+            batches_per_epoch = len(dataloader)
+        except TypeError as exc:
+            raise SystemExit(
+                "Fractional lora epochs require a dataloader with finite length."
+            ) from exc
+        if batches_per_epoch > 0:
+            frac_batches = int(round(fractional * batches_per_epoch))
+            if frac_batches <= 0:
+                frac_batches = 1
+            epoch_plan.append((full_epochs, frac_batches))
+
+    step = 0
+    for epoch_idx, max_batches in epoch_plan:
+        if max_batches is None:
+            epoch_iter = dataloader
+        else:
+            epoch_iter = itertools.islice(dataloader, max_batches)
+        iterator = epoch_iter
+        if tqdm is not None and _tqdm_enabled():
+            if progressive_cycle is not None:
+                if progressive_total is not None:
+                    desc = (
+                        f"LoRA (cycle {progressive_cycle}/{progressive_total}, "
+                        f"epoch {epoch_idx+1})"
+                    )
+                else:
+                    desc = f"LoRA (cycle {progressive_cycle}, epoch {epoch_idx+1})"
+            else:
+                desc = f"LoRA (epoch {epoch_idx+1})"
+            iterator = tqdm(
+                epoch_iter,
+                desc=desc,
+                unit="batch",
+                total=max_batches,
+            )
+        for batch in iterator:
+            input_ids = batch[0].to(args.device)
+            attention_mask = batch[1].to(args.device)
+            autocast_ctx = (
+                torch.autocast(device_type=device_type, dtype=amp_dtype)
+                if use_amp
+                else nullcontext()
+            )
+            with autocast_ctx:
+                outputs = model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    use_cache=False,
+                )
+                logits = outputs.logits
+                shift_logits = logits[:, :-1, :].contiguous()
+                shift_labels = input_ids[:, 1:].contiguous()
+                shift_mask = attention_mask[:, 1:].contiguous()
+                ce_flat = F.cross_entropy(
+                    shift_logits.view(-1, shift_logits.size(-1)),
+                    shift_labels.view(-1),
+                    reduction="none",
+                )
+                ce_denom = shift_mask.sum()
+                if ce_denom.item() == 0:
+                    continue
+                ce_loss = (
+                    ce_flat * shift_mask.view(-1).to(ce_flat.dtype)
+                ).sum() / ce_denom
+                kl_loss = None
+                if use_kl:
+                    set_lora_enabled(lora_modules, False)
+                    with torch.no_grad():
+                        base_outputs = model(
+                            input_ids=input_ids,
+                            attention_mask=attention_mask,
+                            use_cache=False,
+                        )
+                        base_logits = base_outputs.logits
+                    set_lora_enabled(lora_modules, True)
+                    if base_logits.device != shift_logits.device:
+                        base_logits = base_logits.to(shift_logits.device)
+                    shift_base_logits = base_logits[:, :-1, :].contiguous()
+                    log_p_pre = F.log_softmax(shift_base_logits / kl_temp, dim=-1)
+                    log_p_post = F.log_softmax(shift_logits / kl_temp, dim=-1)
+                    p_pre = log_p_pre.exp()
+                    kl_flat = (p_pre * (log_p_pre - log_p_post)).sum(dim=-1)
+                    kl_loss = (
+                        kl_flat * shift_mask.to(kl_flat.dtype)
+                    ).sum() / ce_denom
+
+            total_loss = ce_loss
+            if kl_loss is not None:
+                total_loss = total_loss + (kl_weight * (kl_temp ** 2) * kl_loss)
+
+            if args.lora_grad_accum_steps > 1:
+                total_loss = total_loss / args.lora_grad_accum_steps
+            if use_scaler:
+                scaler.scale(total_loss).backward()
+            else:
+                total_loss.backward()
+
+            if (step + 1) % args.lora_grad_accum_steps == 0:
+                if args.lora_max_grad_norm is not None:
+                    if use_scaler:
+                        scaler.unscale_(optimizer)
+                    torch.nn.utils.clip_grad_norm_(
+                        lora_params,
+                        args.lora_max_grad_norm,
+                    )
+                if use_scaler:
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+                if args.lora_eval_every and (step + 1) % args.lora_eval_every == 0:
+                    prev_mode = model.training
+                    model.eval()
+                    eval_device = args.eval_device or args.device
+                    if eval_dataloaders is not None:
+                        results = ppl_eval.evaluate_ppl_dataloaders(
+                            model,
+                            eval_dataloaders,
+                            eval_device,
+                            max_batches=args.lora_eval_max_batches,
+                        )
+                    else:
+                        results = ppl_eval.evaluate_ppl_datasets(
+                            model,
+                            eval_tokenizer,
+                            datasets=eval_datasets,
+                            configs=eval_configs,
+                            split=args.eval_split,
+                            text_field=args.eval_text_field,
+                            num_samples=args.eval_num_samples,
+                            seq_len=args.eval_seq_len,
+                            batch_size=args.eval_batch_size or args.batch_size,
+                            device=eval_device,
+                            seed=args.seed,
+                            shuffle=False,
+                            model_family=args.eval_model_family,
+                            add_bos=args.eval_add_bos,
+                            max_batches=args.lora_eval_max_batches,
+                            cache_dir=args.eval_cache_dir,
+                            num_workers=args.eval_num_workers,
+                        )
+                    eval_history.append({"step": step + 1, "ppl": results})
+                    print(f"[lora] eval step={step+1}: {results}")
+                    if prev_mode:
+                        model.train()
+
+            if args.lora_log_steps and (
+                step == 0 or (step + 1) % args.lora_log_steps == 0
+            ):
+                log_parts = [f"loss={total_loss.item():.6f}"]
+                if kl_loss is not None:
+                    log_parts.append(f"kl={kl_loss.item():.6f}")
+                print(
+                    f"[lora] epoch={epoch_idx+1} step={step+1} "
+                    + " ".join(log_parts)
+                )
+            step += 1
+
+    merge_lora_adapters(model)
+
+
+def _masked_kl(
+    logits_p: torch.Tensor,
+    logits_q: torch.Tensor,
+    attention_mask: torch.Tensor,
+    temp: float,
+    detach_p: bool = True,
+) -> Optional[torch.Tensor]:
+    shift_mask = attention_mask[:, 1:].contiguous()
+    denom = shift_mask.sum()
+    if denom.item() == 0:
+        return None
+
+    p = logits_p[:, :-1, :].contiguous()
+    q = logits_q[:, :-1, :].contiguous()
+    if p.device != q.device:
+        p = p.to(q.device)
+
+    # Keep dtype to avoid blowing up memory on large vocab models.
+    log_p = F.log_softmax(p / temp, dim=-1)
+    log_q = F.log_softmax(q / temp, dim=-1)
+    if detach_p:
+        log_p = log_p.detach()
+    p_probs = log_p.exp()
+    kl_flat = (p_probs * (log_p - log_q)).sum(dim=-1)
+    return (kl_flat * shift_mask.to(kl_flat.dtype)).sum() / denom
+
+
+def _extract_hidden_tensor(output: object) -> Optional[torch.Tensor]:
+    if isinstance(output, torch.Tensor):
+        return output
+    if isinstance(output, (tuple, list)) and output:
+        first = output[0]
+        if isinstance(first, torch.Tensor):
+            return first
+    return None
+
+
+def _grad_l2_norm(grads: List[Optional[torch.Tensor]]) -> float:
+    total = 0.0
+    for grad in grads:
+        if grad is None:
+            continue
+        total += float(grad.detach().float().pow(2).sum().item())
+    if total <= 0.0:
+        return 0.0
+    return float(math.sqrt(total))
+
+
+def _register_forward_pre_hook_with_optional_kwargs(layer, hook):
+    try:
+        handle = layer.register_forward_pre_hook(hook, with_kwargs=True)
+        return handle
+    except TypeError:
+        def wrapper(module, inputs):
+            return hook(module, inputs, None)
+
+        return layer.register_forward_pre_hook(wrapper)
+
+
+def commutator_precondition(
+    student_model: torch.nn.Module,
+    student_layers: List[torch.nn.Module],
+    teacher_model: torch.nn.Module,
+    dataloader,
+    dwce_scores: Optional[List[float]],
+    args: argparse.Namespace,
+    exclude_pairs: Optional[Set[int]] = None,
+    progressive_cycle: Optional[int] = None,
+    progressive_total: Optional[int] = None,
+) -> Dict[str, object]:
+    """Run commutator-style preconditioning before pair fusion.
+
+    Objective on each sampled pair i:
+      L = T^2 * KL(p_teacher || p_student) + mu * L_interaction(i)
+
+    Interaction loss is computed locally on block (i+1):
+      r1 = B_{i+1}(h_{i+1}) - h_{i+1}
+      r0 = B_{i+1}(h_i) - h_i
+      L_interaction = ||r1-r0||^2  (or relative form).
+    """
+    if not bool(getattr(args, "comm_enabled", False)):
+        return {"enabled": False}
+    if not student_layers or len(student_layers) < 2:
+        return {"enabled": False, "reason": "need_at_least_2_layers"}
+
+    temp = float(getattr(args, "comm_temp", 2.0))
+    steps_ratio = float(getattr(args, "comm_steps_ratio", 0.1))
+    lr_scale = float(getattr(args, "comm_lr_scale", 0.1))
+    sample_eta = float(getattr(args, "comm_sample_eta", 0.5))
+    sample_dwce_scale = float(getattr(args, "comm_sample_dwce_scale", 1.0))
+    top_k = int(getattr(args, "comm_topk", 1))
+    interaction_mode = str(getattr(args, "comm_interaction_mode", "relative")).strip().lower()
+    interaction_eps = float(getattr(args, "comm_interaction_eps", 1e-8))
+    mu_cfg = getattr(args, "comm_mu", None)
+    mu_auto = bool(getattr(args, "comm_mu_auto", False))
+    mu_auto_rho = float(getattr(args, "comm_mu_auto_rho", 0.1))
+    mu_auto_eps = float(getattr(args, "comm_mu_auto_eps", 1e-8))
+    comm_train_mode = str(getattr(args, "comm_train_mode", "lora")).strip().lower()
+    log_steps = int(getattr(args, "comm_log_steps", 50))
+
+    if temp <= 0.0:
+        raise SystemExit("--comm_temp must be > 0")
+    if steps_ratio < 0.0:
+        raise SystemExit("--comm_steps_ratio must be >= 0")
+    if lr_scale <= 0.0:
+        raise SystemExit("--comm_lr_scale must be > 0")
+    if not (0.0 <= sample_eta <= 1.0):
+        raise SystemExit("--comm_sample_eta must be in [0, 1]")
+    if top_k <= 0:
+        raise SystemExit("--comm_topk must be >= 1")
+    if interaction_mode not in {"mse", "relative"}:
+        raise SystemExit("--comm_interaction_mode must be one of: mse, relative")
+    if comm_train_mode not in {"lora", "full"}:
+        raise SystemExit("--comm_train_mode must be one of: lora, full")
+    if interaction_eps <= 0.0:
+        raise SystemExit("--comm_interaction_eps must be > 0")
+    if mu_auto_rho < 0.0:
+        raise SystemExit("--comm_mu_auto_rho must be >= 0")
+    if mu_auto_eps <= 0.0:
+        raise SystemExit("--comm_mu_auto_eps must be > 0")
+
+    if mu_cfg is None:
+        base_mu = 0.5 if interaction_mode == "relative" else 0.1
+    else:
+        base_mu = float(mu_cfg)
+    if base_mu < 0.0:
+        raise SystemExit("--comm_mu must be >= 0")
+
+    distill_epochs = float(getattr(args, "distill_epochs", 1.0))
+    if distill_epochs <= 0.0:
+        distill_epochs = 1.0
+    grad_accum = int(getattr(args, "distill_grad_accum_steps", 1))
+    if grad_accum <= 0:
+        grad_accum = 1
+
+    try:
+        batches_per_epoch = len(dataloader)
+    except TypeError as exc:
+        raise SystemExit(
+            "Commutator preconditioning requires a finite-length distillation dataloader."
+        ) from exc
+    if batches_per_epoch <= 0:
+        return {"enabled": False, "reason": "empty_dataloader"}
+
+    full_epochs = int(distill_epochs)
+    fractional = distill_epochs - full_epochs
+    if fractional < 1e-8:
+        fractional = 0.0
+    total_batches = full_epochs * batches_per_epoch
+    if fractional > 0.0:
+        frac_batches = int(round(fractional * batches_per_epoch))
+        if frac_batches <= 0:
+            frac_batches = 1
+        total_batches += frac_batches
+
+    distill_opt_steps = int(math.ceil(total_batches / float(grad_accum)))
+    target_opt_steps = int(round(steps_ratio * distill_opt_steps))
+    if target_opt_steps <= 0:
+        target_opt_steps = 1
+
+    num_pairs = max(len(student_layers) - 1, 0)
+    exclude_set = {
+        int(idx)
+        for idx in (exclude_pairs or set())
+        if isinstance(idx, int) and 0 <= int(idx) < num_pairs
+    }
+    allowed_pairs = [i for i in range(num_pairs) if i not in exclude_set]
+    if not allowed_pairs:
+        return {"enabled": False, "reason": "all_pairs_excluded"}
+
+    ranked_pairs = list(allowed_pairs)
+    if dwce_scores is not None and len(dwce_scores) >= num_pairs:
+        finite_pairs = []
+        for idx in allowed_pairs:
+            value = float(dwce_scores[idx])
+            if math.isfinite(value):
+                finite_pairs.append(idx)
+        if finite_pairs:
+            ranked_pairs = sorted(finite_pairs, key=lambda i: float(dwce_scores[i]))
+        else:
+            ranked_pairs = list(allowed_pairs)
+    candidate_pairs = ranked_pairs[: min(top_k, len(ranked_pairs))]
+    if not candidate_pairs:
+        return {"enabled": False, "reason": "no_candidate_pairs"}
+
+    layer_trainable_params: List[List[torch.nn.Parameter]] = []
+    trainable_params: List[torch.nn.Parameter] = []
+    if comm_train_mode == "lora":
+        # LoRA comm preconditioning: update LoRA adapters on receiver layer (i+1).
+        lora_modules = apply_lora_adapters(student_model, args)
+        if not lora_modules:
+            return {"enabled": False, "reason": "no_lora_modules"}
+
+        trainable_seen: Set[int] = set()
+        for module in lora_modules:
+            for param in module.lora_parameters():
+                pid = id(param)
+                if pid in trainable_seen:
+                    continue
+                trainable_seen.add(pid)
+                trainable_params.append(param)
+
+        for layer in student_layers:
+            seen: Set[int] = set()
+            params: List[torch.nn.Parameter] = []
+            for module in layer.modules():
+                if not isinstance(module, LoRALinear):
+                    continue
+                for param in module.lora_parameters():
+                    pid = id(param)
+                    if pid in seen:
+                        continue
+                    seen.add(pid)
+                    params.append(param)
+            layer_trainable_params.append(params)
+    else:
+        # Full-weight comm preconditioning: update full receiver-layer weights.
+        for layer in student_layers:
+            seen: Set[int] = set()
+            params: List[torch.nn.Parameter] = []
+            for param in layer.parameters():
+                if not isinstance(param, torch.nn.Parameter):
+                    continue
+                pid = id(param)
+                if pid in seen:
+                    continue
+                seen.add(pid)
+                params.append(param)
+            layer_trainable_params.append(params)
+
+    candidate_pairs = [
+        i
+        for i in candidate_pairs
+        if (i + 1) < len(layer_trainable_params) and layer_trainable_params[i + 1]
+    ]
+    if not candidate_pairs:
+        if comm_train_mode == "lora":
+            merge_lora_adapters(student_model)
+        return {"enabled": False, "reason": "no_trainable_receiver_layers"}
+
+    if comm_train_mode == "full":
+        trainable_seen: Set[int] = set()
+        for pair_idx in candidate_pairs:
+            for param in layer_trainable_params[pair_idx + 1]:
+                pid = id(param)
+                if pid in trainable_seen:
+                    continue
+                trainable_seen.add(pid)
+                trainable_params.append(param)
+        if not trainable_params:
+            return {"enabled": False, "reason": "no_trainable_receiver_layers"}
+
+        # Freeze non-comm params to reduce grad memory.
+        for param in student_model.parameters():
+            param.requires_grad_(False)
+        for param in trainable_params:
+            param.requires_grad_(True)
+
+    if not trainable_params:
+        if comm_train_mode == "lora":
+            merge_lora_adapters(student_model)
+        return {"enabled": False, "reason": "no_trainable_params"}
+
+    candidate_probs = torch.full(
+        (len(candidate_pairs),),
+        1.0 / float(len(candidate_pairs)),
+        dtype=torch.float32,
+    )
+    if dwce_scores is not None and len(dwce_scores) >= num_pairs and sample_eta > 0.0:
+        score_vec = torch.tensor(
+            [float(dwce_scores[i]) for i in candidate_pairs], dtype=torch.float32
+        )
+        score_vec = torch.nan_to_num(score_vec, nan=1e9, posinf=1e9, neginf=-1e9)
+        biased = torch.softmax(-float(sample_dwce_scale) * score_vec, dim=0)
+        candidate_probs = (1.0 - sample_eta) * candidate_probs + sample_eta * biased
+        candidate_probs = candidate_probs / candidate_probs.sum()
+
+    probs_by_pair = [0.0 for _ in range(num_pairs)]
+    for pos, pair_idx in enumerate(candidate_pairs):
+        probs_by_pair[pair_idx] = float(candidate_probs[pos].item())
+
+    lr = float(getattr(args, "distill_lr", 1e-4)) * lr_scale
+    optimizer = torch.optim.AdamW(
+        trainable_params,
+        lr=lr,
+        weight_decay=float(getattr(args, "distill_weight_decay", 0.0)),
+    )
+
+    device_type = torch.device(args.device).type
+    amp_dtype = None
+    if args.dtype == "float16":
+        amp_dtype = torch.float16
+    elif args.dtype == "bfloat16":
+        amp_dtype = torch.bfloat16
+    use_amp = amp_dtype is not None and device_type == "cuda"
+    use_scaler = use_amp and amp_dtype == torch.float16
+    scaler = torch.cuda.amp.GradScaler() if use_scaler else None
+
+    teacher_device = next(teacher_model.parameters()).device
+    teacher_model.eval()
+    student_model.train()
+
+    gen = torch.Generator(device="cpu")
+    seed = int(getattr(args, "seed", 0))
+    if progressive_cycle is not None:
+        seed += int(progressive_cycle) * 100003
+    gen.manual_seed(seed)
+
+    opt_step = 0
+    total_loss_sum = 0.0
+    anchor_sum = 0.0
+    interaction_sum = 0.0
+    mu_sum = 0.0
+    counted = 0
+    pair_counts = [0 for _ in range(num_pairs)]
+
+    desc = "Comm"
+    if progressive_cycle is not None:
+        if progressive_total is not None:
+            desc = f"Comm (cycle {progressive_cycle}/{progressive_total})"
+        else:
+            desc = f"Comm (cycle {progressive_cycle})"
+    iterator = range(target_opt_steps)
+    if tqdm is not None and _tqdm_enabled():
+        iterator = tqdm(iterator, desc=desc, unit="step")
+
+    data_iter = iter(dataloader)
+    autocast_ctx = (
+        torch.autocast(device_type=device_type, dtype=amp_dtype)
+        if use_amp
+        else nullcontext()
+    )
+
+    for _ in iterator:
+        optimizer.zero_grad(set_to_none=True)
+        accum_done = 0
+        while accum_done < grad_accum:
+            try:
+                batch = next(data_iter)
+            except StopIteration:
+                data_iter = iter(dataloader)
+                batch = next(data_iter)
+
+            input_ids = batch[0].to(args.device)
+            attention_mask = batch[1].to(args.device)
+            sampled_pos = int(torch.multinomial(candidate_probs, 1, generator=gen).item())
+            pair_idx = int(candidate_pairs[sampled_pos])
+            pair_counts[pair_idx] += 1
+
+            receiver_params = layer_trainable_params[pair_idx + 1]
+            receiver_param_ids = {id(param) for param in receiver_params}
+
+            teacher_ids = input_ids.to(teacher_device)
+            teacher_mask = attention_mask.to(teacher_device)
+            with torch.no_grad(), autocast_ctx:
+                teacher_outputs = teacher_model(
+                    input_ids=teacher_ids,
+                    attention_mask=teacher_mask,
+                    use_cache=False,
+                )
+                teacher_logits = teacher_outputs.logits
+
+            capture: Dict[str, object] = {
+                "h_l": None,
+                "h_lp1": None,
+                "y1": None,
+                "recv_args": None,
+                "recv_kwargs": None,
+            }
+
+            def _hook_l(_module, inputs, _output):
+                if inputs and isinstance(inputs[0], torch.Tensor):
+                    capture["h_l"] = inputs[0]
+
+            def _hook_recv_pre(_module, inputs, kwargs):
+                capture["recv_args"] = inputs
+                capture["recv_kwargs"] = kwargs
+
+            def _hook_recv(_module, inputs, output):
+                if inputs and isinstance(inputs[0], torch.Tensor):
+                    capture["h_lp1"] = inputs[0]
+                capture["y1"] = _extract_hidden_tensor(output)
+
+            handles: List[object] = [
+                student_layers[pair_idx].register_forward_hook(_hook_l),
+                _register_forward_pre_hook_with_optional_kwargs(
+                    student_layers[pair_idx + 1], _hook_recv_pre
+                ),
+                student_layers[pair_idx + 1].register_forward_hook(_hook_recv),
+            ]
+            try:
+                with autocast_ctx:
+                    student_outputs = student_model(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        use_cache=False,
+                    )
+                    student_logits = student_outputs.logits
+            finally:
+                for handle in handles:
+                    try:
+                        handle.remove()
+                    except Exception:
+                        pass
+
+            with autocast_ctx:
+                anchor_kl = _masked_kl(
+                    teacher_logits,
+                    student_logits,
+                    attention_mask,
+                    temp=temp,
+                    detach_p=True,
+                )
+                if anchor_kl is None:
+                    continue
+                anchor_loss = (temp ** 2) * anchor_kl
+
+                interaction_loss = None
+                h_l = capture.get("h_l")
+                h_lp1 = capture.get("h_lp1")
+                y1 = capture.get("y1")
+                recv_args = capture.get("recv_args")
+                recv_kwargs = capture.get("recv_kwargs")
+                if (
+                    isinstance(h_l, torch.Tensor)
+                    and isinstance(h_lp1, torch.Tensor)
+                    and isinstance(y1, torch.Tensor)
+                    and isinstance(recv_args, tuple)
+                    and len(recv_args) > 0
+                    and isinstance(recv_args[0], torch.Tensor)
+                ):
+                    call_args = list(recv_args)
+                    first_hidden = call_args[0]
+                    h_l_detached = h_l.detach().to(
+                        device=first_hidden.device,
+                        dtype=first_hidden.dtype,
+                    )
+                    call_args[0] = h_l_detached
+                    call_kwargs = dict(recv_kwargs) if isinstance(recv_kwargs, dict) else {}
+
+                    y0_raw = student_layers[pair_idx + 1](*tuple(call_args), **call_kwargs)
+                    y0 = _extract_hidden_tensor(y0_raw)
+                    if isinstance(y0, torch.Tensor):
+                        if y0.device != y1.device:
+                            y0 = y0.to(y1.device)
+                        h_lp1_detached = h_lp1.detach().to(device=y1.device, dtype=y1.dtype)
+                        h_l_for_res = h_l.detach().to(device=y0.device, dtype=y0.dtype)
+                        r1 = y1 - h_lp1_detached
+                        r0 = y0 - h_l_for_res
+                        mask = attention_mask.to(dtype=r1.dtype)
+                        mask_sum = mask.sum()
+                        if mask_sum.item() > 0:
+                            if interaction_mode == "relative":
+                                num = (r1 - r0).float().pow(2).sum(dim=-1)
+                                den = r1.float().pow(2).sum(dim=-1) + float(interaction_eps)
+                                ratio = (num / den) * mask.to(num.dtype)
+                                interaction_loss = ratio.sum() / (mask_sum + 1e-8)
+                            else:
+                                denom = mask_sum * r1.size(-1)
+                                if denom.item() > 0:
+                                    interaction_loss = (
+                                        (r1 - r0).pow(2) * mask.unsqueeze(-1)
+                                    ).sum() / denom
+
+                mu_effective = float(base_mu)
+                if (
+                    mu_auto
+                    and interaction_loss is not None
+                    and receiver_params
+                    and mu_auto_rho > 0.0
+                ):
+                    anchor_grads = torch.autograd.grad(
+                        anchor_loss,
+                        receiver_params,
+                        retain_graph=True,
+                        allow_unused=True,
+                    )
+                    interaction_grads = torch.autograd.grad(
+                        interaction_loss,
+                        receiver_params,
+                        retain_graph=True,
+                        allow_unused=True,
+                    )
+                    anchor_norm = _grad_l2_norm(list(anchor_grads))
+                    interaction_norm = _grad_l2_norm(list(interaction_grads))
+                    if interaction_norm > 0.0:
+                        mu_effective = float(
+                            mu_auto_rho
+                            * (anchor_norm / (interaction_norm + float(mu_auto_eps)))
+                        )
+                    else:
+                        mu_effective = float(base_mu)
+                    if not math.isfinite(mu_effective):
+                        mu_effective = float(base_mu)
+
+                total_loss = anchor_loss
+                if interaction_loss is not None:
+                    total_loss = total_loss + (float(mu_effective) * interaction_loss)
+
+            if grad_accum > 1:
+                total_loss = total_loss / float(grad_accum)
+
+            if use_scaler:
+                scaler.scale(total_loss).backward()
+            else:
+                total_loss.backward()
+
+            # Only the sampled receiver layer updates on this micro-batch.
+            for param in trainable_params:
+                if id(param) in receiver_param_ids:
+                    continue
+                if param.grad is not None:
+                    if comm_train_mode == "lora":
+                        param.grad.zero_()
+                    else:
+                        param.grad = None
+
+            total_loss_sum += float(total_loss.detach().float().item())
+            anchor_sum += float(anchor_loss.detach().float().item())
+            if interaction_loss is not None:
+                interaction_sum += float(interaction_loss.detach().float().item())
+            mu_sum += float(mu_effective)
+            counted += 1
+            accum_done += 1
+
+        if args.distill_max_grad_norm is not None:
+            if use_scaler:
+                scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(
+                trainable_params,
+                float(args.distill_max_grad_norm),
+            )
+
+        if use_scaler:
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            optimizer.step()
+
+        opt_step += 1
+        if log_steps and (opt_step == 1 or opt_step % log_steps == 0):
+            denom = max(counted, 1)
+            print(
+                f"[comm] step={opt_step}/{target_opt_steps} "
+                f"loss={total_loss_sum/denom:.6f} "
+                f"anchor={anchor_sum/denom:.6f} "
+                f"int={interaction_sum/denom:.6f} "
+                f"mu={mu_sum/denom:.6f}"
+            )
+
+    if comm_train_mode == "lora":
+        merge_lora_adapters(student_model)
+
+    stats: Dict[str, object] = {
+        "enabled": True,
+        "train_mode": comm_train_mode,
+        "opt_steps": int(target_opt_steps),
+        "grad_accum_steps": int(grad_accum),
+        "lr": float(lr),
+        "temp": float(temp),
+        "steps_ratio": float(steps_ratio),
+        "lr_scale": float(lr_scale),
+        "interaction_mode": interaction_mode,
+        "interaction_eps": float(interaction_eps),
+        "mu": float(base_mu),
+        "mu_auto": bool(mu_auto),
+        "mu_auto_rho": float(mu_auto_rho),
+        "mu_auto_eps": float(mu_auto_eps),
+        "sample_eta": float(sample_eta),
+        "sample_dwce_scale": float(sample_dwce_scale),
+        "topk": int(top_k),
+        "candidate_pairs": [int(i) for i in candidate_pairs],
+        "trainable_params": int(sum(int(param.numel()) for param in trainable_params)),
+    }
+    total_samples = int(sum(pair_counts))
+    probs_list = [float(x) for x in probs_by_pair]
+    freqs = (
+        [float(c) / float(total_samples) for c in pair_counts]
+        if total_samples > 0
+        else [0.0 for _ in pair_counts]
+    )
+    top_show = min(10, num_pairs)
+    top_indices = sorted(range(num_pairs), key=lambda i: pair_counts[i], reverse=True)[:top_show]
+    top_pairs = [
+        {
+            "pair": int(i),
+            "count": int(pair_counts[i]),
+            "freq": float(freqs[i]),
+            "prob": float(probs_list[i]) if i < len(probs_list) else None,
+        }
+        for i in top_indices
+        if pair_counts[i] > 0
+    ]
+    stats["pair_selection"] = {
+        "num_pairs": int(num_pairs),
+        "excluded_pairs": sorted(exclude_set),
+        "candidate_pairs": [int(i) for i in candidate_pairs],
+        "total_samples": total_samples,
+        "unique_pairs": int(sum(1 for c in pair_counts if c > 0)),
+        "counts": [int(c) for c in pair_counts],
+        "freqs": freqs,
+        "probs": probs_list,
+        "top_pairs": top_pairs,
+    }
+
+    if total_samples > 0 and top_pairs:
+        top_str = ", ".join(
+            f"{entry['pair']}-{entry['pair'] + 1}: {entry['count']} "
+            f"(obs={entry['freq']:.3f}, exp={entry['prob']:.3f})"
+            for entry in top_pairs
+            if entry.get("prob") is not None
+        )
+        if not top_str:
+            top_str = ", ".join(
+                f"{entry['pair']}-{entry['pair'] + 1}: {entry['count']} "
+                f"(obs={entry['freq']:.3f})"
+                for entry in top_pairs
+            )
+        print(
+            f"[comm] Pair sampling stats: total={total_samples} "
+            f"unique={stats['pair_selection']['unique_pairs']}/{num_pairs} "
+            f"top={top_str}"
+        )
+
+    if counted > 0:
+        stats["avg_loss"] = float(total_loss_sum / float(counted))
+        stats["avg_anchor"] = float(anchor_sum / float(counted))
+        stats["avg_interaction"] = float(interaction_sum / float(counted))
+        stats["avg_mu"] = float(mu_sum / float(counted))
+    return stats

src/fuse_layers_model.py

@@ -0,0 +1,595 @@
+#!/usr/bin/env python3
+"""Model and layer helpers for fuse_layers."""
+
+import os
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+try:
+    from tqdm import tqdm
+except Exception:  # pragma: no cover - optional dependency
+    tqdm = None
+
+
+def _tqdm_enabled() -> bool:
+    value = os.environ.get("DISABLE_TQDM", os.environ.get("TQDM_DISABLE", "0"))
+    return value.strip().lower() not in {"1", "true", "yes", "on"}
+
+
+def get_dtype(dtype: str):
+    if dtype == "auto":
+        return None
+    if dtype == "float16":
+        return torch.float16
+    if dtype == "bfloat16":
+        return torch.bfloat16
+    return torch.float32
+
+
+def resolve_attr(root: object, path: str) -> Optional[object]:
+    cur = root
+    for part in path.split("."):
+        if not hasattr(cur, part):
+            return None
+        cur = getattr(cur, part)
+    return cur
+
+
+def resolve_attr_with_parent(root: object, path: str) -> Tuple[object, str, object]:
+    parts = path.split(".")
+    cur = root
+    for part in parts[:-1]:
+        if not hasattr(cur, part):
+            raise ValueError(f"'{path}' not found on model")
+        cur = getattr(cur, part)
+    name = parts[-1]
+    if not hasattr(cur, name):
+        raise ValueError(f"'{path}' not found on model")
+    return cur, name, getattr(cur, name)
+
+
+def find_layer_container(model, layer_path: Optional[str]) -> Tuple[object, str, object]:
+    if layer_path:
+        parent, name, container = resolve_attr_with_parent(model, layer_path)
+        return parent, name, container
+
+    candidate_paths = [
+        "model.layers",  # LLaMA, Mistral, Qwen2, Gemma
+        "model.decoder.layers",  # OPT
+        "transformer.h",  # GPT-2, GPT-J, Bloom, Falcon
+        "transformer.blocks",  # MPT
+        "gpt_neox.layers",  # GPT-NeoX
+        "layers",  # fallback
+    ]
+    for path in candidate_paths:
+        candidate = resolve_attr(model, path)
+        if candidate is None:
+            continue
+        try:
+            list(candidate)
+        except TypeError:
+            continue
+        parent, name, container = resolve_attr_with_parent(model, path)
+        return parent, name, container
+
+    raise ValueError(
+        "Could not locate transformer layers. Pass --layer_path explicitly."
+    )
+
+
+def find_attention_module(layer: torch.nn.Module) -> torch.nn.Module:
+    if hasattr(layer, "self_attn"):
+        return getattr(layer, "self_attn")
+    if hasattr(layer, "attn"):
+        return getattr(layer, "attn")
+    if hasattr(layer, "attention"):
+        return getattr(layer, "attention")
+    for _, module in layer.named_modules():
+        if all(
+            hasattr(module, attr) for attr in ("q_proj", "k_proj", "v_proj", "o_proj")
+        ):
+            return module
+    raise ValueError("Could not find attention module with q_proj/k_proj/v_proj/o_proj")
+
+
+def find_mlp_module(layer: torch.nn.Module) -> torch.nn.Module:
+    if hasattr(layer, "mlp"):
+        return getattr(layer, "mlp")
+    for attr in ("feed_forward", "feedforward", "ffn", "ff"):
+        if hasattr(layer, attr):
+            return getattr(layer, attr)
+    for _, module in layer.named_modules():
+        if all(hasattr(module, attr) for attr in ("gate_proj", "up_proj", "down_proj")):
+            return module
+        if all(hasattr(module, attr) for attr in ("fc1", "fc2")):
+            return module
+        if all(
+            hasattr(module, attr)
+            for attr in ("dense_h_to_4h", "dense_4h_to_h")
+        ):
+            return module
+        if all(hasattr(module, attr) for attr in ("w1", "w2")):
+            return module
+    raise ValueError("Could not find MLP/FFN module on layer")
+
+
+def get_head_info(
+    attn: torch.nn.Module, hidden_size: int, config
+) -> Tuple[int, int, int]:
+    num_heads = getattr(attn, "num_heads", None)
+    if num_heads is None:
+        num_heads = getattr(attn, "num_attention_heads", None)
+    if num_heads is None and config is not None:
+        num_heads = getattr(
+            config,
+            "num_attention_heads",
+            getattr(config, "num_heads", getattr(config, "n_head", None)),
+        )
+
+    num_key_value_heads = getattr(attn, "num_key_value_heads", None)
+    if num_key_value_heads is None:
+        num_key_value_heads = getattr(attn, "num_kv_heads", None)
+    if num_key_value_heads is None and config is not None:
+        num_key_value_heads = getattr(
+            config,
+            "num_key_value_heads",
+            getattr(config, "num_kv_heads", getattr(config, "n_head_kv", None)),
+        )
+
+    head_dim = getattr(attn, "head_dim", None)
+    if head_dim is None and config is not None:
+        head_dim = getattr(config, "head_dim", None)
+
+    if num_heads is None:
+        if hasattr(attn, "q_proj"):
+            q_out = attn.q_proj.weight.shape[0]
+            if head_dim is not None:
+                num_heads = q_out // head_dim
+            elif num_key_value_heads is not None and hasattr(attn, "k_proj"):
+                k_out = attn.k_proj.weight.shape[0]
+                head_dim = k_out // max(int(num_key_value_heads), 1)
+                num_heads = q_out // head_dim
+    if num_heads is None:
+        raise ValueError(
+            "Attention module missing num_heads/num_attention_heads; "
+            "pass --layer_path or add config overrides."
+        )
+
+    if num_key_value_heads is None:
+        num_key_value_heads = num_heads
+
+    if head_dim is None:
+        head_dim = hidden_size // int(num_heads)
+
+    if num_key_value_heads is None and hasattr(attn, "k_proj"):
+        k_out = attn.k_proj.weight.shape[0]
+        num_key_value_heads = k_out // int(head_dim)
+
+    return int(num_heads), int(num_key_value_heads), int(head_dim)
+
+
+def cosine_cost_matrix(
+    a: torch.Tensor, b: torch.Tensor, eps: float = 1e-8
+) -> torch.Tensor:
+    a_norm = a / (a.norm(dim=1, keepdim=True) + eps)
+    b_norm = b / (b.norm(dim=1, keepdim=True) + eps)
+    sim = a_norm @ b_norm.t()
+    return 1.0 - sim
+
+
+def hungarian(cost: torch.Tensor) -> List[int]:
+    # Kuhn-Munkres for square cost matrix (minimization).
+    n = cost.size(0)
+    u = [0.0] * (n + 1)
+    v = [0.0] * (n + 1)
+    p = [0] * (n + 1)
+    way = [0] * (n + 1)
+
+    for i in range(1, n + 1):
+        p[0] = i
+        j0 = 0
+        minv = [float("inf")] * (n + 1)
+        used = [False] * (n + 1)
+        while True:
+            used[j0] = True
+            i0 = p[j0]
+            delta = float("inf")
+            j1 = 0
+            for j in range(1, n + 1):
+                if used[j]:
+                    continue
+                cur = cost[i0 - 1, j - 1].item() - u[i0] - v[j]
+                if cur < minv[j]:
+                    minv[j] = cur
+                    way[j] = j0
+                if minv[j] < delta:
+                    delta = minv[j]
+                    j1 = j
+            for j in range(0, n + 1):
+                if used[j]:
+                    u[p[j]] += delta
+                    v[j] -= delta
+                else:
+                    minv[j] -= delta
+            j0 = j1
+            if p[j0] == 0:
+                break
+        while True:
+            j1 = way[j0]
+            p[j0] = p[j1]
+            j0 = j1
+            if j0 == 0:
+                break
+
+    assignment = [-1] * n
+    for j in range(1, n + 1):
+        if p[j] > 0:
+            assignment[p[j] - 1] = j - 1
+    return assignment
+
+
+def compute_head_means(
+    model,
+    attn_i: torch.nn.Module,
+    attn_j: torch.nn.Module,
+    dataloader,
+    device: str,
+    hidden_size: int,
+) -> Tuple[torch.Tensor, torch.Tensor, int, int, int]:
+    num_heads_i, num_kv_i, head_dim_i = get_head_info(attn_i, hidden_size, model.config)
+    num_heads_j, num_kv_j, head_dim_j = get_head_info(attn_j, hidden_size, model.config)
+    if num_heads_i != num_heads_j or head_dim_i != head_dim_j:
+        raise ValueError("Head counts or head_dim differ between layers; cannot align")
+
+    sums_i = torch.zeros(num_heads_i, head_dim_i, device="cpu")
+    sums_j = torch.zeros(num_heads_j, head_dim_j, device="cpu")
+    count_i = [0]
+    count_j = [0]
+
+    def make_hook(
+        sums: torch.Tensor, count_ref: List[int], num_heads: int, head_dim: int
+    ):
+        def hook(_module, inputs, _output):
+            hidden = inputs[0].detach()
+            if hidden.dim() != 3:
+                return
+            batch, seq, width = hidden.shape
+            if width != num_heads * head_dim:
+                return
+            reshaped = hidden.view(batch, seq, num_heads, head_dim)
+            sums.add_(reshaped.sum(dim=(0, 1)).float().cpu())
+            count_ref[0] += batch * seq
+
+        return hook
+
+    hook_i = attn_i.o_proj.register_forward_hook(
+        make_hook(sums_i, count_i, num_heads_i, head_dim_i)
+    )
+    hook_j = attn_j.o_proj.register_forward_hook(
+        make_hook(sums_j, count_j, num_heads_j, head_dim_j)
+    )
+
+    model.eval()
+    iterator = dataloader
+    if tqdm is not None and _tqdm_enabled():
+        iterator = tqdm(dataloader, desc="Head stats", unit="batch")
+    with torch.no_grad():
+        for batch in iterator:
+            input_ids = batch[0].to(device)
+            _ = model(input_ids=input_ids)
+
+    hook_i.remove()
+    hook_j.remove()
+
+    if count_i[0] == 0 or count_j[0] == 0:
+        raise RuntimeError("Failed to capture head outputs; check attention modules.")
+
+    mean_i = sums_i / count_i[0]
+    mean_j = sums_j / count_j[0]
+    return mean_i, mean_j, num_heads_i, num_kv_i, head_dim_i
+
+
+def build_head_permutation(
+    mean_i: torch.Tensor,
+    mean_j: torch.Tensor,
+    num_heads: int,
+    num_kv_heads: int,
+    eps: float,
+) -> List[int]:
+    group_size = num_heads // num_kv_heads
+    if group_size * num_kv_heads != num_heads:
+        raise ValueError("num_heads must be divisible by num_key_value_heads")
+
+    perm = list(range(num_heads))
+    for g in range(num_kv_heads):
+        start = g * group_size
+        end = start + group_size
+        cost = cosine_cost_matrix(mean_i[start:end], mean_j[start:end], eps=eps)
+        assignment = hungarian(cost)
+        for local_idx, match in enumerate(assignment):
+            perm[start + local_idx] = start + match
+    return perm
+
+
+def permute_attention_heads(
+    attn: torch.nn.Module,
+    perm: List[int],
+    num_heads: int,
+    num_kv_heads: int,
+    head_dim: int,
+) -> None:
+    hidden_size = num_heads * head_dim
+
+    def permute_out_proj_weight(weight: torch.Tensor) -> torch.Tensor:
+        out_features, in_features = weight.shape
+        if in_features != hidden_size:
+            raise ValueError(
+                "o_proj in_features ({} ) != num_heads*head_dim ({})".format(
+                    in_features, hidden_size
+                )
+            )
+        reshaped = weight.view(out_features, num_heads, head_dim)
+        reshaped = reshaped[:, perm, :]
+        return reshaped.reshape(out_features, in_features)
+
+    def permute_proj_weight(weight: torch.Tensor) -> torch.Tensor:
+        out_features, in_features = weight.shape
+        if out_features != hidden_size:
+            raise ValueError(
+                "proj out_features ({}) != num_heads*head_dim ({})".format(
+                    out_features, hidden_size
+                )
+            )
+        reshaped = weight.view(num_heads, head_dim, in_features)
+        reshaped = reshaped[perm, :, :]
+        return reshaped.reshape(out_features, in_features)
+
+    def permute_proj_bias(bias: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        if bias is None:
+            return None
+        reshaped = bias.view(num_heads, head_dim)
+        reshaped = reshaped[perm, :]
+        return reshaped.reshape(num_heads * head_dim)
+
+    with torch.no_grad():
+        attn.q_proj.weight.copy_(permute_proj_weight(attn.q_proj.weight))
+        if attn.q_proj.bias is not None:
+            attn.q_proj.bias.copy_(permute_proj_bias(attn.q_proj.bias))
+
+        if num_kv_heads == num_heads:
+            attn.k_proj.weight.copy_(permute_proj_weight(attn.k_proj.weight))
+            if attn.k_proj.bias is not None:
+                attn.k_proj.bias.copy_(permute_proj_bias(attn.k_proj.bias))
+            attn.v_proj.weight.copy_(permute_proj_weight(attn.v_proj.weight))
+            if attn.v_proj.bias is not None:
+                attn.v_proj.bias.copy_(permute_proj_bias(attn.v_proj.bias))
+
+        attn.o_proj.weight.copy_(permute_out_proj_weight(attn.o_proj.weight))
+
+
+def compute_fisher(
+    model,
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    dataloader,
+    fisher_mode: str,
+    device: str,
+) -> Tuple[List[Dict[str, object]], int, List[Dict[str, int]]]:
+    for param in model.parameters():
+        param.requires_grad_(False)
+    for layer in (layer_a, layer_b):
+        for param in layer.parameters():
+            param.requires_grad_(True)
+
+    fisher_sums: List[Dict[str, object]] = []
+    param_numels: List[Dict[str, int]] = []
+    for layer in (layer_a, layer_b):
+        layer_sums: Dict[str, object] = {}
+        layer_numels: Dict[str, int] = {}
+        for name, param in layer.named_parameters():
+            if not param.requires_grad:
+                continue
+            if fisher_mode == "param":
+                layer_sums[name] = torch.zeros_like(
+                    param, dtype=torch.float32, device="cpu"
+                )
+            else:
+                layer_sums[name] = 0.0
+                layer_numels[name] = param.numel()
+        fisher_sums.append(layer_sums)
+        param_numels.append(layer_numels)
+
+    num_batches = 0
+    model.eval()
+    iterator = dataloader
+    if tqdm is not None and _tqdm_enabled():
+        iterator = tqdm(dataloader, desc="Fisher", unit="batch")
+    for batch in iterator:
+        input_ids = batch[0].to(device)
+        outputs = model(input_ids=input_ids, labels=input_ids)
+        loss = outputs.loss
+        loss.backward()
+        for layer_idx, layer in enumerate((layer_a, layer_b)):
+            layer_sums = fisher_sums[layer_idx]
+            for name, param in layer.named_parameters():
+                if not param.requires_grad or param.grad is None:
+                    continue
+                grad_sq = param.grad.detach().float().pow(2)
+                if fisher_mode == "param":
+                    layer_sums[name] += grad_sq.cpu()
+                else:
+                    layer_sums[name] += float(grad_sq.sum().item())
+        model.zero_grad(set_to_none=True)
+        num_batches += 1
+
+    if num_batches == 0:
+        raise RuntimeError("No batches processed; check dataset or text inputs.")
+
+    return fisher_sums, num_batches, param_numels
+
+
+def merge_layers(
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    fisher_a: Dict[str, object],
+    fisher_b: Dict[str, object],
+    num_batches: int,
+    numels_a: Dict[str, int],
+    numels_b: Dict[str, int],
+    fisher_mode: str,
+    eps: float,
+) -> int:
+    merged = 0
+    params_b = {name: param for name, param in layer_b.named_parameters()}
+    with torch.no_grad():
+        for name, param_a in layer_a.named_parameters():
+            param_b = params_b.get(name)
+            if param_b is None or param_b.shape != param_a.shape:
+                continue
+            if fisher_mode == "param":
+                fa = fisher_a[name] / num_batches
+                fb = fisher_b[name] / num_batches
+                # Fisher tensors are accumulated on CPU to save VRAM; move to the
+                # parameter device for the actual merge.
+                if isinstance(fa, torch.Tensor) and fa.device != param_a.device:
+                    fa = fa.to(param_a.device)
+                if isinstance(fb, torch.Tensor) and fb.device != param_a.device:
+                    fb = fb.to(param_a.device)
+                denom = fa + fb
+                denom_mean = float(denom.mean().item())
+                if denom_mean <= eps:
+                    merged_param = 0.5 * (param_a.float() + param_b.float())
+                else:
+                    merged_param = (fa * param_a.float() + fb * param_b.float()) / (
+                        denom + eps
+                    )
+            else:
+                fa = fisher_a[name] / (num_batches * numels_a[name])
+                fb = fisher_b[name] / (num_batches * numels_b[name])
+                denom = fa + fb
+                if denom <= eps:
+                    merged_param = 0.5 * (param_a.float() + param_b.float())
+                else:
+                    merged_param = (
+                        fa * param_a.float() + fb * param_b.float()
+                    ) / (denom + eps)
+            param_a.copy_(merged_param.to(dtype=param_a.dtype))
+            merged += 1
+    return merged
+
+
+def merge_layers_with_gates(
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    gates: Dict[str, torch.Tensor],
+) -> int:
+    """Merge layer_b into layer_a using precomputed gates.
+
+    Each gate is a lambda in [0, 1] that mixes parameters as:
+      W = lambda * W_a + (1 - lambda) * W_b
+
+    Gate tensors may be scalars (per-tensor gating) or full tensors matching the
+    parameter shape (per-parameter gating).
+    """
+    merged = 0
+    params_b = {name: param for name, param in layer_b.named_parameters()}
+    with torch.no_grad():
+        for name, param_a in layer_a.named_parameters():
+            gate = gates.get(name)
+            if gate is None:
+                continue
+            param_b = params_b.get(name)
+            if param_b is None or param_b.shape != param_a.shape:
+                continue
+            lam = gate
+            if not isinstance(lam, torch.Tensor):
+                lam = torch.tensor(lam)
+            if lam.device != param_a.device:
+                lam = lam.to(param_a.device)
+            merged_param = lam * param_a.float() + (1.0 - lam) * param_b.float()
+            param_a.copy_(merged_param.to(dtype=param_a.dtype))
+            merged += 1
+    return merged
+
+
+def drop_layer(container: object, index: int) -> object:
+    if isinstance(container, torch.nn.ModuleList):
+        return torch.nn.ModuleList(
+            [layer for idx, layer in enumerate(container) if idx != index]
+        )
+    if isinstance(container, list):
+        del container[index]
+        return container
+    raise TypeError("Layer container must be ModuleList or list")
+
+
+def decrement_config(config) -> None:
+    for attr in ("num_hidden_layers", "n_layer", "num_layers"):
+        if hasattr(config, attr):
+            value = getattr(config, attr)
+            if isinstance(value, int) and value > 0:
+                setattr(config, attr, value - 1)
+    normalize_config(config)
+
+
+def normalize_config(config) -> None:
+    num_hidden_layers = getattr(config, "num_hidden_layers", None)
+    layer_types = getattr(config, "layer_types", None)
+    if (
+        isinstance(num_hidden_layers, int)
+        and num_hidden_layers >= 0
+        and isinstance(layer_types, (list, tuple))
+        and len(layer_types) != num_hidden_layers
+    ):
+        config.layer_types = list(layer_types[:num_hidden_layers])
+
+
+def find_colon_modules(module: torch.nn.Module) -> List[str]:
+    found: List[str] = []
+    for name, child in module._modules.items():
+        if ":" in name:
+            found.append(name)
+        if isinstance(child, torch.nn.Module):
+            for sub in find_colon_modules(child):
+                found.append(f"{name}.{sub}")
+    return found
+
+
+def get_norm_pair(
+    layer: torch.nn.Module,
+) -> Tuple[
+    Optional[torch.nn.Module],
+    Optional[torch.nn.Module],
+    Tuple[Optional[str], Optional[str]],
+]:
+    candidates = [
+        ("input_layernorm", "post_attention_layernorm"),
+        ("ln_1", "ln_2"),
+        ("norm1", "norm2"),
+        ("norm_1", "norm_2"),
+        ("layer_norm_1", "layer_norm_2"),
+        ("self_attn_layer_norm", "final_layer_norm"),
+    ]
+    for n1, n2 in candidates:
+        if hasattr(layer, n1) and hasattr(layer, n2):
+            return getattr(layer, n1), getattr(layer, n2), (n1, n2)
+    return None, None, (None, None)
+
+
+def clone_state_dict(module: torch.nn.Module) -> Dict[str, torch.Tensor]:
+    return {k: v.detach().clone() for k, v in module.state_dict().items()}
+
+
+def apply_norm_policy(
+    layer: torch.nn.Module,
+    norm_policy: str,
+    norm1_state: Optional[Dict[str, torch.Tensor]],
+    norm2_state: Optional[Dict[str, torch.Tensor]],
+    norm_names: Tuple[Optional[str], Optional[str]],
+) -> None:
+    norm1, norm2, _ = get_norm_pair(layer)
+    if norm_policy in {"copy_n1", "hybrid"} and norm1_state is not None and norm1 is not None:
+        norm1.load_state_dict(norm1_state)
+    if norm_policy == "copy_n1_n2" and norm2_state is not None and norm2 is not None:
+        norm2.load_state_dict(norm2_state)

src/fuse_layers_select.py

@@ -0,0 +1,1152 @@
+#!/usr/bin/env python3
+"""Automatic adjacent-pair selection via configurable scoring metrics."""
+
+import copy
+import math
+from contextlib import contextmanager
+from typing import Dict, List, Optional, Set, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from fuse_layers_model import (
+    build_head_permutation,
+    compute_fisher,
+    compute_head_means,
+    find_attention_module,
+    find_layer_container,
+    merge_layers,
+    permute_attention_heads,
+)
+
+_DWCE_GRAD_CACHE_MAX_BYTES = 1 << 30
+
+
+class _DwceGradCacheOverflow(RuntimeError):
+    """Raised when shared-backward DWCE caching exceeds the configured budget."""
+
+
+def _get_hidden_size(model) -> int:
+    hidden_size = getattr(model.config, "hidden_size", None)
+    if hidden_size is None:
+        hidden_size = getattr(model.config, "n_embd", None)
+    if hidden_size is None:
+        raise SystemExit("Model config missing hidden_size/n_embd")
+    return int(hidden_size)
+
+
+def _detach_arg(arg):
+    if torch.is_tensor(arg):
+        return arg.detach()
+    if isinstance(arg, (list, tuple)):
+        return type(arg)(_detach_arg(x) for x in arg)
+    if isinstance(arg, dict):
+        return {k: _detach_arg(v) for k, v in arg.items()}
+    return arg
+
+
+def _register_forward_hook(layer, hook):
+    try:
+        def wrapper(module, inputs, kwargs, output):
+            return hook(module, inputs, output, kwargs)
+
+        handle = layer.register_forward_hook(wrapper, with_kwargs=True)
+        return handle, True
+    except TypeError:
+        def wrapper(module, inputs, output):
+            return hook(module, inputs, output, None)
+        handle = layer.register_forward_hook(wrapper)
+        return handle, False
+
+
+@contextmanager
+def _temporary_layers(parent: object, name: str, new_layers: object):
+    original = getattr(parent, name)
+    setattr(parent, name, new_layers)
+    try:
+        yield
+    finally:
+        setattr(parent, name, original)
+
+
+def _extract_hidden(output):
+    if torch.is_tensor(output):
+        return output
+    if isinstance(output, (tuple, list)):
+        if output and all(torch.is_tensor(item) for item in output):
+            return output[0]
+        for item in output:
+            hidden = _extract_hidden(item)
+            if hidden is not None:
+                return hidden
+        return None
+    if isinstance(output, dict):
+        for key in ("hidden_states", "last_hidden_state", "hidden_state"):
+            if key in output:
+                value = output[key]
+                if isinstance(value, (tuple, list)) and value and all(
+                    torch.is_tensor(item) for item in value
+                ):
+                    return value[-1]
+                hidden = _extract_hidden(value)
+                if hidden is not None:
+                    return hidden
+        for value in output.values():
+            hidden = _extract_hidden(value)
+            if hidden is not None:
+                return hidden
+        return None
+    for attr in ("hidden_states", "last_hidden_state"):
+        if hasattr(output, attr):
+            value = getattr(output, attr)
+            if isinstance(value, (tuple, list)) and value and all(
+                torch.is_tensor(item) for item in value
+            ):
+                return value[-1]
+            hidden = _extract_hidden(value)
+            if hidden is not None:
+                return hidden
+    return None
+
+
+def _build_fused_layer_for_pair(
+    model,
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    dataloader,
+    device: str,
+    fisher_mode: str,
+    eps: float,
+    hidden_size: int,
+    enable_head_permute: bool = True,
+) -> Tuple[torch.nn.Module, Dict[str, float]]:
+    attn_a = find_attention_module(layer_a)
+    attn_b = find_attention_module(layer_b)
+    perm = None
+    inv_perm = None
+    num_heads = None
+    num_kv_heads = None
+    head_dim = None
+    if enable_head_permute:
+        mean_a, mean_b, num_heads, num_kv_heads, head_dim = compute_head_means(
+            model,
+            attn_a,
+            attn_b,
+            dataloader,
+            device,
+            hidden_size,
+        )
+
+        perm = build_head_permutation(
+            mean_a,
+            mean_b,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            eps=eps,
+        )
+
+    layer_a_copy = copy.deepcopy(layer_a)
+    layer_b_copy = copy.deepcopy(layer_b)
+    attn_b_copy = find_attention_module(layer_b_copy)
+    if perm is not None:
+        permute_attention_heads(
+            attn_b_copy, perm, num_heads, num_kv_heads, head_dim=head_dim
+        )
+
+        inv_perm = [0] * len(perm)
+        for idx, mapped in enumerate(perm):
+            inv_perm[mapped] = idx
+
+        permute_attention_heads(attn_b, perm, num_heads, num_kv_heads, head_dim=head_dim)
+    try:
+        fisher_sums, num_batches, param_numels = compute_fisher(
+            model,
+            layer_a,
+            layer_b,
+            dataloader,
+            fisher_mode=fisher_mode,
+            device=device,
+        )
+    finally:
+        if inv_perm is not None:
+            permute_attention_heads(
+                attn_b, inv_perm, num_heads, num_kv_heads, head_dim=head_dim
+            )
+
+    merge_layers(
+        layer_a_copy,
+        layer_b_copy,
+        fisher_sums[0],
+        fisher_sums[1],
+        num_batches,
+        param_numels[0],
+        param_numels[1],
+        fisher_mode=fisher_mode,
+        eps=eps,
+    )
+
+    # Scalar mixing coefficients per parameter tensor; used by pressure redistribution
+    # to simulate future fusions without running another Fisher pass.
+    fuse_priors: Dict[str, float] = {}
+    params_b = {name: param for name, param in layer_b.named_parameters()}
+    clamp_eps = 1e-4
+    for name, param_a in layer_a.named_parameters():
+        param_b = params_b.get(name)
+        if param_b is None or param_b.shape != param_a.shape:
+            continue
+        if fisher_mode == "param":
+            fa = fisher_sums[0][name] / max(num_batches, 1)
+            fb = fisher_sums[1][name] / max(num_batches, 1)
+            if isinstance(fa, torch.Tensor):
+                fa_val = float(fa.mean().item())
+            else:
+                fa_val = float(fa)
+            if isinstance(fb, torch.Tensor):
+                fb_val = float(fb.mean().item())
+            else:
+                fb_val = float(fb)
+        else:
+            fa_val = float(
+                fisher_sums[0][name]
+                / (max(num_batches, 1) * max(param_numels[0].get(name, 1), 1))
+            )
+            fb_val = float(
+                fisher_sums[1][name]
+                / (max(num_batches, 1) * max(param_numels[1].get(name, 1), 1))
+            )
+        denom = fa_val + fb_val
+        if denom <= eps:
+            lam = 0.5
+        else:
+            lam = fa_val / (denom + eps)
+        lam = min(max(lam, clamp_eps), 1.0 - clamp_eps)
+        fuse_priors[name] = lam
+
+    layer_a_copy.eval()
+    return layer_a_copy, fuse_priors
+
+
+def _init_fisher_accumulators(
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    fisher_mode: str,
+    device: str,
+) -> Tuple[List[Dict[str, object]], List[Dict[str, int]]]:
+    fisher_sums: List[Dict[str, object]] = []
+    param_numels: List[Dict[str, int]] = []
+    for layer in (layer_a, layer_b):
+        layer_sums: Dict[str, object] = {}
+        layer_numels: Dict[str, int] = {}
+        for name, param in layer.named_parameters():
+            if not param.requires_grad:
+                continue
+            if fisher_mode == "param":
+                layer_sums[name] = torch.zeros_like(
+                    param, dtype=torch.float32, device="cpu"
+                )
+            else:
+                layer_sums[name] = torch.zeros((), dtype=torch.float32, device=device)
+                layer_numels[name] = param.numel()
+        fisher_sums.append(layer_sums)
+        param_numels.append(layer_numels)
+    return fisher_sums, param_numels
+
+
+def _accumulate_fisher_from_grads(
+    layer: torch.nn.Module,
+    layer_sums: Dict[str, object],
+    fisher_mode: str,
+) -> None:
+    for name, param in layer.named_parameters():
+        if not param.requires_grad or param.grad is None:
+            continue
+        grad_sq = param.grad.detach().float().pow(2)
+        if fisher_mode == "param":
+            layer_sums[name] += grad_sq.cpu()
+        else:
+            layer_sums[name] += grad_sq.sum()
+
+
+def _finalize_fisher_sums(
+    fisher_sums: List[Dict[str, object]],
+    fisher_mode: str,
+) -> List[Dict[str, object]]:
+    if fisher_mode == "param":
+        return fisher_sums
+
+    finalized: List[Dict[str, object]] = []
+    for layer_sums in fisher_sums:
+        finalized_layer: Dict[str, object] = {}
+        for name, value in layer_sums.items():
+            if isinstance(value, torch.Tensor):
+                finalized_layer[name] = float(value.detach().cpu().item())
+            else:
+                finalized_layer[name] = float(value)
+        finalized.append(finalized_layer)
+    return finalized
+
+
+def _compute_fuse_priors(
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    fisher_sums: List[Dict[str, object]],
+    num_batches: int,
+    param_numels: List[Dict[str, int]],
+    fisher_mode: str,
+    eps: float,
+) -> Dict[str, float]:
+    fuse_priors: Dict[str, float] = {}
+    params_b = {name: param for name, param in layer_b.named_parameters()}
+    clamp_eps = 1e-4
+    for name, param_a in layer_a.named_parameters():
+        param_b = params_b.get(name)
+        if param_b is None or param_b.shape != param_a.shape:
+            continue
+        if fisher_mode == "param":
+            fa = fisher_sums[0][name] / max(num_batches, 1)
+            fb = fisher_sums[1][name] / max(num_batches, 1)
+            fa_val = float(fa.mean().item()) if isinstance(fa, torch.Tensor) else float(fa)
+            fb_val = float(fb.mean().item()) if isinstance(fb, torch.Tensor) else float(fb)
+        else:
+            fa_val = float(
+                fisher_sums[0][name]
+                / (max(num_batches, 1) * max(param_numels[0].get(name, 1), 1))
+            )
+            fb_val = float(
+                fisher_sums[1][name]
+                / (max(num_batches, 1) * max(param_numels[1].get(name, 1), 1))
+            )
+        denom = fa_val + fb_val
+        lam = 0.5 if denom <= eps else fa_val / (denom + eps)
+        fuse_priors[name] = min(max(lam, clamp_eps), 1.0 - clamp_eps)
+    return fuse_priors
+
+
+def _score_dwce_with_shared_backward(
+    model,
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    dataloader,
+    device: str,
+    fisher_mode: str,
+    max_batches: int,
+    eps: float,
+    norm: str,
+    hidden_size: int,
+    enable_head_permute: bool = True,
+) -> Tuple[float, Dict[str, object]]:
+    attn_a = find_attention_module(layer_a)
+    attn_b = find_attention_module(layer_b)
+    perm = None
+    inv_perm = None
+    num_heads = None
+    num_kv_heads = None
+    head_dim = None
+    if enable_head_permute:
+        mean_a, mean_b, num_heads, num_kv_heads, head_dim = compute_head_means(
+            model,
+            attn_a,
+            attn_b,
+            dataloader,
+            device,
+            hidden_size,
+        )
+        perm = build_head_permutation(
+            mean_a,
+            mean_b,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            eps=eps,
+        )
+
+    layer_a_copy = copy.deepcopy(layer_a)
+    layer_b_copy = copy.deepcopy(layer_b)
+    attn_b_copy = find_attention_module(layer_b_copy)
+    if perm is not None:
+        permute_attention_heads(
+            attn_b_copy, perm, num_heads, num_kv_heads, head_dim=head_dim
+        )
+
+        inv_perm = [0] * len(perm)
+        for idx, mapped in enumerate(perm):
+            inv_perm[mapped] = idx
+
+    cache: Dict[str, Optional[torch.Tensor]] = {"teacher": None}
+    grad_sq_cache: List[torch.Tensor] = []
+    cached_bytes = 0
+
+    def hook_b(_module, _inputs, output, _kwargs=None):
+        teacher_hidden = _extract_hidden(output)
+        if teacher_hidden is None:
+            raise RuntimeError("Failed to extract teacher hidden state output.")
+        cache["teacher"] = teacher_hidden
+        if teacher_hidden.requires_grad:
+            teacher_hidden.retain_grad()
+        return output
+
+    handle_b, _ = _register_forward_hook(layer_b, hook_b)
+    for param in model.parameters():
+        param.requires_grad_(False)
+    for layer in (layer_a, layer_b):
+        for param in layer.parameters():
+            param.requires_grad_(True)
+    fisher_sums, param_numels = _init_fisher_accumulators(
+        layer_a, layer_b, fisher_mode, device
+    )
+    num_batches = 0
+
+    if perm is not None:
+        permute_attention_heads(attn_b, perm, num_heads, num_kv_heads, head_dim=head_dim)
+    try:
+        model.eval()
+        for batch_idx, batch in enumerate(dataloader):
+            if max_batches and batch_idx >= max_batches:
+                break
+            cache["teacher"] = None
+            input_ids = batch[0].to(device)
+            attention_mask = batch[1].to(device) if len(batch) > 1 else None
+
+            model.zero_grad(set_to_none=True)
+            outputs = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=input_ids,
+            )
+            outputs.loss.backward()
+
+            teacher = cache["teacher"]
+            grad = None if teacher is None else teacher.grad
+            if teacher is None or grad is None:
+                raise RuntimeError(
+                    "Auto selection hooks failed to capture outputs/gradients. "
+                    "Try updating PyTorch or run with --layer <index>."
+                )
+            grad_sq = grad.detach().pow(2).to(device=device, dtype=torch.float16)
+            cached_bytes += grad_sq.numel() * grad_sq.element_size()
+            if cached_bytes > _DWCE_GRAD_CACHE_MAX_BYTES:
+                raise _DwceGradCacheOverflow(
+                    "DWCE grad cache exceeded device-memory budget during shared-backward scoring."
+                )
+            grad_sq_cache.append(grad_sq)
+            _accumulate_fisher_from_grads(layer_a, fisher_sums[0], fisher_mode)
+            _accumulate_fisher_from_grads(layer_b, fisher_sums[1], fisher_mode)
+            model.zero_grad(set_to_none=True)
+            num_batches += 1
+    finally:
+        handle_b.remove()
+        if inv_perm is not None:
+            permute_attention_heads(
+                attn_b, inv_perm, num_heads, num_kv_heads, head_dim=head_dim
+            )
+        for param in model.parameters():
+            param.requires_grad_(True)
+
+    if num_batches == 0:
+        raise RuntimeError("No batches processed; check dataset or text inputs.")
+
+    fisher_sums = _finalize_fisher_sums(fisher_sums, fisher_mode)
+    merge_layers(
+        layer_a_copy,
+        layer_b_copy,
+        fisher_sums[0],
+        fisher_sums[1],
+        num_batches,
+        param_numels[0],
+        param_numels[1],
+        fisher_mode=fisher_mode,
+        eps=eps,
+    )
+    fuse_priors = _compute_fuse_priors(
+        layer_a,
+        layer_b,
+        fisher_sums,
+        num_batches,
+        param_numels,
+        fisher_mode,
+        eps,
+    )
+
+    fused_layer = layer_a_copy
+    fused_layer.eval()
+    phase2_cache = {"teacher": None, "fused": None}
+
+    def hook_a(_module, inputs, output, kwargs=None):
+        with torch.no_grad():
+            detached_inputs = tuple(_detach_arg(arg) for arg in inputs)
+            if kwargs:
+                detached_kwargs = {k: _detach_arg(v) for k, v in kwargs.items()}
+                fused_out = fused_layer(*detached_inputs, **detached_kwargs)
+            else:
+                fused_out = fused_layer(*detached_inputs)
+        fused_hidden = _extract_hidden(fused_out)
+        if fused_hidden is None:
+            raise RuntimeError("Failed to extract fused hidden state output.")
+        phase2_cache["fused"] = fused_hidden
+        return output
+
+    def hook_b_eval(_module, _inputs, output, _kwargs=None):
+        teacher_hidden = _extract_hidden(output)
+        if teacher_hidden is None:
+            raise RuntimeError("Failed to extract teacher hidden state output.")
+        phase2_cache["teacher"] = teacher_hidden
+        return output
+
+    handle_a, has_kwargs_a = _register_forward_hook(layer_a, hook_a)
+    handle_b_eval, has_kwargs_b = _register_forward_hook(layer_b, hook_b_eval)
+    supports_kwargs = has_kwargs_a and has_kwargs_b
+
+    score_num = 0.0
+    score_den = 0.0
+    token_count = 0.0
+    try:
+        model.eval()
+        for batch_idx, batch in enumerate(dataloader):
+            if batch_idx >= num_batches:
+                break
+            phase2_cache["teacher"] = None
+            phase2_cache["fused"] = None
+            input_ids = batch[0].to(device)
+            attention_mask = batch[1].to(device) if len(batch) > 1 else None
+
+            with torch.no_grad():
+                model(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    use_cache=False,
+                )
+
+            teacher = phase2_cache["teacher"]
+            fused = phase2_cache["fused"]
+            if teacher is None or fused is None:
+                raise RuntimeError(
+                    "Auto selection hooks failed to capture outputs during DWCE replay."
+                )
+            grad_sq = grad_sq_cache[batch_idx].to(dtype=torch.float32)
+            if attention_mask is not None:
+                mask = attention_mask.to(dtype=torch.float32).unsqueeze(-1)
+                batch_tokens = float(mask.sum().item())
+                grad_sq = grad_sq * mask
+            else:
+                mask = None
+                batch_tokens = float(input_ids.numel())
+            token_count += batch_tokens
+
+            delta = fused - teacher
+            if mask is not None:
+                delta = delta * mask
+            score_num += (delta.float().pow(2) * grad_sq).sum().item()
+            score_den += (teacher.float().pow(2) * grad_sq).sum().item()
+    finally:
+        handle_a.remove()
+        handle_b_eval.remove()
+
+    score = (
+        score_num / (score_den + eps)
+        if norm == "relative"
+        else score_num / max(token_count, 1.0)
+    )
+    meta = {
+        "num_batches": num_batches,
+        "token_count": token_count,
+        "norm": norm,
+        "supports_kwargs": supports_kwargs,
+        "fuse_priors": fuse_priors,
+        "metric": "dwce",
+        "dwce_mode": "shared",
+    }
+    return score, meta
+
+
+def _compute_dwce_for_pair(
+    model,
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    fused_layer: torch.nn.Module,
+    dataloader,
+    device: str,
+    max_batches: int,
+    eps: float,
+    norm: str,
+) -> Tuple[float, Dict[str, object]]:
+    cache = {"teacher": None, "fused": None}
+    supports_kwargs = True
+
+    def hook_a(_module, inputs, output, kwargs=None):
+        with torch.no_grad():
+            detached_inputs = tuple(_detach_arg(arg) for arg in inputs)
+            if kwargs is not None and len(kwargs) > 0:
+                detached_kwargs = {k: _detach_arg(v) for k, v in kwargs.items()}
+                fused_out = fused_layer(*detached_inputs, **detached_kwargs)
+            else:
+                fused_out = fused_layer(*detached_inputs)
+        fused_hidden = _extract_hidden(fused_out)
+        if fused_hidden is None:
+            raise RuntimeError("Failed to extract fused hidden state output.")
+        cache["fused"] = fused_hidden
+        return output
+
+    def hook_b(_module, _inputs, output, _kwargs=None):
+        teacher_hidden = _extract_hidden(output)
+        if teacher_hidden is None:
+            raise RuntimeError("Failed to extract teacher hidden state output.")
+        cache["teacher"] = teacher_hidden
+        if teacher_hidden.requires_grad:
+            teacher_hidden.retain_grad()
+        return output
+
+    handle_a, has_kwargs_a = _register_forward_hook(layer_a, hook_a)
+    handle_b, has_kwargs_b = _register_forward_hook(layer_b, hook_b)
+    supports_kwargs = has_kwargs_a and has_kwargs_b
+
+    score_num = 0.0
+    score_den = 0.0
+    token_count = 0.0
+    num_batches = 0
+
+    model.eval()
+    for batch_idx, batch in enumerate(dataloader):
+        if max_batches and batch_idx >= max_batches:
+            break
+        cache["teacher"] = None
+        cache["fused"] = None
+
+        input_ids = batch[0].to(device)
+        attention_mask = batch[1].to(device) if len(batch) > 1 else None
+
+        model.zero_grad(set_to_none=True)
+        outputs = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=input_ids,
+        )
+        loss = outputs.loss
+        loss.backward()
+
+        teacher = cache["teacher"]
+        fused = cache["fused"]
+        grad = None if teacher is None else teacher.grad
+        if teacher is None or fused is None or grad is None:
+            raise RuntimeError(
+                "Auto selection hooks failed to capture outputs/gradients. "
+                "Try updating PyTorch or run with --layer <index>."
+            )
+        if not teacher.requires_grad:
+            raise RuntimeError(
+                "Teacher hidden state does not require grad. "
+                "Ensure model parameters require grad for DWCE."
+            )
+
+        with torch.no_grad():
+            if attention_mask is not None:
+                mask = attention_mask.to(dtype=torch.float32).unsqueeze(-1)
+                batch_tokens = float(mask.sum().item())
+            else:
+                mask = None
+                batch_tokens = float(input_ids.numel())
+            token_count += batch_tokens
+
+            delta = fused - teacher
+            grad_sq = grad.pow(2)
+            if mask is not None:
+                delta = delta * mask
+                grad_sq = grad_sq * mask
+
+            score_num += (delta.pow(2) * grad_sq).sum().item()
+            score_den += (teacher.pow(2) * grad_sq).sum().item()
+            num_batches += 1
+
+    handle_a.remove()
+    handle_b.remove()
+
+    if norm == "relative":
+        score = score_num / (score_den + eps)
+    else:
+        denom = token_count if token_count > 0 else 1.0
+        score = score_num / denom
+
+    meta = {
+        "num_batches": num_batches,
+        "token_count": token_count,
+        "norm": norm,
+        "supports_kwargs": supports_kwargs,
+    }
+    return score, meta
+
+
+def _compute_cosine_for_pair(
+    model,
+    layer_a: torch.nn.Module,
+    layer_b: torch.nn.Module,
+    dataloader,
+    device: str,
+    max_batches: int,
+    eps: float,
+) -> Tuple[float, Dict[str, object]]:
+    cache = {"a": None, "b": None}
+    supports_kwargs = True
+
+    def hook_a(_module, _inputs, output, _kwargs=None):
+        hidden = _extract_hidden(output)
+        if hidden is None:
+            raise RuntimeError("Failed to extract layer_a hidden state output.")
+        cache["a"] = hidden
+        return output
+
+    def hook_b(_module, _inputs, output, _kwargs=None):
+        hidden = _extract_hidden(output)
+        if hidden is None:
+            raise RuntimeError("Failed to extract layer_b hidden state output.")
+        cache["b"] = hidden
+        return output
+
+    handle_a, has_kwargs_a = _register_forward_hook(layer_a, hook_a)
+    handle_b, has_kwargs_b = _register_forward_hook(layer_b, hook_b)
+    supports_kwargs = has_kwargs_a and has_kwargs_b
+
+    score_sum = 0.0
+    token_count = 0.0
+    num_batches = 0
+
+    model.eval()
+    for batch_idx, batch in enumerate(dataloader):
+        if max_batches and batch_idx >= max_batches:
+            break
+        cache["a"] = None
+        cache["b"] = None
+
+        input_ids = batch[0].to(device)
+        attention_mask = batch[1].to(device) if len(batch) > 1 else None
+
+        with torch.no_grad():
+            model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                use_cache=False,
+            )
+
+        hidden_a = cache["a"]
+        hidden_b = cache["b"]
+        if hidden_a is None or hidden_b is None:
+            raise RuntimeError(
+                "Auto selection hooks failed to capture outputs for cosine scoring."
+            )
+
+        with torch.no_grad():
+            a = hidden_a.float()
+            b = hidden_b.float()
+            cos = F.cosine_similarity(a, b, dim=-1, eps=eps)
+            distance = 1.0 - cos
+
+            if attention_mask is not None:
+                mask = attention_mask.to(dtype=torch.float32)
+                batch_tokens = float(mask.sum().item())
+                distance = distance * mask
+            else:
+                batch_tokens = float(distance.numel())
+
+            token_count += batch_tokens
+            score_sum += float(distance.sum().item())
+            num_batches += 1
+
+    handle_a.remove()
+    handle_b.remove()
+
+    denom = token_count if token_count > 0 else 1.0
+    score = score_sum / denom
+    meta = {
+        "num_batches": num_batches,
+        "token_count": token_count,
+        "metric": "cosine",
+        "supports_kwargs": supports_kwargs,
+    }
+    return score, meta
+
+
+def _compute_global_rel_change_for_pair(
+    model,
+    layers: List[torch.nn.Module],
+    pair_idx: int,
+    dataloader,
+    args,
+    max_batches: int,
+    eps: float,
+) -> Tuple[float, Dict[str, object]]:
+    hidden_size = _get_hidden_size(model)
+    head_permute_select = not bool(getattr(args, "no_head_permute_select", False))
+    layer_a = layers[pair_idx]
+    layer_b = layers[pair_idx + 1]
+    fused_layer, fuse_priors = _build_fused_layer_for_pair(
+        model,
+        layer_a,
+        layer_b,
+        dataloader,
+        device=args.device,
+        fisher_mode=args.fisher_mode,
+        eps=eps,
+        hidden_size=hidden_size,
+        enable_head_permute=head_permute_select,
+    )
+    fused_layer.to(args.device)
+    fused_layer.eval()
+
+    parent, name, container = find_layer_container(model, getattr(args, "layer_path", None))
+    if len(list(container)) != len(layers):
+        raise RuntimeError("Layer container changed during auto-selection; aborting rerank.")
+
+    virtual_layers = list(layers)
+    virtual_layers[pair_idx] = fused_layer
+    del virtual_layers[pair_idx + 1]
+    if isinstance(container, torch.nn.ModuleList):
+        virtual_container = torch.nn.ModuleList(virtual_layers)
+    elif isinstance(container, list):
+        virtual_container = virtual_layers
+    else:
+        raise TypeError("Layer container must be ModuleList or list")
+
+    teacher_cache = {"pair": None, "final": None}
+    supports_kwargs = True
+
+    def hook_pair(_module, _inputs, output, _kwargs=None):
+        hidden = _extract_hidden(output)
+        if hidden is None:
+            raise RuntimeError("Failed to extract pair output for global relation rerank.")
+        teacher_cache["pair"] = hidden
+        return output
+
+    handle_pair, has_kwargs_pair = _register_forward_hook(layer_b, hook_pair)
+    supports_kwargs = supports_kwargs and has_kwargs_pair
+
+    score_sum = 0.0
+    token_count = 0.0
+    num_batches = 0
+
+    model.eval()
+    for batch_idx, batch in enumerate(dataloader):
+        if max_batches and batch_idx >= max_batches:
+            break
+
+        teacher_cache["pair"] = None
+
+        input_ids = batch[0].to(args.device)
+        attention_mask = batch[1].to(args.device) if len(batch) > 1 else None
+
+        with torch.no_grad():
+            teacher_outputs = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+                use_cache=False,
+            )
+            teacher_hidden_states = getattr(teacher_outputs, "hidden_states", None)
+            if not teacher_hidden_states:
+                raise RuntimeError("Teacher forward did not return hidden_states.")
+            teacher_final = teacher_hidden_states[-1]
+            teacher_pair = teacher_cache["pair"]
+
+        if teacher_pair is None or teacher_final is None:
+            raise RuntimeError(
+                "Failed to capture teacher pair/final hidden states for global rerank."
+            )
+
+        with torch.no_grad(), _temporary_layers(parent, name, virtual_container):
+            fused_outputs = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+                use_cache=False,
+            )
+            fused_hidden_states = getattr(fused_outputs, "hidden_states", None)
+            if not fused_hidden_states:
+                raise RuntimeError("Fused forward did not return hidden_states.")
+            fused_final = fused_hidden_states[-1]
+
+        if fused_final is None:
+            raise RuntimeError("Failed to capture fused final hidden state for global rerank.")
+
+        with torch.no_grad():
+            teacher_pair_f = teacher_pair.float()
+            teacher_final_f = teacher_final.float()
+            fused_final_f = fused_final.float()
+
+            teacher_rel = F.cosine_similarity(
+                teacher_pair_f, teacher_final_f, dim=-1, eps=eps
+            )
+            fused_rel = F.cosine_similarity(
+                teacher_pair_f, fused_final_f, dim=-1, eps=eps
+            )
+            rel_change = (teacher_rel - fused_rel).abs()
+
+            if attention_mask is not None:
+                mask = attention_mask.to(dtype=torch.float32)
+                batch_tokens = float(mask.sum().item())
+                rel_change = rel_change * mask
+            else:
+                batch_tokens = float(rel_change.numel())
+
+            token_count += batch_tokens
+            score_sum += float(rel_change.sum().item())
+            num_batches += 1
+
+    handle_pair.remove()
+    del fused_layer
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    denom = token_count if token_count > 0 else 1.0
+    score = score_sum / denom
+    meta = {
+        "num_batches": num_batches,
+        "token_count": token_count,
+        "metric": "global_rel_change",
+        "supports_kwargs": supports_kwargs,
+        "fuse_priors": fuse_priors,
+    }
+    return score, meta
+
+
+def select_layer_auto(
+    model,
+    layers: List[torch.nn.Module],
+    dataloader,
+    args,
+    previous_scores: Optional[List[float]] = None,
+    start_index: int = 0,
+    exclude_pairs: Optional[Set[int]] = None,
+) -> Tuple[int, List[float], Dict[str, object]]:
+    num_layers = len(layers)
+    if num_layers < 2:
+        raise SystemExit("Model must have at least 2 layers for auto selection.")
+
+    hidden_size = _get_hidden_size(model)
+    num_pairs = num_layers - 1
+    scores: List[float] = [float("inf")] * num_pairs
+    meta_per_pair: List[Optional[Dict[str, object]]] = [None] * num_pairs
+    supports_kwargs_all = True
+    head_permute_select = not bool(getattr(args, "no_head_permute_select", False))
+    exclude_set: Set[int] = {
+        int(idx)
+        for idx in (exclude_pairs or set())
+        if isinstance(idx, int) and 0 <= int(idx) < num_pairs
+    }
+
+    max_batches = args.auto_max_batches
+    start_index = max(0, min(start_index, num_pairs))
+    auto_metric = str(getattr(args, "auto_metric", "dwce")).strip().lower()
+    if auto_metric == "hybrid":
+        auto_metric = "hybrid_cosine"
+    if auto_metric not in {
+        "dwce",
+        "cosine",
+        "hybrid_cosine",
+        "hybrid_global_rel",
+    }:
+        raise SystemExit(
+            "--auto_metric must be one of: dwce, cosine, hybrid, "
+            "hybrid_cosine, hybrid_global_rel"
+        )
+    auto_cosine_topk = int(getattr(args, "auto_cosine_topk", 3))
+    if auto_cosine_topk <= 0:
+        raise SystemExit("--auto_cosine_topk must be >= 1")
+    print(
+        f"[auto] metric={auto_metric}; using "
+        f"{('all' if max_batches == 0 else max_batches)} batches "
+        "from calibration samples."
+    )
+
+    reuse_upto = 0
+    allow_reuse = auto_metric == "dwce"
+    if previous_scores:
+        reuse_upto = min(start_index, len(previous_scores), num_pairs) if allow_reuse else 0
+        for idx in range(reuse_upto):
+            if idx in exclude_set:
+                scores[idx] = float("inf")
+                meta_per_pair[idx] = {"excluded": True}
+                print(f"[auto] skipped excluded pair {idx}-{idx+1}.")
+                continue
+            scores[idx] = previous_scores[idx]
+            meta_per_pair[idx] = (
+                {
+                    "num_batches": 0,
+                    "token_count": 0.0,
+                    "norm": args.auto_norm,
+                    "metric": auto_metric,
+                    "supports_kwargs": True,
+                    "reused": True,
+                }
+            )
+            print(f"[auto] reused pair {idx}-{idx+1}: {scores[idx]:.6e}")
+
+    compute_start = start_index if reuse_upto == start_index else reuse_upto
+    pairs_to_score: List[int] = []
+    for idx in range(compute_start, num_pairs):
+        if idx in exclude_set:
+            scores[idx] = float("inf")
+            meta_per_pair[idx] = {"excluded": True}
+            print(f"[auto] skipped excluded pair {idx}-{idx+1}.")
+            continue
+        pairs_to_score.append(idx)
+
+    def _score_dwce_for_pair(idx: int) -> Tuple[float, Dict[str, object]]:
+        print(f"[auto] building fused pair {idx}-{idx+1} for DWCE...")
+        layer_a = layers[idx]
+        layer_b = layers[idx + 1]
+        dwce_mode = str(getattr(args, "auto_dwce_mode", "separate")).strip().lower()
+        if dwce_mode == "shared":
+            try:
+                return _score_dwce_with_shared_backward(
+                    model,
+                    layer_a,
+                    layer_b,
+                    dataloader,
+                    device=args.device,
+                    fisher_mode=args.fisher_mode,
+                    max_batches=max_batches,
+                    eps=args.eps,
+                    norm=args.auto_norm,
+                    hidden_size=hidden_size,
+                    enable_head_permute=head_permute_select,
+                )
+            except _DwceGradCacheOverflow:
+                print(
+                    "[auto] shared-backward DWCE cache exceeded budget; "
+                    "falling back to separate mode."
+                )
+        fused, fuse_priors = _build_fused_layer_for_pair(
+            model,
+            layer_a,
+            layer_b,
+            dataloader,
+            device=args.device,
+            fisher_mode=args.fisher_mode,
+            eps=args.eps,
+            hidden_size=hidden_size,
+            enable_head_permute=head_permute_select,
+        )
+        fused.to(args.device)
+        fused.eval()
+        for param in model.parameters():
+            param.requires_grad_(True)
+        score, meta = _compute_dwce_for_pair(
+            model,
+            layer_a,
+            layer_b,
+            fused,
+            dataloader,
+            device=args.device,
+            max_batches=max_batches,
+            eps=args.eps,
+            norm=args.auto_norm,
+        )
+        meta["fuse_priors"] = fuse_priors
+        meta["metric"] = "dwce"
+        del fused
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return score, meta
+
+    def _score_cosine_for_pair(idx: int) -> Tuple[float, Dict[str, object]]:
+        print(f"[auto] scoring cosine for pair {idx}-{idx+1}...")
+        layer_a = layers[idx]
+        layer_b = layers[idx + 1]
+        return _compute_cosine_for_pair(
+            model,
+            layer_a,
+            layer_b,
+            dataloader,
+            device=args.device,
+            max_batches=max_batches,
+            eps=args.eps,
+        )
+
+    def _score_global_rel_for_pair(idx: int) -> Tuple[float, Dict[str, object]]:
+        print(f"[auto] scoring global relation change for pair {idx}-{idx+1}...")
+        return _compute_global_rel_change_for_pair(
+            model,
+            layers,
+            idx,
+            dataloader,
+            args=args,
+            max_batches=max_batches,
+            eps=args.eps,
+        )
+
+    if auto_metric in {"dwce", "cosine"}:
+        for idx in pairs_to_score:
+            if auto_metric == "dwce":
+                score, meta = _score_dwce_for_pair(idx)
+            else:
+                score, meta = _score_cosine_for_pair(idx)
+            supports_kwargs_all = supports_kwargs_all and meta.get("supports_kwargs", True)
+            scores[idx] = score
+            meta_per_pair[idx] = meta
+            print(f"[auto] {auto_metric} pair {idx}-{idx+1}: {score:.6e}")
+    else:
+        dwce_prefilter: Dict[int, float] = {}
+        for idx in pairs_to_score:
+            score, meta = _score_dwce_for_pair(idx)
+            dwce_prefilter[idx] = score
+            supports_kwargs_all = supports_kwargs_all and meta.get("supports_kwargs", True)
+            meta_per_pair[idx] = {
+                "prefilter_dwce": score,
+                "dwce_meta": meta,
+                "metric": "hybrid",
+            }
+            print(f"[auto] hybrid prefilter DWCE pair {idx}-{idx+1}: {score:.6e}")
+        ranked = sorted(pairs_to_score, key=lambda i: float(dwce_prefilter[i]))
+        shortlist = ranked[: min(auto_cosine_topk, len(ranked))]
+        print(f"[auto] hybrid shortlist (dwce top-{len(shortlist)}): {shortlist}")
+        for idx in shortlist:
+            if auto_metric == "hybrid_global_rel":
+                score, rerank_meta = _score_global_rel_for_pair(idx)
+                score_metric = "global_rel_change"
+            else:
+                score, rerank_meta = _score_cosine_for_pair(idx)
+                score_metric = "cosine"
+            supports_kwargs_all = supports_kwargs_all and rerank_meta.get(
+                "supports_kwargs", True
+            )
+            scores[idx] = score
+            pair_meta = meta_per_pair[idx] or {}
+            pair_meta["rerank_meta"] = rerank_meta
+            pair_meta["score_metric"] = score_metric
+            meta_per_pair[idx] = pair_meta
+            print(f"[auto] hybrid {score_metric} pair {idx}-{idx+1}: {score:.6e}")
+
+    if not supports_kwargs_all:
+        print(
+            "[auto] Warning: forward hooks did not capture kwargs; "
+            "fused-layer calls may be approximate."
+        )
+
+    print(f"[auto] score summary (metric={auto_metric}, norm={args.auto_norm}):")
+    for idx, score in enumerate(scores):
+        if idx in exclude_set:
+            print(f"[auto]   pair {idx}-{idx+1}: excluded")
+        elif math.isfinite(float(score)):
+            print(f"[auto]   pair {idx}-{idx+1}: {score:.6e}")
+        else:
+            print(f"[auto]   pair {idx}-{idx+1}: {score}")
+
+    candidates = [i for i in range(num_pairs) if i not in exclude_set]
+    if not candidates:
+        raise SystemExit("All pairs are excluded; cannot auto-select a fusion layer.")
+    best_idx = min(candidates, key=lambda i: scores[i])
+    best_score = float(scores[best_idx])
+    if not math.isfinite(best_score):
+        raise SystemExit(
+            "Auto selection failed: all candidate pairs have non-finite scores "
+            "(check --exclude_pairs and data)."
+        )
+    print(f"[auto] Selected layer {best_idx} (score={best_score:.6e})")
+
+    meta = {
+        "per_pair": meta_per_pair,
+        "supports_kwargs": supports_kwargs_all,
+        "max_batches": max_batches,
+        "norm": args.auto_norm,
+        "metric": auto_metric,
+        "cosine_topk": auto_cosine_topk,
+        "start_index": start_index,
+        "excluded_pairs": sorted(exclude_set),
+    }
+    return best_idx, scores, meta

src/loratune.py

@@ -0,0 +1,430 @@
+#!/usr/bin/env python3
+"""Centralized Alpaca LoRA finetuning for post-pruned models."""
+
+import argparse
+import itertools
+import json
+import os
+from types import SimpleNamespace
+from pathlib import Path
+
+import torch
+from contextlib import nullcontext
+
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, PretrainedConfig
+from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+import ppl_eval
+
+from fuse_layers_data import FixedSeqDataset, load_instruction_records
+from fuse_layers_distill import LoRALinear, apply_lora_adapters, merge_lora_adapters
+
+try:
+    from tqdm import tqdm
+except Exception:  # pragma: no cover
+    tqdm = None
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run centralized Alpaca LoRA finetuning.")
+    parser.add_argument("--base_model", required=True, help="Path or HF model id to finetune")
+    parser.add_argument("--output_dir", required=True, help="Directory to save merged model")
+    parser.add_argument("--device", default="cuda", help="Training device")
+    parser.add_argument(
+        "--dtype",
+        default="bfloat16",
+        choices=["float32", "float16", "bfloat16"],
+        help="Model load/training dtype",
+    )
+    parser.add_argument("--trust_remote_code", action="store_true", help="Enable trust_remote_code")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+
+    parser.add_argument(
+        "--instruction_dataset",
+        default="yahma/alpaca-cleaned",
+        help="HF dataset name for Alpaca-style instruction data",
+    )
+    parser.add_argument("--instruction_config", default=None, help="Optional dataset config")
+    parser.add_argument("--instruction_split", default="train", help="Dataset split")
+    parser.add_argument("--instruction_field_instruction", default="instruction")
+    parser.add_argument("--instruction_field_input", default="input")
+    parser.add_argument("--instruction_field_output", default="output")
+    parser.add_argument("--max_samples", type=int, default=0, help="Limit instruction samples (0 = all)")
+    parser.add_argument("--seq_len", type=int, default=1024, help="Training sequence length")
+    parser.add_argument("--batch_size", type=int, default=64, help="Global batch size")
+    parser.add_argument("--micro_batch_size", type=int, default=4, help="Per-step micro-batch size")
+    parser.add_argument("--epochs", type=float, default=1.0, help="Training epochs")
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate")
+    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay")
+    parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Gradient clipping norm")
+    parser.add_argument("--log_steps", type=int, default=100, help="Log every N optimizer steps")
+    parser.add_argument(
+        "--save_steps",
+        type=int,
+        default=200,
+        help="Save LoRA adapter checkpoints every N optimizer steps (0 = disable)",
+    )
+    parser.add_argument(
+        "--no_wikitext2_ppl_on_log",
+        dest="wikitext2_ppl_on_log",
+        action="store_false",
+        help="Disable Wikitext-2 perplexity evaluation at loss log steps",
+    )
+    parser.set_defaults(wikitext2_ppl_on_log=True)
+    parser.add_argument("--wikitext2_ppl_seq_len", type=int, default=128)
+    parser.add_argument("--wikitext2_ppl_batch_size", type=int, default=8)
+    parser.add_argument("--wikitext2_ppl_max_batches", type=int, default=None)
+
+    parser.add_argument("--lora_rank", type=int, default=8, help="LoRA rank")
+    parser.add_argument("--lora_alpha", type=float, default=16.0, help="LoRA alpha")
+    parser.add_argument("--lora_dropout", type=float, default=0.0, help="LoRA dropout")
+    parser.add_argument(
+        "--lora_target_modules",
+        nargs="*",
+        default=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "down_proj", "up_proj"],
+        help="Linear module suffixes to LoRA-wrap",
+    )
+
+    return parser.parse_args()
+
+
+def get_dtype(name: str) -> torch.dtype:
+    return {
+        "float32": torch.float32,
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+    }[name]
+
+
+def seed_all(seed: int) -> None:
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def normalize_config(config):
+    layer_types = getattr(config, "layer_types", None)
+    num_hidden_layers = getattr(config, "num_hidden_layers", None)
+    if layer_types is not None and num_hidden_layers is not None and len(layer_types) != num_hidden_layers:
+        config.layer_types = list(layer_types[:num_hidden_layers])
+    if getattr(config, "_attn_implementation", None) is None:
+        config._attn_implementation = "eager"
+    return config
+
+
+def load_normalized_config(base_model: str, trust_remote_code: bool):
+    config_dict, unused_kwargs = PretrainedConfig.get_config_dict(base_model, trust_remote_code=trust_remote_code)
+    layer_types = config_dict.get("layer_types")
+    num_hidden_layers = config_dict.get("num_hidden_layers")
+    if layer_types is not None and num_hidden_layers is not None and len(layer_types) != num_hidden_layers:
+        config_dict["layer_types"] = list(layer_types[:num_hidden_layers])
+    if config_dict.get("_attn_implementation") is None:
+        config_dict["_attn_implementation"] = "eager"
+    model_type = config_dict["model_type"]
+    config_class = CONFIG_MAPPING[model_type]
+    config = config_class.from_dict(config_dict, **unused_kwargs)
+    return normalize_config(config)
+
+
+def validate_local_model_dir(base_path: Path) -> None:
+    if not base_path.exists() or not base_path.is_dir():
+        return
+
+    has_config = (base_path / "config.json").is_file()
+    has_weights = any(
+        (base_path / name).is_file()
+        for name in (
+            "model.safetensors",
+            "model.safetensors.index.json",
+            "pytorch_model.bin",
+            "pytorch_model.bin.index.json",
+        )
+    )
+    if has_config and has_weights:
+        return
+
+    raise SystemExit(
+        "Local --base_model points to an incomplete HF model directory: "
+        f"{base_path}. Expected at least config.json and model weights. "
+        "Set --base_model/BASE_MODEL to a saved HF model directory."
+    )
+
+
+def load_base_artifacts(args: argparse.Namespace):
+    base_path = Path(args.base_model)
+    if base_path.is_file() and base_path.suffix == ".bin":
+        checkpoint = torch.load(str(base_path), map_location="cpu", weights_only=False)
+        if not isinstance(checkpoint, dict) or "model" not in checkpoint or "tokenizer" not in checkpoint:
+            raise SystemExit("Expected a .bin checkpoint dict with `model` and `tokenizer` entries.")
+        model = checkpoint["model"]
+        tokenizer = checkpoint["tokenizer"]
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token
+        return model, tokenizer
+
+    validate_local_model_dir(base_path)
+    tokenizer = AutoTokenizer.from_pretrained(args.base_model, trust_remote_code=args.trust_remote_code)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token
+    config = load_normalized_config(args.base_model, trust_remote_code=args.trust_remote_code)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.base_model,
+        config=config,
+        torch_dtype=get_dtype(args.dtype),
+        trust_remote_code=args.trust_remote_code,
+    )
+    return model, tokenizer
+
+
+def build_training_loader(args: argparse.Namespace, tokenizer) -> torch.utils.data.DataLoader:
+    num_samples = args.max_samples if args.max_samples > 0 else 0
+    records = load_instruction_records(args, num_samples)
+    if not records:
+        raise SystemExit("No instruction records were loaded.")
+    dataset = FixedSeqDataset(records, tokenizer, args.seq_len)
+    return torch.utils.data.DataLoader(dataset, batch_size=args.micro_batch_size, shuffle=True)
+
+
+def save_lora_adapters(
+    model: torch.nn.Module, args: argparse.Namespace, subdir: str = "lora_adapter"
+) -> str:
+    adapter_dir = os.path.join(args.output_dir, subdir)
+    os.makedirs(adapter_dir, exist_ok=True)
+
+    adapter_state = {}
+    adapter_modules = {}
+    for module_name, module in model.named_modules():
+        if not isinstance(module, LoRALinear):
+            continue
+        adapter_modules[module_name] = {
+            "rank": module.rank,
+            "alpha": module.alpha,
+            "scaling": module.scaling,
+            "dropout": getattr(module.dropout, "p", 0.0),
+            "base_layer_class": type(module.base).__name__,
+            "in_features": module.base.in_features,
+            "out_features": module.base.out_features,
+        }
+        adapter_state[f"{module_name}.lora_A.weight"] = module.lora_A.weight.detach().cpu()
+        adapter_state[f"{module_name}.lora_B.weight"] = module.lora_B.weight.detach().cpu()
+
+    torch.save(adapter_state, os.path.join(adapter_dir, "adapter_model.bin"))
+    with open(os.path.join(adapter_dir, "adapter_config.json"), "w", encoding="utf-8") as handle:
+        json.dump(
+            {
+                "base_model": args.base_model,
+                "lora_rank": args.lora_rank,
+                "lora_alpha": args.lora_alpha,
+                "lora_dropout": args.lora_dropout,
+                "lora_target_modules": list(args.lora_target_modules),
+                "batch_size": args.batch_size,
+                "micro_batch_size": args.micro_batch_size,
+                "grad_accum_steps": args.grad_accum_steps,
+                "modules": adapter_modules,
+            },
+            handle,
+            indent=2,
+        )
+    return adapter_dir
+
+
+def prepare_wikitext2_eval(args: argparse.Namespace, model, tokenizer):
+    if not args.wikitext2_ppl_on_log:
+        return None
+    return ppl_eval.prepare_ppl_dataloaders(
+        tokenizer=tokenizer,
+        datasets=["wikitext"],
+        configs=["wikitext-2-raw-v1"],
+        split="test",
+        text_field=None,
+        num_samples=0,
+        seq_len=args.wikitext2_ppl_seq_len,
+        batch_size=args.wikitext2_ppl_batch_size,
+        seed=args.seed,
+        shuffle=False,
+        model_family="auto",
+        add_bos="auto",
+        cache_dir=None,
+        num_workers=0,
+        model=model,
+    )
+
+
+def train(model: torch.nn.Module, dataloader, args: argparse.Namespace, wikitext2_eval_dataloaders=None) -> dict:
+    lora_args = SimpleNamespace(
+        lora_rank=args.lora_rank,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        lora_target_modules=args.lora_target_modules,
+        lora_respect_exclude_pairs=False,
+        layer_path=None,
+        exclude_pairs=None,
+    )
+    lora_modules = apply_lora_adapters(model, lora_args)
+    lora_params = [param for module in lora_modules for param in module.lora_parameters()]
+
+    optimizer = torch.optim.AdamW(
+        lora_params,
+        lr=args.learning_rate,
+        weight_decay=args.weight_decay,
+    )
+    model.train()
+
+    device = torch.device(args.device)
+    device_type = device.type
+    amp_dtype = None
+    if args.dtype == "float16":
+        amp_dtype = torch.float16
+    elif args.dtype == "bfloat16":
+        amp_dtype = torch.bfloat16
+    use_amp = amp_dtype is not None and device_type == "cuda"
+    use_scaler = use_amp and amp_dtype == torch.float16
+    scaler = torch.cuda.amp.GradScaler() if use_scaler else None
+
+    full_epochs = int(args.epochs)
+    fractional = args.epochs - full_epochs
+    epoch_plan = [None] * full_epochs
+    if fractional > 1e-8:
+        frac_batches = max(1, int(round(fractional * len(dataloader))))
+        epoch_plan.append(frac_batches)
+
+    optimizer.zero_grad(set_to_none=True)
+    optimizer_step = 0
+    seen_batches = 0
+    last_loss = None
+    ppl_history = []
+
+    for epoch_idx, max_batches in enumerate(epoch_plan, start=1):
+        iterator = dataloader if max_batches is None else itertools.islice(dataloader, max_batches)
+        if tqdm is not None:
+            iterator = tqdm(iterator, desc=f"LoRA epoch {epoch_idx}", unit="batch", total=max_batches)
+        for batch in iterator:
+            input_ids = batch[0].to(args.device)
+            attention_mask = batch[1].to(args.device)
+
+            autocast_ctx = (
+                torch.autocast(device_type=device_type, dtype=amp_dtype)
+                if use_amp
+                else nullcontext()
+            )
+            with autocast_ctx:
+                outputs = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
+                logits = outputs.logits[:, :-1, :].contiguous()
+                labels = input_ids[:, 1:].contiguous()
+                mask = attention_mask[:, 1:].contiguous()
+                ce_flat = torch.nn.functional.cross_entropy(
+                    logits.view(-1, logits.size(-1)),
+                    labels.view(-1),
+                    reduction="none",
+                )
+                denom = mask.sum()
+                if denom.item() == 0:
+                    continue
+                loss = (ce_flat * mask.reshape(-1).to(ce_flat.dtype)).sum() / denom
+
+            last_loss = float(loss.detach().item())
+            scaled_loss = loss / max(args.grad_accum_steps, 1)
+            if use_scaler:
+                scaler.scale(scaled_loss).backward()
+            else:
+                scaled_loss.backward()
+
+            seen_batches += 1
+            if seen_batches % max(args.grad_accum_steps, 1) != 0:
+                continue
+
+            if args.max_grad_norm is not None:
+                if use_scaler:
+                    scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(lora_params, args.max_grad_norm)
+            if use_scaler:
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                optimizer.step()
+            optimizer.zero_grad(set_to_none=True)
+            optimizer_step += 1
+
+            if args.log_steps and optimizer_step % args.log_steps == 0:
+                print(f"[loratune] step={optimizer_step} loss={last_loss:.6f}")
+                if wikitext2_eval_dataloaders is not None:
+                    prev_mode = model.training
+                    model.eval()
+                    ppl_results = ppl_eval.evaluate_ppl_dataloaders(
+                        model,
+                        wikitext2_eval_dataloaders,
+                        args.device,
+                        max_batches=args.wikitext2_ppl_max_batches,
+                    )
+                    ppl_history.append({"step": optimizer_step, "ppl": ppl_results})
+                    print(f"[loratune] ppl step={optimizer_step} {ppl_results}")
+                    if prev_mode:
+                        model.train()
+
+            if args.save_steps and optimizer_step % args.save_steps == 0:
+                checkpoint_dir = save_lora_adapters(
+                    model,
+                    args,
+                    subdir=os.path.join("checkpoints", f"step_{optimizer_step}"),
+                )
+                print(f"[loratune] saved adapter checkpoint to {checkpoint_dir}")
+
+    adapter_dir = save_lora_adapters(model, args)
+    merge_lora_adapters(model)
+    return {
+        "adapter_dir": adapter_dir,
+        "optimizer_steps": optimizer_step,
+        "seen_batches": seen_batches,
+        "last_loss": last_loss,
+        "wikitext2_ppl_history": ppl_history,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    if args.batch_size < 1:
+        raise SystemExit("--batch_size must be >= 1")
+    if args.micro_batch_size < 1:
+        raise SystemExit("--micro_batch_size must be >= 1")
+    args.grad_accum_steps = args.batch_size // args.micro_batch_size
+    if args.grad_accum_steps < 1:
+        raise SystemExit("--batch_size must be >= --micro_batch_size")
+
+    seed_all(args.seed)
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    model, tokenizer = load_base_artifacts(args)
+    if args.dtype != "float32":
+        model = model.to(get_dtype(args.dtype))
+    model.to(args.device)
+
+    dataloader = build_training_loader(args, tokenizer)
+    wikitext2_eval_dataloaders = prepare_wikitext2_eval(args, model, tokenizer)
+    metrics = train(model, dataloader, args, wikitext2_eval_dataloaders=wikitext2_eval_dataloaders)
+
+    model.save_pretrained(args.output_dir)
+    tokenizer.save_pretrained(args.output_dir)
+
+    with open(os.path.join(args.output_dir, "loratune_metrics.json"), "w", encoding="utf-8") as handle:
+        json.dump(
+            {
+                "base_model": args.base_model,
+                "instruction_dataset": args.instruction_dataset,
+                "seq_len": args.seq_len,
+                "batch_size": args.batch_size,
+                "micro_batch_size": args.micro_batch_size,
+                "grad_accum_steps": args.grad_accum_steps,
+                "epochs": args.epochs,
+                "learning_rate": args.learning_rate,
+                "save_steps": args.save_steps,
+                "lora_rank": args.lora_rank,
+                "lora_alpha": args.lora_alpha,
+                "lora_dropout": args.lora_dropout,
+                **metrics,
+            },
+            handle,
+            indent=2,
+        )
+
+
+if __name__ == "__main__":
+    main()

src/loratune_config.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+"""Configuration helpers for centralized LoRA finetuning."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from types import SimpleNamespace
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class LoRATuneConfig:
+    """Structured config matching the current loratune.py CLI surface."""
+
+    base_model: str = ""
+    output_dir: str = ""
+    device: str = "cuda"
+    dtype: str = "bfloat16"
+    trust_remote_code: bool = False
+    seed: int = 42
+
+    instruction_dataset: str = "tatsu-lab/alpaca"
+    instruction_config: Optional[str] = None
+    instruction_split: str = "train"
+    instruction_field_instruction: str = "instruction"
+    instruction_field_input: str = "input"
+    instruction_field_output: str = "output"
+    max_samples: int = 0
+    seq_len: int = 1024
+    batch_size: int = 64
+    micro_batch_size: int = 4
+    epochs: float = 1.0
+    learning_rate: float = 1e-4
+    weight_decay: float = 0.0
+    max_grad_norm: float = 1.0
+    log_steps: int = 100
+
+    wikitext2_ppl_on_log: bool = True
+    wikitext2_ppl_seq_len: int = 128
+    wikitext2_ppl_batch_size: int = 8
+    wikitext2_ppl_max_batches: Optional[int] = None
+
+    lora_rank: int = 8
+    lora_alpha: float = 16.0
+    lora_dropout: float = 0.0
+    lora_target_modules: List[str] = field(
+        default_factory=lambda: [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "down_proj",
+            "up_proj",
+        ]
+    )
+
+    @property
+    def grad_accum_steps(self) -> int:
+        if self.batch_size < 1:
+            raise ValueError("batch_size must be >= 1")
+        if self.micro_batch_size < 1:
+            raise ValueError("micro_batch_size must be >= 1")
+        if self.batch_size < self.micro_batch_size:
+            raise ValueError("batch_size must be >= micro_batch_size")
+        return self.batch_size // self.micro_batch_size
+
+    def validate(self) -> "LoRATuneConfig":
+        _ = self.grad_accum_steps
+        if not self.base_model:
+            raise ValueError("base_model must be set")
+        if not self.output_dir:
+            raise ValueError("output_dir must be set")
+        return self
+
+    def to_dict(self) -> Dict[str, Any]:
+        data = asdict(self)
+        data["grad_accum_steps"] = self.grad_accum_steps
+        return data
+
+    def to_namespace(self) -> SimpleNamespace:
+        return SimpleNamespace(**self.to_dict())
+
+    @classmethod
+    def from_dict(cls, values: Dict[str, Any]) -> "LoRATuneConfig":
+        return cls(**values)

src/ppl_eval.py

@@ -0,0 +1,576 @@
+#!/usr/bin/env python3
+"""Perplexity evaluation for causal LMs on HF datasets or provided text."""
+
+import argparse
+import json
+import math
+import os
+from typing import Dict, Iterable, List, Optional
+
+import torch
+
+try:
+    from datasets import load_dataset
+except Exception:  # pragma: no cover - optional dependency
+    load_dataset = None
+
+try:
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+except Exception as exc:  # pragma: no cover - fail early with clear error
+    raise SystemExit("transformers is required: pip install transformers") from exc
+
+try:
+    from tqdm import tqdm
+except Exception:  # pragma: no cover - optional dependency
+    tqdm = None
+
+
+def _tqdm_enabled() -> bool:
+    value = os.environ.get("DISABLE_TQDM", os.environ.get("TQDM_DISABLE", "0"))
+    return value.strip().lower() not in {"1", "true", "yes", "on"}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Compute perplexity for a causal LM on one or more datasets."
+    )
+    parser.add_argument("--model", required=True, help="HF model id or local path")
+    parser.add_argument(
+        "--dataset",
+        action="append",
+        default=[],
+        help="HF dataset name (repeatable).",
+    )
+    parser.add_argument(
+        "--dataset_config",
+        action="append",
+        default=[],
+        help="Optional dataset config (repeatable or single shared config).",
+    )
+    parser.add_argument(
+        "--dataset_split",
+        default="test",
+        help="Dataset split to use (default: test)",
+    )
+    parser.add_argument(
+        "--dataset_text_field",
+        default=None,
+        help="Text field in dataset (default: auto-detect, applies to all datasets)",
+    )
+    parser.add_argument(
+        "--text",
+        action="append",
+        default=[],
+        help="Inline text samples (can pass multiple)",
+    )
+    parser.add_argument(
+        "--text_file",
+        default=None,
+        help="Path to a text file for evaluation data",
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=0,
+        help="Number of token sequences to use per dataset (0 = all)",
+    )
+    parser.add_argument(
+        "--seq_len", type=int, default=2048, help="Sequence length"
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=2, help="Batch size"
+    )
+    parser.add_argument(
+        "--max_batches",
+        type=int,
+        default=None,
+        help="Optional max number of batches to evaluate per dataset",
+    )
+    parser.add_argument(
+        "--model_family",
+        type=str,
+        choices=["auto", "llama", "qwen"],
+        default="auto",
+        help="Model family for BOS handling",
+    )
+    parser.add_argument(
+        "--add_bos",
+        type=str,
+        choices=["auto", "always", "never"],
+        default="auto",
+        help="Whether to prepend BOS to each sample",
+    )
+    parser.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device for model + compute",
+    )
+    parser.add_argument(
+        "--dtype",
+        default="auto",
+        choices=["auto", "float32", "float16", "bfloat16"],
+        help="Model dtype",
+    )
+    parser.add_argument(
+        "--seed", type=int, default=0, help="Random seed for shuffling"
+    )
+    parser.add_argument(
+        "--shuffle",
+        action="store_true",
+        help="Shuffle dataset before sampling",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=0,
+        help="DataLoader workers",
+    )
+    parser.add_argument(
+        "--cache_dir",
+        default=None,
+        help="Optional datasets cache directory",
+    )
+    parser.add_argument(
+        "--trust_remote_code",
+        action="store_true",
+        help="Allow custom model code from hub",
+    )
+    parser.add_argument(
+        "--output",
+        default=None,
+        help="Optional JSON output path",
+    )
+    return parser.parse_args()
+
+
+def _normalize_config(config: Optional[str]) -> Optional[str]:
+    if config is None:
+        return None
+    if config.strip().lower() in {"none", "null", "-"}:
+        return None
+    return config
+
+
+def _expand_dataset_configs(
+    datasets: List[str], configs: List[str]
+) -> List[Optional[str]]:
+    if not configs:
+        return [None] * len(datasets)
+    if len(configs) == 1 and len(datasets) > 1:
+        return [_normalize_config(configs[0])] * len(datasets)
+    if len(configs) != len(datasets):
+        raise SystemExit(
+            "Provide zero, one, or matching-count --dataset_config values."
+        )
+    return [_normalize_config(cfg) for cfg in configs]
+
+
+def guess_text_field(dataset) -> str:
+    if hasattr(dataset, "column_names") and dataset.column_names:
+        if "text" in dataset.column_names:
+            return "text"
+        return dataset.column_names[0]
+    if hasattr(dataset, "features"):
+        names = list(dataset.features.keys())
+        if "text" in names:
+            return "text"
+        if names:
+            return names[0]
+    return "text"
+
+
+def _infer_model_family(model) -> str:
+    model_type = str(getattr(getattr(model, "config", None), "model_type", "")).lower()
+    architectures = getattr(getattr(model, "config", None), "architectures", [])
+    arch_lower = " ".join(str(name).lower() for name in architectures)
+    if "qwen" in model_type or "qwen" in arch_lower:
+        return "qwen"
+    if "llama" in model_type or "llama" in arch_lower:
+        return "llama"
+    return "unknown"
+
+
+def _resolve_add_bos(setting: str, model_family: str, tokenizer) -> bool:
+    if setting == "always":
+        return True
+    if setting == "never":
+        return False
+    if model_family == "llama":
+        return True
+    if model_family == "qwen":
+        return False
+    if hasattr(tokenizer, "add_bos_token"):
+        return bool(getattr(tokenizer, "add_bos_token"))
+    init_kwargs = getattr(tokenizer, "init_kwargs", None)
+    if isinstance(init_kwargs, dict) and "add_bos_token" in init_kwargs:
+        return bool(init_kwargs["add_bos_token"])
+    return False
+
+
+def build_token_chunks(
+    texts: Iterable[str],
+    tokenizer,
+    seq_len: int,
+    num_samples: int,
+    add_bos: bool = False,
+) -> List[torch.Tensor]:
+    chunks: List[torch.Tensor] = []
+    buffer: List[int] = []
+    for text in texts:
+        ids = tokenizer.encode(text, add_special_tokens=False)
+        if add_bos and tokenizer.bos_token_id is not None:
+            ids = [tokenizer.bos_token_id] + ids
+        if not ids:
+            continue
+        buffer.extend(ids)
+        while len(buffer) >= seq_len and len(chunks) < num_samples:
+            chunk = buffer[:seq_len]
+            buffer = buffer[seq_len:]
+            chunks.append(torch.tensor(chunk, dtype=torch.long))
+        if len(chunks) >= num_samples:
+            break
+    return chunks
+
+
+def get_dtype(dtype: str):
+    if dtype == "auto":
+        return None
+    if dtype == "float16":
+        return torch.float16
+    if dtype == "bfloat16":
+        return torch.bfloat16
+    return torch.float32
+
+
+def compute_ppl(model, dataloader, device: str, max_batches: Optional[int]) -> float:
+    model.eval()
+    nll_sum = 0.0
+    token_count = 0
+    iterator = dataloader
+    if tqdm is not None and _tqdm_enabled():
+        iterator = tqdm(dataloader, desc="PPL", unit="batch")
+    with torch.no_grad():
+        for step, batch in enumerate(iterator):
+            if isinstance(batch, dict):
+                input_ids = batch["input_ids"].to(device)
+            else:
+                input_ids = batch[0].to(device)
+            outputs = model(input_ids=input_ids)
+            logits = outputs.logits
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = input_ids[:, 1:].contiguous()
+            loss = torch.nn.functional.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                reduction="sum",
+            )
+            nll_sum += float(loss.item())
+            token_count += shift_labels.numel()
+            if max_batches is not None and step + 1 >= max_batches:
+                break
+
+    if token_count == 0:
+        raise RuntimeError("No tokens processed; check evaluation inputs.")
+
+    return math.exp(nll_sum / token_count)
+
+
+def _load_lm_dataset(
+    tokenizer,
+    dataset_name: str,
+    config: Optional[str],
+    split: str,
+    text_field: Optional[str],
+    seq_len: int,
+    add_bos: bool,
+    cache_dir: Optional[str],
+):
+    dataset = load_dataset(
+        dataset_name,
+        config,
+        split=split,
+        trust_remote_code=True,
+        cache_dir=cache_dir,
+    )
+
+    field = text_field or guess_text_field(dataset)
+
+    def is_valid_text(example) -> bool:
+        value = example.get(field)
+        return isinstance(value, str) and value.strip() != ""
+
+    dataset = dataset.filter(is_valid_text, desc=f"filter-{dataset_name}")
+
+    def tokenize_fn(examples):
+        tokenized = tokenizer(
+            examples[field],
+            add_special_tokens=False,
+            return_attention_mask=False,
+        )
+        if add_bos and tokenizer.bos_token_id is not None:
+            tokenized["input_ids"] = [
+                [tokenizer.bos_token_id] + ids for ids in tokenized["input_ids"]
+            ]
+        return tokenized
+
+    tokenized = dataset.map(
+        tokenize_fn,
+        batched=True,
+        remove_columns=dataset.column_names,
+        desc=f"tokenize-{dataset_name}",
+    )
+
+    def group_texts(examples):
+        concatenated = []
+        for ids in examples["input_ids"]:
+            concatenated.extend(ids)
+        total_length = (len(concatenated) // seq_len) * seq_len
+        if total_length == 0:
+            return {"input_ids": []}
+        return {
+            "input_ids": [
+                concatenated[i : i + seq_len] for i in range(0, total_length, seq_len)
+            ]
+        }
+
+    lm_dataset = tokenized.map(
+        group_texts,
+        batched=True,
+        batch_size=1000,
+        remove_columns=tokenized.column_names,
+        desc=f"group-{dataset_name}",
+    )
+    lm_dataset.set_format(type="torch", columns=["input_ids"])
+    return lm_dataset
+
+
+def prepare_ppl_dataloaders(
+    tokenizer,
+    datasets: List[str],
+    configs: List[Optional[str]],
+    split: str,
+    text_field: Optional[str],
+    num_samples: int,
+    seq_len: int,
+    batch_size: int,
+    seed: int,
+    shuffle: bool,
+    model_family: str = "auto",
+    add_bos: str = "auto",
+    cache_dir: Optional[str] = None,
+    num_workers: int = 0,
+    model=None,
+) -> Dict[str, torch.utils.data.DataLoader]:
+    if load_dataset is None:
+        raise SystemExit("datasets is required for dataset evaluation")
+
+    resolved_family = model_family
+    if resolved_family == "auto":
+        if model is None:
+            raise SystemExit("model is required when model_family is 'auto'")
+        resolved_family = _infer_model_family(model)
+    use_bos = _resolve_add_bos(add_bos, resolved_family, tokenizer)
+    if use_bos and tokenizer.bos_token_id is None:
+        use_bos = False
+
+    dataloaders: Dict[str, torch.utils.data.DataLoader] = {}
+    for idx, (dataset_name, config) in enumerate(zip(datasets, configs)):
+        lm_dataset = _load_lm_dataset(
+            tokenizer=tokenizer,
+            dataset_name=dataset_name,
+            config=config,
+            split=split,
+            text_field=text_field,
+            seq_len=seq_len,
+            add_bos=use_bos,
+            cache_dir=cache_dir,
+        )
+        if shuffle:
+            try:
+                lm_dataset = lm_dataset.shuffle(seed=seed + idx)
+            except Exception:
+                pass
+        if num_samples and hasattr(lm_dataset, "__len__"):
+            limit = min(num_samples, len(lm_dataset))
+            lm_dataset = lm_dataset.select(range(limit))
+
+        data_loader = torch.utils.data.DataLoader(
+            lm_dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+        )
+        label = dataset_name if config is None else f"{dataset_name}:{config}"
+        dataloaders[label] = data_loader
+
+    return dataloaders
+
+
+def evaluate_ppl_dataloaders(
+    model,
+    dataloaders: Dict[str, torch.utils.data.DataLoader],
+    device: str,
+    max_batches: Optional[int] = None,
+) -> Dict[str, float]:
+    results: Dict[str, float] = {}
+    for label, data_loader in dataloaders.items():
+        ppl = compute_ppl(model, data_loader, device, max_batches=max_batches)
+        results[label] = ppl
+    return results
+
+
+def evaluate_ppl_datasets(
+    model,
+    tokenizer,
+    datasets: List[str],
+    configs: List[Optional[str]],
+    split: str,
+    text_field: Optional[str],
+    num_samples: int,
+    seq_len: int,
+    batch_size: int,
+    device: str,
+    seed: int,
+    shuffle: bool,
+    model_family: str = "auto",
+    add_bos: str = "auto",
+    max_batches: Optional[int] = None,
+    cache_dir: Optional[str] = None,
+    num_workers: int = 0,
+) -> Dict[str, float]:
+    if load_dataset is None:
+        raise SystemExit("datasets is required for dataset evaluation")
+
+    resolved_family = model_family
+    if resolved_family == "auto":
+        resolved_family = _infer_model_family(model)
+    use_bos = _resolve_add_bos(add_bos, resolved_family, tokenizer)
+    if use_bos and tokenizer.bos_token_id is None:
+        use_bos = False
+
+    results: Dict[str, float] = {}
+    for idx, (dataset_name, config) in enumerate(zip(datasets, configs)):
+        lm_dataset = _load_lm_dataset(
+            tokenizer=tokenizer,
+            dataset_name=dataset_name,
+            config=config,
+            split=split,
+            text_field=text_field,
+            seq_len=seq_len,
+            add_bos=use_bos,
+            cache_dir=cache_dir,
+        )
+        if shuffle:
+            try:
+                lm_dataset = lm_dataset.shuffle(seed=seed + idx)
+            except Exception:
+                pass
+        if num_samples and hasattr(lm_dataset, "__len__"):
+            limit = min(num_samples, len(lm_dataset))
+            lm_dataset = lm_dataset.select(range(limit))
+
+        data_loader = torch.utils.data.DataLoader(
+            lm_dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+        )
+        label = dataset_name if config is None else f"{dataset_name}:{config}"
+        ppl = compute_ppl(model, data_loader, device, max_batches=max_batches)
+        results[label] = ppl
+    return results
+
+
+def main() -> None:
+    args = parse_args()
+    torch.manual_seed(args.seed)
+
+    dtype = get_dtype(args.dtype)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model,
+        torch_dtype=dtype,
+        trust_remote_code=args.trust_remote_code,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model, trust_remote_code=args.trust_remote_code
+    )
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model.to(args.device)
+
+    results: Dict[str, float] = {}
+    resolved_family = args.model_family
+    if resolved_family == "auto":
+        resolved_family = _infer_model_family(model)
+    use_bos = _resolve_add_bos(args.add_bos, resolved_family, tokenizer)
+    if use_bos and tokenizer.bos_token_id is None:
+        use_bos = False
+
+    if args.dataset:
+        datasets = list(args.dataset)
+        configs = _expand_dataset_configs(datasets, list(args.dataset_config))
+        results.update(
+            evaluate_ppl_datasets(
+                model,
+                tokenizer,
+                datasets=datasets,
+                configs=configs,
+                split=args.dataset_split,
+                text_field=args.dataset_text_field,
+                num_samples=args.num_samples,
+                seq_len=args.seq_len,
+                batch_size=args.batch_size,
+                device=args.device,
+                seed=args.seed,
+                shuffle=args.shuffle,
+                model_family=resolved_family,
+                add_bos="always" if use_bos else "never",
+                max_batches=args.max_batches,
+                cache_dir=args.cache_dir,
+                num_workers=args.num_workers,
+            )
+        )
+
+    if args.text_file or args.text:
+        custom_texts: List[str] = []
+        if args.text_file:
+            with open(args.text_file, "r", encoding="utf-8") as handle:
+                custom_texts.extend([line.strip() for line in handle if line.strip()])
+        if args.text:
+            custom_texts.extend([t for t in args.text if t])
+        if custom_texts:
+            chunks = build_token_chunks(
+                custom_texts,
+                tokenizer,
+                args.seq_len,
+                args.num_samples if args.num_samples > 0 else 1_000_000,
+                add_bos=use_bos,
+            )
+            if not chunks:
+                raise SystemExit(
+                    "Not enough custom text to build token sequences. "
+                    "Provide more --text/--text_file content or reduce --seq_len."
+                )
+            dataset = torch.utils.data.TensorDataset(torch.stack(chunks))
+            dataloader = torch.utils.data.DataLoader(
+                dataset, batch_size=args.batch_size, shuffle=False
+            )
+            results["custom"] = compute_ppl(
+                model, dataloader, args.device, max_batches=args.max_batches
+            )
+
+    if not results:
+        raise SystemExit("Provide --dataset and/or --text/--text_file for evaluation")
+
+    print("Perplexity results:")
+    for name, ppl in results.items():
+        print(f"{name}: {ppl:.4f}")
+
+    if args.output:
+        with open(args.output, "w", encoding="utf-8") as handle:
+            json.dump({"model": args.model, "results": results}, handle, indent=2)
+
+
+if __name__ == "__main__":
+    main()

src/ppl_eval_progressive.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+"""Evaluate perplexity for a progressive-pruned model assembled from cycles."""
+
+import argparse
+
+import torch
+
+try:
+    import ppl_eval
+except Exception as exc:  # pragma: no cover - optional dependency
+    raise SystemExit("ppl_eval.py is required (missing or invalid)") from exc
+
+try:
+    from transformers import AutoTokenizer
+except Exception as exc:  # pragma: no cover - fail early with clear error
+    raise SystemExit("transformers is required: pip install transformers") from exc
+
+from progressive_loader import load_progressive_model
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Evaluate PPL for a model reconstructed from progressive cycles."
+    )
+    parser.add_argument("--base_model", required=True, help="Base HF model id or path")
+    parser.add_argument(
+        "--progressive_dir",
+        required=True,
+        help="Output directory from progressive pruning",
+    )
+    parser.add_argument(
+        "--cycle",
+        type=int,
+        default=None,
+        help="Cycle to load (default: final)",
+    )
+    parser.add_argument(
+        "--dataset",
+        action="append",
+        default=[],
+        help="Evaluation dataset name (repeatable). Defaults to wikitext.",
+    )
+    parser.add_argument(
+        "--dataset_config",
+        action="append",
+        default=[],
+        help="Evaluation dataset config (repeatable or single shared config).",
+    )
+    parser.add_argument(
+        "--dataset_split",
+        default="test",
+        help="Evaluation dataset split (default: test)",
+    )
+    parser.add_argument(
+        "--dataset_text_field",
+        default=None,
+        help="Evaluation text field override (default: auto-detect)",
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=0,
+        help="Number of token sequences per dataset (0 = all)",
+    )
+    parser.add_argument(
+        "--seq_len",
+        type=int,
+        default=2048,
+        help="Sequence length for eval",
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=4,
+        help="Batch size for eval",
+    )
+    parser.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device for eval",
+    )
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument(
+        "--model_family",
+        type=str,
+        choices=["auto", "llama", "qwen"],
+        default="auto",
+        help="Model family for BOS handling",
+    )
+    parser.add_argument(
+        "--add_bos",
+        type=str,
+        choices=["auto", "always", "never"],
+        default="auto",
+        help="Whether to prepend BOS to each sample",
+    )
+    parser.add_argument(
+        "--max_batches",
+        type=int,
+        default=None,
+        help="Optional max number of eval batches per dataset",
+    )
+    parser.add_argument(
+        "--cache_dir",
+        default=None,
+        help="Optional datasets cache dir for eval",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=0,
+        help="Eval DataLoader workers",
+    )
+    parser.add_argument(
+        "--dtype",
+        default="auto",
+        choices=["auto", "float32", "float16", "bfloat16"],
+        help="Model dtype",
+    )
+    parser.add_argument(
+        "--trust_remote_code",
+        action="store_true",
+        help="Allow custom model code from hub",
+    )
+    parser.add_argument(
+        "--layer_path",
+        default=None,
+        help="Override layer attribute path if needed",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    torch.manual_seed(args.seed)
+
+    datasets = args.dataset or ["wikitext"]
+    configs = args.dataset_config or ["wikitext-2-raw-v1"]
+    configs = ppl_eval._expand_dataset_configs(datasets, configs)
+
+    model = load_progressive_model(
+        args.base_model,
+        args.progressive_dir,
+        cycle=args.cycle,
+        device=args.device,
+        dtype=args.dtype,
+        trust_remote_code=args.trust_remote_code,
+        layer_path=args.layer_path,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.base_model, trust_remote_code=args.trust_remote_code
+    )
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    results = ppl_eval.evaluate_ppl_datasets(
+        model,
+        tokenizer,
+        datasets=datasets,
+        configs=configs,
+        split=args.dataset_split,
+        text_field=args.dataset_text_field,
+        num_samples=args.num_samples,
+        seq_len=args.seq_len,
+        batch_size=args.batch_size,
+        device=args.device,
+        seed=args.seed,
+        shuffle=False,
+        model_family=args.model_family,
+        add_bos=args.add_bos,
+        max_batches=args.max_batches,
+        cache_dir=args.cache_dir,
+        num_workers=args.num_workers,
+    )
+
+    print("Perplexity results:")
+    for name, ppl in results.items():
+        print(f"{name}: {ppl:.4f}")
+
+
+if __name__ == "__main__":
+    main()

src/print_progressive_ppl_csv.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+"""Print progressive PPL stats as CSV from progressive_metadata.json.
+
+Expected (current) metadata shape:
+  - data["eval"]["pre_ppl"]
+  - data["cycles"][i]["redistrib_post_ppl"] (optional; legacy key)
+  - data["cycles"][i]["comm_post_ppl"] (optional; current key)
+  - data["cycles"][i]["distill_post_ppl"]
+  - data["cycles"][i]["lora_post_ppl"] (typically only set on the last cycle)
+  - data["cycles"][i]["post_ppl"]
+"""
+
+import argparse
+import csv
+import json
+import os
+import shlex
+import sys
+from typing import Any, List, Optional
+
+
+def _cell(value: Any) -> str:
+    if value is None:
+        return ""
+    if isinstance(value, dict):
+        if not value:
+            return ""
+        return ";".join(str(value[key]) for key in sorted(value))
+    if isinstance(value, (list, tuple)):
+        return ";".join(str(item) for item in value)
+    return str(value)
+
+
+def _read_run_command_tokens(metadata_path: str) -> Optional[List[str]]:
+    meta_dir = os.path.dirname(os.path.abspath(metadata_path))
+    run_args_path = os.path.join(meta_dir, "run_args.txt")
+    if not os.path.exists(run_args_path):
+        return None
+
+    try:
+        with open(run_args_path, "r", encoding="utf-8") as handle:
+            lines = handle.read().splitlines()
+    except OSError:
+        return None
+
+    cmd_line = None
+    for idx, line in enumerate(lines):
+        if line.strip() == "command:":
+            if idx + 1 < len(lines):
+                cmd_line = lines[idx + 1].strip()
+            break
+
+    if not cmd_line:
+        return None
+
+    try:
+        return shlex.split(cmd_line)
+    except ValueError:
+        return None
+
+
+def _parse_exclude_pairs_from_tokens(tokens: List[str]) -> Optional[List[int]]:
+    start = None
+    for idx, tok in enumerate(tokens):
+        if tok in ("--exclude_pairs", "--exclude_layers"):
+            start = idx + 1
+            break
+    if start is None:
+        return None
+
+    raw: List[int] = []
+    for tok in tokens[start:]:
+        if tok.startswith("--"):
+            break
+        # Legacy bug: run_args.txt used to print "python" before every token.
+        if tok == "python":
+            continue
+        try:
+            raw.append(int(tok))
+        except ValueError:
+            continue
+    return raw
+
+
+def _normalize_excluded_pairs(raw: List[int], num_pairs: int) -> List[int]:
+    exclude: List[int] = []
+    for idx in raw:
+        if idx < 0:
+            idx = num_pairs + idx
+        if 0 <= idx < num_pairs:
+            exclude.append(idx)
+    return sorted(set(exclude))
+
+
+def _read_excluded_pairs_from_cycle_meta(meta_dir: str, cycle_idx: int) -> Optional[List[int]]:
+    path = os.path.join(meta_dir, f"cycle_{cycle_idx}", "cycle_metadata.json")
+    try:
+        with open(path, "r", encoding="utf-8") as handle:
+            cycle_meta = json.load(handle)
+    except (FileNotFoundError, json.JSONDecodeError, OSError):
+        return None
+
+    dwce_meta = cycle_meta.get("dwce_meta") or {}
+    excluded = dwce_meta.get("excluded_pairs")
+    if isinstance(excluded, list) and all(isinstance(x, int) for x in excluded):
+        return excluded
+    return None
+
+
+def _num_pairs_for_cycle(data: dict, meta_dir: str, cycle_idx: int) -> Optional[int]:
+    num_progressive = data.get("num_progressive")
+    final_num_layers = data.get("final_num_layers")
+    if isinstance(num_progressive, int) and isinstance(final_num_layers, int):
+        initial_layers = final_num_layers + num_progressive
+        return max(initial_layers - cycle_idx, 0)
+
+    cycle_meta_path = os.path.join(meta_dir, f"cycle_{cycle_idx}", "cycle_metadata.json")
+    try:
+        with open(cycle_meta_path, "r", encoding="utf-8") as handle:
+            cycle_meta = json.load(handle)
+    except (FileNotFoundError, json.JSONDecodeError, OSError):
+        return None
+
+    num_layers_before = cycle_meta.get("num_layers_before")
+    if isinstance(num_layers_before, int):
+        return max(num_layers_before - 1, 0)
+    return None
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Print progressive PPL values as CSV from progressive_metadata.json"
+    )
+    parser.add_argument("path", help="Path to progressive_metadata.json")
+    args = parser.parse_args()
+
+    try:
+        with open(args.path, "r", encoding="utf-8") as handle:
+            data = json.load(handle)
+    except FileNotFoundError as exc:
+        raise SystemExit(f"File not found: {args.path}") from exc
+    except json.JSONDecodeError as exc:
+        raise SystemExit(f"Invalid JSON: {args.path}") from exc
+
+    meta_dir = os.path.dirname(os.path.abspath(args.path))
+    run_tokens = _read_run_command_tokens(args.path)
+    raw_exclude = (
+        _parse_exclude_pairs_from_tokens(run_tokens) if run_tokens is not None else None
+    )
+
+    writer = csv.writer(sys.stdout)
+    writer.writerow(
+        [
+            "cycle",
+            "layer_merged",
+            "layer_pair",
+            "excluded_pairs",
+            "redistrib_post_ppl",
+            "distill_post_ppl",
+            "lora_post_ppl",
+            "post_ppl",
+        ]
+    )
+
+    pre_ppl = data.get("eval", {}).get("pre_ppl")
+    if pre_ppl is not None:
+        writer.writerow(["pre", "", "", "", "", "", "", _cell(pre_ppl)])
+
+    cycles = data.get("cycles") or data.get("cycle_summaries") or []
+    for cycle in cycles:
+        cycle_idx = cycle.get("cycle", "")
+        layer_merged = cycle.get("layer_merged")
+        layer_pair = ""
+        if isinstance(layer_merged, int):
+            layer_pair = f"{layer_merged}-{layer_merged + 1}"
+
+        excluded_pairs = _read_excluded_pairs_from_cycle_meta(
+            meta_dir, cycle_idx if isinstance(cycle_idx, int) else -1
+        )
+        if excluded_pairs is None and raw_exclude is not None and isinstance(cycle_idx, int):
+            num_pairs = _num_pairs_for_cycle(data, meta_dir, cycle_idx)
+            if isinstance(num_pairs, int):
+                excluded_pairs = _normalize_excluded_pairs(raw_exclude, num_pairs)
+        redistrib_post_ppl = cycle.get("redistrib_post_ppl")
+        if redistrib_post_ppl is None:
+            redistrib_post_ppl = cycle.get("comm_post_ppl")
+
+        writer.writerow(
+            [
+                cycle_idx,
+                layer_merged if layer_merged is not None else "",
+                layer_pair,
+                _cell(excluded_pairs),
+                _cell(redistrib_post_ppl),
+                _cell(cycle.get("distill_post_ppl")),
+                _cell(cycle.get("lora_post_ppl")),
+                _cell(cycle.get("post_ppl")),
+            ]
+        )
+
+
+if __name__ == "__main__":
+    main()

src/progressive_loader.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+"""Utilities to reconstruct models from progressive pruning cycles."""
+
+import json
+import os
+from typing import Optional
+
+import torch
+
+try:
+    from transformers import AutoModelForCausalLM, PretrainedConfig
+    from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+except Exception as exc:  # pragma: no cover - fail early with clear error
+    raise SystemExit("transformers is required: pip install transformers") from exc
+
+from fuse_layers_model import (
+    decrement_config,
+    drop_layer,
+    find_layer_container,
+    get_dtype,
+    normalize_config,
+)
+
+
+def load_progressive_metadata(output_dir: str) -> dict:
+    path = os.path.join(output_dir, "progressive_metadata.json")
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"Missing progressive metadata at {path}")
+    with open(path, "r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def load_normalized_config(model_path: str, trust_remote_code: bool):
+    config_dict, unused_kwargs = PretrainedConfig.get_config_dict(
+        model_path,
+        trust_remote_code=trust_remote_code,
+    )
+    num_hidden_layers = config_dict.get("num_hidden_layers")
+    layer_types = config_dict.get("layer_types")
+    if (
+        isinstance(num_hidden_layers, int)
+        and num_hidden_layers >= 0
+        and isinstance(layer_types, list)
+        and len(layer_types) != num_hidden_layers
+    ):
+        config_dict["layer_types"] = list(layer_types[:num_hidden_layers])
+    model_type = config_dict["model_type"]
+    config_class = CONFIG_MAPPING[model_type]
+    config = config_class.from_dict(config_dict, **unused_kwargs)
+    normalize_config(config)
+    return config
+
+
+def load_causal_lm(
+    model_path_or_id: str,
+    *,
+    torch_dtype,
+    trust_remote_code: bool,
+    **kwargs,
+) -> torch.nn.Module:
+    config = None
+    config_path = os.path.join(model_path_or_id, "config.json")
+    if os.path.isdir(model_path_or_id) and os.path.isfile(config_path):
+        config = load_normalized_config(model_path_or_id, trust_remote_code)
+    return AutoModelForCausalLM.from_pretrained(
+        model_path_or_id,
+        config=config,
+        torch_dtype=torch_dtype,
+        trust_remote_code=trust_remote_code,
+        **kwargs,
+    )
+
+
+def load_progressive_model(
+    base_model_id: str,
+    output_dir: str,
+    cycle: Optional[int] = None,
+    device: Optional[str] = None,
+    dtype: str = "auto",
+    trust_remote_code: bool = False,
+    layer_path: Optional[str] = None,
+) -> torch.nn.Module:
+    meta = load_progressive_metadata(output_dir)
+    num_cycles = int(meta.get("num_progressive", 0))
+    if cycle is None:
+        cycle = num_cycles
+    if cycle < 0 or cycle > num_cycles:
+        raise ValueError(f"Cycle {cycle} is outside [0, {num_cycles}]")
+
+    if cycle > 0:
+        full_model_dir = os.path.join(output_dir, f"cycle_{cycle}", "full_model")
+        if os.path.isdir(full_model_dir):
+            model = load_causal_lm(
+                full_model_dir,
+                torch_dtype=get_dtype(dtype),
+                trust_remote_code=trust_remote_code,
+            )
+            if device:
+                model.to(device)
+            return model
+
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_id,
+        torch_dtype=get_dtype(dtype),
+        trust_remote_code=trust_remote_code,
+    )
+    active_layer_path = layer_path or meta.get("layer_path")
+    parent, name, container = find_layer_container(model, active_layer_path)
+
+    for idx in range(1, cycle + 1):
+        cycle_dir = os.path.join(output_dir, f"cycle_{idx}")
+        cycle_meta_path = os.path.join(cycle_dir, "cycle_metadata.json")
+        if not os.path.exists(cycle_meta_path):
+            raise FileNotFoundError(f"Missing cycle metadata at {cycle_meta_path}")
+        with open(cycle_meta_path, "r", encoding="utf-8") as handle:
+            cycle_meta = json.load(handle)
+
+        layer_idx = int(cycle_meta["layer_merged"])
+        fused_state = cycle_meta.get("fused_layer_state", "fused_layer.pt")
+        fused_state_path = os.path.join(cycle_dir, fused_state)
+        if not os.path.exists(fused_state_path):
+            raise FileNotFoundError(f"Missing fused layer at {fused_state_path}")
+
+        layers = list(container)
+        if layer_idx < 0 or layer_idx >= len(layers):
+            raise ValueError(
+                f"Cycle {idx} layer index {layer_idx} out of range for {len(layers)} layers"
+            )
+
+        state = torch.load(fused_state_path, map_location="cpu")
+        layers[layer_idx].load_state_dict(state)
+
+        new_container = drop_layer(container, layer_idx + 1)
+        setattr(parent, name, new_container)
+        decrement_config(model.config)
+
+        container = new_container
+
+    if device:
+        model.to(device)
+
+    return model