# -*- coding: utf-8 -*-
"""
ESFT-Token 专家筛选（Qwen3-30B-A3B）
对应论文《Let the Expert Stick to His Last》(Wang et al., 2024):
  - 式(7) Token Selection Ratio: 统计每个专家被 Top-K 选中的 token 比例
  - 式(8) 累积阈值筛选: 按相关度降序累积，首个 >= p 的 Top 专家子集

用法:
    python expert_selection.py \
        --data_path data/gov_corpus.jsonl \
        --threshold_p 0.2 \
        --output selected_experts.json

硬件: 单张 80G 卡即可（bf16 权重约 60G），device_map="auto" 也支持多卡。
"""
import argparse
import json

import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model_path", default="Qwen/Qwen3-30B-A3B")
    ap.add_argument("--data_path", required=True, help="JSONL，每行含 text 字段")
    ap.add_argument("--num_samples", type=int, default=32, help="论文默认采样 32 段")
    ap.add_argument("--seq_len", type=int, default=4096, help="论文默认 L=4096")
    ap.add_argument("--threshold_p", type=float, default=0.2, help="ESFT-Token 论文取 0.2")
    ap.add_argument("--output", default="selected_experts.json")
    args = ap.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",  # 没装 flash-attn 就改 "sdpa"
        device_map="auto",
    )
    model.eval()

    cfg = model.config
    L, E, K = cfg.num_hidden_layers, cfg.num_experts, cfg.num_experts_per_tok
    # Qwen3-30B-A3B: L=48, E=128, K=8

    # ---------- 1. 构造采样数据：拼接语料后切成 num_samples 段定长序列（与论文一致） ----------
    ds = load_dataset("json", data_files=args.data_path, split="train")
    ids = []
    for t in ds["text"]:
        ids.extend(tokenizer(t, add_special_tokens=False)["input_ids"])
        if len(ids) >= args.num_samples * args.seq_len:
            break
    ids = ids[: args.num_samples * args.seq_len]
    samples = torch.tensor(
        [ids[i * args.seq_len:(i + 1) * args.seq_len] for i in range(args.num_samples)]
    )

    # ---------- 2. 给每层 router(gate) 注册 hook，统计各专家被选中的 token 次数 ----------
    counts = [torch.zeros(E, dtype=torch.long) for _ in range(L)]
    total_tokens = 0

    def make_hook(li):
        def hook(module, inputs, output):
            # output: router_logits, shape (tokens, E)
            # 式(7) 中 1(g>0) 即"进入 Top-K"
            topk_idx = torch.topk(output.float(), K, dim=-1).indices
            counts[li] += torch.bincount(topk_idx.reshape(-1).cpu(), minlength=E)

        return hook

    handles = []
    for i in range(L):
        mlp = model.model.layers[i].mlp
        assert hasattr(mlp, "gate"), f"layer {i} 不是 MoE 层"
        handles.append(mlp.gate.register_forward_hook(make_hook(i)))

    embed_device = model.model.embed_tokens.weight.device
    bs = 1  # 显存宽裕可加大
    for i in range(0, samples.size(0), bs):
        batch = samples[i:i + bs].to(embed_device)
        total_tokens += batch.numel()
        model(input_ids=batch)

    for h in handles:
        h.remove()

    # ---------- 3. 式(7) 归一化 + 式(8) 累积阈值筛选 ----------
    selected, stats = {}, []
    for li in range(L):
        r = counts[li].float() / (total_tokens * K)  # 全层专家求和 = 1
        r_sorted, order = torch.sort(r, descending=True)
        cum = torch.cumsum(r_sorted, dim=0)
        n = int((cum < args.threshold_p).sum().item()) + 1  # 首个累积 >= p 的子集
        selected[str(li)] = sorted(order[:n].tolist())
        stats.append(n)

    result = {
        "model": args.model_path,
        "threshold_p": args.threshold_p,
        "num_samples": args.num_samples,
        "seq_len": args.seq_len,
        "total_tokens": total_tokens,
        "num_experts_per_layer": E,
        "top_k": K,
        "avg_selected": sum(stats) / L,
        "min_selected": min(stats),
        "max_selected": max(stats),
        "selected_experts": selected,
    }
    with open(args.output, "w") as f:
        json.dump(result, f, indent=2, ensure_ascii=False)

    print(f"[Done] 每层选中专家: 平均 {sum(stats)/L:.1f} / {E} "
          f"({sum(stats)/L/E:.1%}), 最少 {min(stats)}, 最多 {max(stats)}")
    print(f"[Done] 已保存 -> {args.output}")


if __name__ == "__main__":
    main()

