first commit

Codes/3-2_model_training/_tuning.fix.modified.py

@@ -0,0 +1,599 @@
+"""
+Head-LoRA Finetune (Refactored v2)
+==========================================
+Target: Align specific attention heads to ignore 'data' tokens using Flash Attention Interception.
+
+Key features:
+- LoRA head masking on Q proj: only selected heads' dimensions are updated.
+- LoRA head masking on K proj: only KV groups corresponding to selected heads
+  are updated (GQA-aware).
+- Preservation loss: heads that share a KV group with selected heads but are
+  NOT themselves selected get a KL loss to stay close to the original model,
+  since their K changed but they weren't intended training targets.
+- Heads whose KV group is completely untouched need no loss at all.
+"""
+
+import os
+import json
+import csv
+import argparse
+import math
+import random
+import re
+from typing import List, Tuple, Dict, Set
+
+import torch
+import numpy as np
+from tqdm import tqdm
+from torch.utils.data import Dataset, DataLoader
+from evallib import quick_eval_mmlu, quick_eval_asr_util
+from transformers import (
+    AutoConfig, AutoTokenizer, AutoModelForCausalLM,
+    BitsAndBytesConfig, get_linear_schedule_with_warmup,
+)
+from transformers.models.llama.modeling_llama import ALL_ATTENTION_FUNCTIONS
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel
+from peft.tuners.lora.layer import Linear as LoraLinear
+from lib.tokenize_data_mask import apply_chat_tokenize_with_strip_and_mark
+from lib.printcolor import print_tok_color, print_data_color_in_batch
+
+# -----------------------------------------------------------------------------
+# 1. Flash Attention Spy Interceptor
+# -----------------------------------------------------------------------------
+print("🥷 Injecting Spy Interceptor into Flash Attention 2...")
+
+if "flash_attention_2" not in ALL_ATTENTION_FUNCTIONS:
+    raise RuntimeError("Current environment does not support Flash Attention 2!")
+
+orig_flash_attn = ALL_ATTENTION_FUNCTIONS["flash_attention_2"]
+
+def wrapped_flash_attn(module, query, key, value, attention_mask, scaling, **kwargs):
+    q_last = query[..., -1:, :]
+
+    num_heads = query.shape[1]
+    num_kv_heads = key.shape[1]
+
+    k_expanded = key
+    v_expanded = value
+
+    if num_heads != num_kv_heads:
+        n_rep = num_heads // num_kv_heads
+        k_expanded = key.repeat_interleave(n_rep, dim=1)
+        v_expanded = value.repeat_interleave(n_rep, dim=1)
+
+    attn_weights_tiny = torch.matmul(q_last, k_expanded.transpose(2, 3)) * scaling
+
+    if attention_mask is not None:
+        if attention_mask.dim() == 4:
+            if attention_mask.size(2) > 1:
+                mask_slice = attention_mask[..., -1:, :]
+                attn_weights_tiny = attn_weights_tiny + mask_slice
+            else:
+                attn_weights_tiny = attn_weights_tiny + attention_mask
+        elif attention_mask.dim() == 2:
+            mask_expanded = attention_mask[:, None, None, :]
+            if attention_mask.dtype == torch.bool or (attention_mask.max() <= 1.0 and attention_mask.min() >= 0.0):
+                min_dtype = torch.finfo(attn_weights_tiny.dtype).min
+                attn_weights_tiny = attn_weights_tiny.masked_fill(mask_expanded == 0, min_dtype)
+            else:
+                attn_weights_tiny = attn_weights_tiny + mask_expanded
+
+    attn_probs = torch.nn.functional.softmax(attn_weights_tiny, dim=-1, dtype=torch.float32).to(query.dtype)
+
+    retrieve_map = getattr(module.config, "retrieve_attn_map", None)
+    if retrieve_map is not None and hasattr(module, "layer_idx") and module.layer_idx in retrieve_map:
+        target_heads_dict = retrieve_map[module.layer_idx]
+        for h_idx in target_heads_dict.keys():
+            target_heads_dict[h_idx] = attn_probs[:, h_idx, :, :]
+
+    orig_result = orig_flash_attn(module, query, key, value, attention_mask, scaling=scaling, **kwargs)
+    return orig_result
+
+ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = wrapped_flash_attn
+print("✅ Spy Interceptor installed successfully!")
+
+# -----------------------------------------------------------------------------
+# 2. Utils: Model & Heads
+# -----------------------------------------------------------------------------
+SEED = 42
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+
+def get_lora_targets(model, layers: List[int]) -> List[str]:
+    mtype = (getattr(model.config, "model_type", "") or "").lower()
+    if "llama" in mtype or "mistral" in mtype:
+        return [f"model.layers.{i}.self_attn.{p}" for i in layers for p in ("q_proj", "k_proj")]
+    cand = []
+    for name, _ in model.named_modules():
+        if any(f".{i}." in name for i in layers) and name.split(".")[-1] in {"q_proj", "k_proj"}:
+            cand.append(name)
+    return cand
+
+def load_model(model_path: str):
+    cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=True, padding_side="left")
+
+    if tok.pad_token_id is None:
+        tok.pad_token = tok.eos_token
+        tok.pad_token_id = tok.eos_token_id
+
+    bnb_cfg = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        config=cfg,
+        quantization_config=bnb_cfg,
+        device_map="auto",
+        trust_remote_code=True,
+        attn_implementation="flash_attention_2",
+    )
+    return model, tok, tok
+
+def get_target_structure(heads_file: str, topk_spec: str) -> Dict[int, List[int]]:
+    if not os.path.exists(heads_file):
+        raise FileNotFoundError(f"Heads file not found: {heads_file}")
+
+    with open(heads_file, "r") as f:
+        ihead_list = json.load(f)
+        if type(ihead_list[0]) == str:
+            all_heads = [(str(tag), 0.0) for tag in ihead_list]
+        else:
+            all_heads = ihead_list
+
+    count = len(all_heads)
+    if topk_spec.endswith("p"):
+        count = max(1, math.ceil(float(topk_spec[:-1]) / 100.0 * len(all_heads)))
+    else:
+        count = int(topk_spec)
+
+    target_heads = all_heads[:min(count, len(all_heads))]
+
+    structure = {}
+    for tag, _ in target_heads:
+        try:
+            parts = tag.split('_')
+            l = int(parts[0][1:])
+            h = int(parts[1][1:])
+            structure.setdefault(l, []).append(h)
+        except Exception:
+            continue
+
+    print(f"🎯 Selected {len(target_heads)} heads from {topk_spec}.")
+    return structure
+
+def discover_existing_adapter(out_dir: str):
+    if not os.path.isdir(out_dir): return None, 0, 0
+    candidates = []
+    if os.path.exists(os.path.join(out_dir, "adapter_config.json")):
+        candidates.append((0, 0, out_dir))
+
+    for entry in os.listdir(out_dir):
+        path = os.path.join(out_dir, entry)
+        if os.path.isdir(path) and os.path.exists(os.path.join(path, "adapter_config.json")):
+            m = re.match(r"^batch_(\d+)_(\d+)$", entry)
+            if m:
+                epoch_idx = int(m.group(1))
+                batch_idx = int(m.group(2))
+                candidates.append((epoch_idx, batch_idx, path))
+
+    if not candidates: return None, 0, 0
+    candidates.sort(key=lambda x: (x[0], x[1]))
+    return candidates[-1][2], candidates[-1][0], candidates[-1][1]
+
+# -----------------------------------------------------------------------------
+# 3. GQA-Aware Head Analysis
+# -----------------------------------------------------------------------------
+def compute_gqa_affected_heads(target_structure: Dict[int, List[int]],
+                               num_q_heads: int,
+                               num_kv_heads: int) -> Dict[int, List[int]]:
+    """
+    For each layer, find Q heads that are NOT selected but share a KV group
+    with at least one selected head. These heads will have their K changed
+    by LoRA but were not intended as training targets.
+
+    Returns: {layer_idx: [affected_but_unselected_head_indices]}
+    """
+    group_size = num_q_heads // num_kv_heads
+
+    affected_structure = {}
+    for l, selected_heads in target_structure.items():
+        selected_set = set(selected_heads)
+
+        # Find which KV groups are touched
+        touched_kv_groups = set()
+        for h in selected_heads:
+            touched_kv_groups.add(h // group_size)
+
+        # Find unselected heads in those touched groups
+        affected = []
+        for kv_g in touched_kv_groups:
+            for offset in range(group_size):
+                q_head = kv_g * group_size + offset
+                if q_head not in selected_set:
+                    affected.append(q_head)
+
+        if affected:
+            affected_structure[l] = affected
+
+    return affected_structure
+
+# -----------------------------------------------------------------------------
+# 4. LoRA Head Masking (Monkey-Patch)
+# -----------------------------------------------------------------------------
+def apply_head_mask_to_lora(model, target_structure: Dict[int, List[int]],
+                            num_q_heads: int, num_kv_heads: int, head_dim: int):
+    """
+    Monkey-patch LoRA layers so that:
+    - q_proj: LoRA output is zeroed for unselected Q heads
+    - k_proj: LoRA output is zeroed for KV groups not associated with any selected head
+    """
+    group_size = num_q_heads // num_kv_heads
+
+    for name, module in model.named_modules():
+        if not isinstance(module, LoraLinear):
+            continue
+
+        # Parse layer index from name
+        layer_idx = None
+        parts = name.split(".")
+        for i, part in enumerate(parts):
+            if part == "layers" and i + 1 < len(parts) and parts[i + 1].isdigit():
+                layer_idx = int(parts[i + 1])
+                break
+
+        if layer_idx is None or layer_idx not in target_structure:
+            continue
+
+        selected_heads = set(target_structure[layer_idx])
+        is_q = name.endswith("q_proj")
+        is_k = name.endswith("k_proj")
+
+        if not (is_q or is_k):
+            continue
+
+        if is_q:
+            mask = torch.zeros(num_q_heads * head_dim)
+            for h in selected_heads:
+                mask[h * head_dim : (h + 1) * head_dim] = 1.0
+            desc = f"Q mask: {len(selected_heads)}/{num_q_heads} heads"
+        else:
+            touched_kv_groups = set()
+            for h in selected_heads:
+                touched_kv_groups.add(h // group_size)
+
+            mask = torch.zeros(num_kv_heads * head_dim)
+            for g in touched_kv_groups:
+                mask[g * head_dim : (g + 1) * head_dim] = 1.0
+            desc = f"K mask: {len(touched_kv_groups)}/{num_kv_heads} KV groups"
+
+        module.register_buffer("head_mask", mask)
+
+        orig_forward = module.forward
+
+        def make_masked_forward(orig_fn, mod):
+            def masked_forward(x, *args, **kwargs):
+                result = orig_fn(x, *args, **kwargs)
+                base_out = torch.nn.functional.linear(x, mod.base_layer.weight, mod.base_layer.bias)
+                lora_delta = result - base_out
+                masked_delta = lora_delta * mod.head_mask.to(lora_delta.device)
+                return base_out + masked_delta
+            return masked_forward
+
+        module.forward = make_masked_forward(orig_forward, module)
+        print(f"  🎭 Patched {name}: {desc}")
+
+# -----------------------------------------------------------------------------
+# 5. Loss Functions
+# -----------------------------------------------------------------------------
+def head_attention_loss(tuned_map, base_map, data_mask, lambda_data=1.0, eps=1e-8):
+    """Loss for SELECTED heads: KL toward blind teacher + data mass penalty."""
+    mask_data = data_mask.bool()
+    mask_valid = ~mask_data
+
+    if mask_valid.dim() == 2:
+        mask_valid = mask_valid.unsqueeze(1)
+        mask_data = mask_data.unsqueeze(1)
+
+    total = 0.0
+    n = 0
+
+    for l, heads in tuned_map.items():
+        for h, tuned in heads.items():
+            base = base_map.get(l, {}).get(h, None)
+            if tuned is None or base is None:
+                continue
+
+            tuned = tuned.float()
+            base = base.float().to(tuned.device)
+
+            base_v = base * mask_valid.float()
+            base_v = base_v / base_v.sum(dim=-1, keepdim=True).clamp_min(eps)
+
+            tuned_v = tuned * mask_valid.float()
+            tuned_v_sum = tuned_v.sum(dim=-1, keepdim=True).clamp_min(eps)
+            tuned_v = tuned_v / tuned_v_sum
+
+            kl = (base_v * (torch.log(base_v + eps) - torch.log(tuned_v + eps))).sum(dim=-1).mean()
+            data_mass = (tuned * mask_data.float()).sum(dim=-1).mean()
+
+            total = total + kl + lambda_data * data_mass
+            n += 1
+
+    return total / max(n, 1)
+
+
+def head_preservation_loss(tuned_map, orig_map, eps=1e-8):
+    """Loss for KV-group-affected but unselected heads: KL toward original model."""
+    total = 0.0
+    n = 0
+    for l, heads in orig_map.items():
+        for h, orig in heads.items():
+            tuned = tuned_map.get(l, {}).get(h, None)
+            if tuned is None or orig is None:
+                continue
+            tuned = tuned.float()
+            orig = orig.float().to(tuned.device)
+            kl = (orig * (torch.log(orig + eps) - torch.log(tuned + eps))).sum(dim=-1).mean()
+            total += kl
+            n += 1
+    return total / max(n, 1)
+
+
+# -----------------------------------------------------------------------------
+# 6. Data Loading
+# -----------------------------------------------------------------------------
+class JsonlMessagesDS(Dataset):
+    def __init__(self, jsonl_path: str):
+        self.samples = []
+        with open(jsonl_path, "r", encoding="utf-8") as f:
+            for line in f:
+                if line.strip():
+                    obj = json.loads(line)
+                    self.samples.append(obj)
+        random.shuffle(self.samples)
+        print(f"📊 Loaded {len(self.samples)} samples.")
+
+    def __len__(self): return len(self.samples)
+    def __getitem__(self, idx): return self.samples[idx]
+
+def collate(batch, tokenizer):
+    input_ids, attention_mask, data_mask = apply_chat_tokenize_with_strip_and_mark(
+        batch, tokenizer, device="cpu", add_generation_prompt=True,
+        encode_kwargs={"padding_side": "left"},
+        mode="custom_mask == 'inst'",
+        custom_mask_identifier={"data": ["<data>", "</data>"], "inst": ["<inst>", "</inst>"]},
+        return_tensors="pt"
+    )
+    return {"input_ids": input_ids, "attention_mask": attention_mask, "data_mask": data_mask}
+
+# -----------------------------------------------------------------------------
+# 7. Training Loop
+# -----------------------------------------------------------------------------
+def save_model(model, tok, out_dir, epoch, batch_idx, data_path):
+    print("Saving checkpoint...")
+
+    temp_map = getattr(model.config, "retrieve_attn_map", None)
+    if temp_map is not None:
+        del model.config.retrieve_attn_map
+
+    save_dir = os.path.join(out_dir, f"batch_{epoch}_{batch_idx}")
+    model.save_pretrained(save_dir)
+    tok.save_pretrained(save_dir)
+
+    print(f"✅ Saved to {save_dir}")
+
+    mllu_result = quick_eval_mmlu(model, tok)
+    asr_result = quick_eval_asr_util(
+        model, tok, training_data_path=data_path,
+    )
+    log_path = os.path.join(out_dir, "training_log.csv")
+    write_header = not os.path.exists(log_path)
+    with open(log_path, "a", newline="") as f:
+        writer = csv.writer(f)
+        attack_columns = [
+            ("naive", "a_naive", "v_naive"),
+            ("ignore", "a_ignore", "v_ignore"),
+            ("escape_separation", "a_escape", "v_escape"),
+            ("completion_realcmb", "a_cmb", "v_cmb"),
+            ("conv_attack", "a_conv", "v_conv"),
+            ("none", "a_none", "v_none"),
+        ]
+        if write_header:
+            header = ["epoch", "batch", "mmlu score", "asr_mode", "asr_status"]
+            for _, a_col, v_col in attack_columns:
+                header.append(a_col)
+                header.append(v_col)
+            writer.writerow(header)
+        row = [epoch, batch_idx, mllu_result.get("accuracy"), asr_result.get("eval_mode"), asr_result.get("status")]
+        metrics = asr_result.get("metrics", {}) if asr_result.get("status") == "ok" else {}
+        for attack, _, _ in attack_columns:
+            values = metrics.get(attack, {})
+            row.append(values.get("asr"))
+            row.append(values.get("valid_rate"))
+        writer.writerow(row)
+    return save_dir
+
+
+def tune_new(model, tok, tok_inf, target_structure, data_path, out_dir, epochs, bs, lr,
+             resume_adapter=None, start_epoch=0, start_batch_idx=0, batch_save_interval=-1,
+             lambda_preserve=1.0):
+    # -------------------------------------------------------------------------
+    # GQA analysis
+    # -------------------------------------------------------------------------
+    num_q_heads = model.config.num_attention_heads
+    num_kv_heads = getattr(model.config, "num_key_value_heads", num_q_heads)
+    head_dim = model.config.hidden_size // num_q_heads
+    group_size = num_q_heads // num_kv_heads
+
+    print(f"📐 GQA config: {num_q_heads} Q heads, {num_kv_heads} KV heads, group_size={group_size}")
+
+    affected_structure = compute_gqa_affected_heads(target_structure, num_q_heads, num_kv_heads)
+    has_affected = len(affected_structure) > 0
+    if has_affected:
+        total_affected = sum(len(v) for v in affected_structure.values())
+        print(f"🛡️ Preservation targets: {total_affected} KV-group-affected unselected heads across {len(affected_structure)} layers (λ={lambda_preserve})")
+        for l, heads in sorted(affected_structure.items()):
+            print(f"   Layer {l}: affected heads {heads} (selected: {target_structure[l]})")
+    else:
+        print("✅ No KV-group collateral — all touched KV groups are fully selected. No preservation loss needed.")
+
+    # -------------------------------------------------------------------------
+    # Setup LoRA
+    # -------------------------------------------------------------------------
+    layers = sorted(target_structure.keys())
+    targets = get_lora_targets(model, layers)
+
+    lora_cfg = LoraConfig(r=32, lora_alpha=16, bias="none", target_modules=targets, task_type="CAUSAL_LM")
+    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False)
+
+    if resume_adapter:
+        model = PeftModel.from_pretrained(model, resume_adapter, is_trainable=True)
+        print(f"♻️ Resumed LoRA: {resume_adapter}")
+    else:
+        model = get_peft_model(model, lora_cfg)
+
+    # -------------------------------------------------------------------------
+    # Apply head mask to LoRA layers (physical isolation)
+    # -------------------------------------------------------------------------
+    print("🎭 Applying GQA-aware head masks to LoRA layers...")
+    apply_head_mask_to_lora(model, target_structure, num_q_heads, num_kv_heads, head_dim)
+
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    ratio = trainable / total * 100
+    print(f"🔧 Trainable parameters: {trainable:,} / {total:,} ({ratio:.4f}% of total)")
+
+    # -------------------------------------------------------------------------
+    # Dataloader & Optimizer
+    # -------------------------------------------------------------------------
+    dl = DataLoader(JsonlMessagesDS(data_path), batch_size=bs, shuffle=True, collate_fn=lambda b: collate(b, tok))
+    opt = torch.optim.AdamW(model.parameters(), lr=lr)
+    sch = get_linear_schedule_with_warmup(opt, int(0.05 * epochs * len(dl)), epochs * len(dl))
+
+    current_idx = start_epoch
+
+    for ep in range(epochs):
+        if ep < start_epoch:
+            continue
+        pbar = tqdm(dl, desc=f"Ep {ep+1}/{epochs}")
+        for idx, batch in enumerate(pbar):
+            if ep == start_epoch and idx < start_batch_idx:
+                continue
+            device = next(model.parameters()).device
+            MAX_LEN = 50000
+
+            input_ids = batch["input_ids"][:, :MAX_LEN].to(device)
+            attn_mask = batch["attention_mask"][:, :MAX_LEN].to(device)
+            data_mask = batch["data_mask"][:, :MAX_LEN].to(device)
+
+            blind_attn_mask = attn_mask * (~data_mask).long()
+
+            print_data_color_in_batch(5,tok,input_ids,attn_mask)
+            print_data_color_in_batch(5,tok,input_ids,data_mask)
+            os.exit(0)
+
+            # --- A. Base Model Pass (no LoRA, blind mask) ---
+            # Teacher for selected heads
+            model.eval()
+            base_map_container = {l: {h: None for h in h_list} for l, h_list in target_structure.items()}
+            model.config.retrieve_attn_map = base_map_container
+
+            with torch.no_grad(), model.disable_adapter():
+                model(input_ids=input_ids, attention_mask=blind_attn_mask, output_attentions=False, use_cache=False)
+            base_attns_map = base_map_container
+
+            # --- A2. Original Model Pass (no LoRA, full mask) ---
+            # Teacher for KV-group-affected unselected heads
+            if has_affected:
+                orig_map_container = {l: {h: None for h in h_list} for l, h_list in affected_structure.items()}
+                model.config.retrieve_attn_map = orig_map_container
+
+                with torch.no_grad(), model.disable_adapter():
+                    model(input_ids=input_ids, attention_mask=attn_mask, output_attentions=False, use_cache=False)
+                orig_attns_map = orig_map_container
+            else:
+                orig_attns_map = None
+
+            # --- B. Tuned Model Pass (with LoRA, full mask) ---
+            # Collect attention for both selected and affected heads
+            model.train()
+            tuned_map_container = {}
+            for l in target_structure.keys():
+                all_heads_needed = set(target_structure[l])
+                if l in affected_structure:
+                    all_heads_needed.update(affected_structure[l])
+                tuned_map_container[l] = {h: None for h in all_heads_needed}
+            model.config.retrieve_attn_map = tuned_map_container
+
+            model(input_ids=input_ids, attention_mask=attn_mask, output_attentions=False, use_cache=False)
+            tuned_attns_map = model.config.retrieve_attn_map
+
+            # --- C. Loss & Step ---
+            loss_selected = head_attention_loss(tuned_attns_map, base_attns_map, data_mask)
+
+            loss_preserve = torch.tensor(0.0, device=device)
+            if has_affected and orig_attns_map is not None:
+                loss_preserve = head_preservation_loss(tuned_attns_map, orig_attns_map)
+
+            loss = loss_selected + lambda_preserve * loss_preserve
+
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step()
+            sch.step()
+            opt.zero_grad()
+
+            del model.config.retrieve_attn_map
+            pbar.set_postfix(
+                loss=f"{loss.item():.4f}",
+                sel=f"{loss_selected.item():.4f}",
+                pres=f"{loss_preserve.item():.4f}" if has_affected else "N/A",
+            )
+            if batch_save_interval > 0 and (idx + 1) % batch_save_interval == 0:
+                save_model(model, tok, out_dir, current_idx, idx, data_path)
+        save_model(model, tok, out_dir, current_idx, idx, data_path)
+        current_idx += 1
+
+# -----------------------------------------------------------------------------
+# 8. Main
+# -----------------------------------------------------------------------------
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", required=True)
+    parser.add_argument("--data_path", default="../3-1_model_training_preprocess/inj-likechen/crafted_instruction_data_tri_injection_qa.jsonl")
+    parser.add_argument("--head_path", default="../2-2_head_identification/head_scoring/llama31-8b_injsq_dev/heads_sorted/user_prop_inst_0.1.json")
+
+    parser.add_argument("--output_dir", default="outputs_lora")
+    parser.add_argument("--lora_path", default=None)
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--batch-save-interval", type=int, default=-1)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--topk", type=str, default="18.75p")
+    parser.add_argument("--lambda_preserve", type=float, default=1.0,
+                        help="Weight for KV-group-affected preservation loss")
+    args = parser.parse_args()
+
+    resume_adapter, start_epoch, start_batch_idx = discover_existing_adapter(args.output_dir)
+    if args.lora_path:
+        resume_adapter, start_epoch, start_batch_idx = args.lora_path, 0, 0
+
+    model, tok, tok_inf = load_model(args.model_path)
+    target_structure = get_target_structure(args.head_path, args.topk)
+    os.makedirs(args.output_dir, exist_ok=True)
+    tune_new(
+        model, tok, tok_inf, target_structure,
+        args.data_path, args.output_dir,
+        args.epochs, args.batch_size, args.lr,
+        resume_adapter=resume_adapter, start_epoch=start_epoch, start_batch_idx=start_batch_idx,
+        batch_save_interval=args.batch_save_interval,
+        lambda_preserve=args.lambda_preserve,
+    )
+
+if __name__ == "__main__":
+    main()