"""
ESFT + LoRA 训练（Qwen3-30B-A3B, Transformers + PEFT + DeepSpeed + FlashAttention）

- mode=lora（默认）: 只对 expert_selection.py 选出的专家注入 LoRA，其余全部参数冻结
- mode=esft: 论文原版做法，选中专家全量更新，不加 LoRA（建议 --lr 1e-5）

启动示例（8 卡 ZeRO-2）:
    deepspeed --num_gpus 8 train_esft_lora.py \
        --data_path data/gov_sft.jsonl \
        --selected_experts selected_experts.json \
        --deepspeed ds_zero2.json \
        --output_dir ./ckpt_esft_lora

依赖: transformers>=4.51, peft, datasets, accelerate, deepspeed, flash-attn
"""
import argparse
import json

import torch
from datasets import load_dataset
from peft import LoraConfig, get_peft_model
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
)


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model_path", default="Qwen/Qwen3-30B-A3B")
    ap.add_argument("--data_path", required=True, help="JSONL，每行含 text 字段")
    ap.add_argument("--selected_experts", default="selected_experts.json")
    ap.add_argument("--mode", choices=["lora", "esft"], default="lora")
    ap.add_argument("--seq_len", type=int, default=4096)
    ap.add_argument("--output_dir", default="./ckpt_esft_lora")
    ap.add_argument("--deepspeed", default=None)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--epochs", type=float, default=2.0)
    ap.add_argument("--micro_bs", type=int, default=1)
    ap.add_argument("--grad_accum", type=int, default=16)
    ap.add_argument("--lora_r", type=int, default=16)
    ap.add_argument("--lora_alpha", type=int, default=32)
    return ap.parse_args()


def main():
    args = parse_args()
    with open(args.selected_experts) as f:
        selected = json.load(f)["selected_experts"]  # {"0": [3, 17, ...], ...}

    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
    )
    model.config.use_cache = False
    model.config.output_router_logits = False  # 不需要 MoE 辅助负载均衡 loss
    model.gradient_checkpointing_enable(
        gradient_checkpointing_kwargs={"use_reentrant": False}
    )

    if args.mode == "lora":
        # ---- 只对选中专家的三个投影注入 LoRA；PEFT 默认冻结其余全部参数 ----
        target_modules = [
            f"layers.{li}.mlp.experts.{eid}.{proj}"
            for li, eids in selected.items()
            for eid in eids
            for proj in ("gate_proj", "up_proj", "down_proj")
        ]
        lora_cfg = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            lora_dropout=0.05,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=target_modules,
        )
        model = get_peft_model(model, lora_cfg)

        # 校验：注入 LoRA 的 Linear 数必须等于目标模块数，防止 peft 版本匹配逻辑差异
        n_lora = sum(1 for n, _ in model.named_parameters() if "lora_A" in n)
        assert n_lora == len(target_modules), (
            f"LoRA 注入数 {n_lora} != 目标模块数 {len(target_modules)}，"
            f"请检查 peft 版本的 target_modules 匹配逻辑"
        )
        model.print_trainable_parameters()
    else:
        # ---- 论文原版 ESFT：全模型冻结，仅解冻选中专家的权重 ----
        for p in model.parameters():
            p.requires_grad_(False)
        n_train = 0
        for li, eids in selected.items():
            for eid in eids:
                expert = model.model.layers[int(li)].mlp.experts[eid]
                for p in expert.parameters():
                    p.requires_grad_(True)
                    n_train += p.numel()
        n_total = sum(p.numel() for p in model.parameters())
        print(f"[ESFT] 可训参数 {n_train/1e9:.2f}B / {n_total/1e9:.1f}B "
              f"({n_train/n_total:.1%})")

    # 梯度检查点 + 冻结嵌入时必须加这行，否则第一层收不到梯度
    model.enable_input_require_grads()

    # ---------- 数据 ----------
    ds = load_dataset("json", data_files=args.data_path, split="train")

    def tokenize(batch):
        return tokenizer(batch["text"], truncation=True, max_length=args.seq_len)

    ds = ds.map(tokenize, batched=True, num_proc=8, remove_columns=ds.column_names)
    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    targs = TrainingArguments(
        output_dir=args.output_dir,
        per_device_train_batch_size=args.micro_bs,
        gradient_accumulation_steps=args.grad_accum,
        num_train_epochs=args.epochs,
        learning_rate=args.lr,
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        weight_decay=0.0,
        logging_steps=5,
        save_strategy="steps",
        save_steps=500,
        save_total_limit=3,
        bf16=True,
        optim="adamw_torch",
        deepspeed=args.deepspeed,
        report_to="none",
        seed=42,
    )

    trainer = Trainer(
        model=model,
        args=targs,
        train_dataset=ds,
        data_collator=collator,
    )
    trainer.train()

    trainer.save_model(args.output_dir)      # LoRA 模式下只保存 adapter
    tokenizer.save_pretrained(args.output_dir)


if __name__ == "__main__":
    main()