预训练（PT）→ 持续预训练（CPT）→ 指令微调（SFT）→ 偏好对齐（DPO/RLHF）
              ↑
         本文重点

from typing import List, Dict
import json
import os

class DomainDataBuilder:
    """领域预训练数据构建器"""
    
    def __init__(self, domain: str):
        self.domain = domain  # "medical", "legal", "financial", etc.

    def collect_sources(self) -> Dict[str, List[str]]:
        """定义数据来源"""
        sources = {
            "medical": {
                "textbooks": ["医学教科书 PDF", "诊疗指南"],
                "papers": ["PubMed 论文", "中文医学期刊"],
                "records": ["脱敏病历", "医学百科"],
                "guidelines": ["临床路径", "用药指南"],
                "qa": ["医学问答", "丁香园问答"],
            },
            "legal": {
                "textbooks": ["法学教材", "法律释义"],
                "papers": ["法学期刊论文"],
                "records": ["裁判文书", "法律案例"],
                "guidelines": ["司法解释", "部门规章"],
                "qa": ["法律咨询问答"],
            },
            "financial": {
                "textbooks": ["金融学教材", "CFA教材"],
                "papers": ["金融研究论文"],
                "reports": ["年报", "券商研报"],
                "guidelines": ["会计准则", "监管文件"],
                "qa": ["金融问答", "投资者教育"],
            },
        }
        return sources.get(self.domain, {})

    def estimate_tokens(self, file_sizes_mb: List[float]) -> int:
        """估算 token 数（中文约 1.5 字/token，英文约 0.75 词/token）"""
        total_chars = sum(size * 1024 * 1024 * 3 for size in file_sizes_mb)  # 假设中文3字节
        return int(total_chars / 1.5)  # 粗略估算

import re
from typing import Optional

class DataCleaner:
    """预训练数据清洗"""
    
    @staticmethod
    def clean_text(text: str) -> Optional[str]:
        """单文档清洗"""
        # 1. 去除 HTML 标签
        text = re.sub(r'<[^>]+>', '', text)
        
        # 2. 去除 URL
        text = re.sub(r'https?://\S+', '', text)
        
        # 3. 去除多余空白
        text = re.sub(r'\s+', ' ', text).strip()
        
        # 4. 去除乱码（非可打印字符）
        text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]', '', text)
        
        # 5. 去重检测（文档级）
        if len(text) < 100:  # 太短，丢弃
            return None
        
        # 6. 质量过滤：字母/汉字占比
        alpha_ratio = len(re.findall(r'[a-zA-Z\u4e00-\u9fff]', text)) / len(text)
        if alpha_ratio < 0.5:  # 非文字内容太多
            return None
        
        return text
    
    @staticmethod
    def deduplicate(texts: List[str], threshold: float = 0.8) -> List[str]:
        """MinHash 去重"""
        from datasketch import MinHash, MinHashLSH
        
        lsh = MinHashLSH(threshold=threshold, num_perm=128)
        unique_texts = []
        
        for i, text in enumerate(texts):
            m = MinHash(num_perm=128)
            # 按 5-gram 分词
            words = [text[j:j+5] for j in range(0, len(text)-4, 3)]
            for word in words:
                m.update(word.encode('utf-8'))
            
            # 检查是否重复
            result = lsh.query(m)
            if not result:
                lsh.insert(str(i), m)
                unique_texts.append(text)
        
        return unique_texts

def build_mixed_dataset(
    domain_data: List[str],     # 领域文本
    general_data: List[str],    # 通用文本（如预训练数据子集）
    domain_ratio: float = 0.7,  # 领域数据占比
) -> List[str]:
    """
    混合领域数据和通用数据，防止灾难性遗忘
    """
    import random
    
    n_domain = int(len(domain_data) * domain_ratio / (1 - domain_ratio))
    n_domain = min(n_domain, len(domain_data))
    
    sampled_domain = random.sample(domain_data, n_domain)
    sampled_general = random.sample(general_data, len(domain_data) - n_domain + len(domain_data))
    
    # 交错排列（避免连续大量同类型数据）
    mixed = []
    domain_iter = iter(sampled_domain)
    general_iter = iter(sampled_general)
    
    while True:
        try:
            mixed.append(next(domain_iter))
        except StopIteration:
            pass
        try:
            mixed.append(next(general_iter))
        except StopIteration:
            pass
        if not (domain_iter or general_iter):
            break
    
    return mixed

# 推荐混合比例
ratio_guide = {
    "medical": {"domain": 0.6, "general": 0.3, "code": 0.1},   # 医疗
    "legal":   {"domain": 0.65, "general": 0.25, "code": 0.10}, # 法律
    "financial": {"domain": 0.55, "general": 0.35, "code": 0.10}, # 金融
    "code":    {"domain": 0.7, "general": 0.2, "code": 0.10},   # 代码
    "general_domain": {"domain": 0.5, "general": 0.4, "code": 0.1}, # 通用领域
}

# 学习率对比
lr_comparison = {
    "预训练 (PT)":      "3e-4 ~ 6e-4",
    "持续预训练 (CPT)": "5e-5 ~ 2e-5",   # 预训练的 1/10 ~ 1/20
    "指令微调 (SFT)":   "1e-5 ~ 5e-6",
    "DPO":              "5e-7 ~ 1e-6",
}

# 余弦衰减 + Warmup
def get_cpt_lr_scheduler(
    optimizer,
    warmup_steps: int = 2000,
    total_steps: int = 100000,
    min_lr_ratio: float = 0.1,    # 最低学习率为峰值的 10%
):
    """
    CPT 专用学习率调度器：Warmup → Cosine Decay
    """
    from transformers import get_cosine_schedule_with_warmup
    
    scheduler = get_cosine_schedule_with_warmup(
        optimizer=optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=total_steps,
    )
    return scheduler

# 不同领域的推荐配置
domain_configs = {
    "medical": {"lr": 1e-5, "warmup": 2000, "epochs": 1, "batch_size": 256},
    "legal":   {"lr": 1.5e-5, "warmup": 1500, "epochs": 1, "batch_size": 256},
    "financial": {"lr": 1e-5, "warmup": 2000, "epochs": 2, "batch_size": 128},
    "code":    {"lr": 2e-5, "warmup": 1000, "epochs": 1, "batch_size": 256},
}

class AntiForgettingTrainer:
    """带遗忘缓解的 CPT 训练器"""
    
    def __init__(
        self,
        model,
        domain_dataloader,
        general_dataloader,
        replay_ratio: float = 0.3,    # 通用数据回放比例
        ewc_lambda: float = 0.01,     # EWC 正则化系数
    ):
        self.model = model
        self.domain_loader = domain_dataloader
        self.general_loader = general_dataloader
        self.replay_ratio = replay_ratio
        self.ewc_lambda = ewc_lambda
        self.fisher_info = None  # Fisher 信息矩阵
        
    def compute_fisher_information(self, sample_data):
        """
        计算 Fisher 信息矩阵（用于 EWC 正则化）
        在通用数据上计算，约束重要参数不大幅偏移
        """
        fisher = {}
        self.model.eval()
        
        for batch in sample_data:
            self.model.zero_grad()
            outputs = self.model(**batch)
            loss = outputs.loss
            loss.backward()
            
            for name, param in self.model.named_parameters():
                if param.grad is not None:
                    if name not in fisher:
                        fisher[name] = param.grad.data.clone() ** 2
                    else:
                        fisher[name] += param.grad.data.clone() ** 2
        
        # 平均
        for name in fisher:
            fisher[name] /= len(sample_data)
        
        self.fisher_info = fisher
        self.model.train()
    
    def training_step(self, batch):
        """带 EWC 正则化的训练步"""
        # 标准 LM loss
        outputs = self.model(**batch)
        loss = outputs.loss
        
        # EWC 正则化
        if self.fisher_info and self.ewc_lambda > 0:
            ewc_loss = 0
            for name, param in self.model.named_parameters():
                if name in self.fisher_info:
                    # 约束参数不偏离初始值太远
                    ewc_loss += (self.fisher_info[name] * (param - self.initial_params[name]) ** 2).sum()
            loss += self.ewc_lambda * ewc_loss
        
        return loss

class ReplayDataLoader:
    """混合领域数据 + 通用数据回放"""
    
    def __init__(
        self,
        domain_loader,
        general_loader,
        replay_ratio: float = 0.3,
    ):
        self.domain_loader = domain_loader
        self.general_loader = general_loader
        self.replay_ratio = replay_ratio
        
    def __iter__(self):
        domain_iter = iter(self.domain_loader)
        general_iter = iter(self.general_loader)
        
        domain_steps = int(1 / (1 - self.replay_ratio) * self.replay_ratio)
        
        while True:
            # 1 batch 领域数据
            try:
                domain_batch = next(domain_iter)
                yield {"batch": domain_batch, "type": "domain"}
            except StopIteration:
                break
            
            # replay_ratio 比例的通用数据
            for _ in range(domain_steps):
                try:
                    general_batch = next(general_iter)
                    yield {"batch": general_batch, "type": "general"}
                except StopIteration:
                    break

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset, Dataset
from torch.optim import AdamW

def run_cpt(
    base_model: str = "meta-llama/Meta-Llama-3-8B",
    domain_data_path: str = "./data/domain",
    general_data_path: str = "./data/general",
    output_dir: str = "./output/cpt-model",
):
    """完整 CPT 训练流程"""
    
    # ===== 1. 加载模型 =====
    tokenizer = AutoTokenizer.from_pretrained(base_model)
    tokenizer.pad_token = tokenizer.eos_token
    
    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        attn_implementation="flash_attention_2",
    )
    
    # ===== 2. 加载数据 =====
    def tokenize_function(examples):
        # 拼接文本，添加 EOS
        texts = [t + tokenizer.eos_token for t in examples["text"]]
        result = tokenizer(
            texts,
            truncation=True,
            max_length=2048,
            return_overflowing_tokens=True,
        )
        return result
    
    domain_dataset = load_dataset("text", data_dir=domain_data_path, split="train")
    general_dataset = load_dataset("text", data_dir=general_data_path, split="train")
    
    domain_tokenized = domain_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
    general_tokenized = general_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
    
    # 混合数据（70% 领域 + 30% 通用）
    from datasets import concatenate_datasets
    domain_sampled = domain_tokenized.shuffle().select(range(int(len(domain_tokenized) * 0.7)))
    general_sampled = general_tokenized.shuffle().select(range(int(len(domain_tokenized) * 0.3)))
    mixed_dataset = concatenate_datasets([domain_sampled, general_sampled]).shuffle()
    
    # ===== 3. 数据收集器 =====
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,  # 因果语言模型
    )
    
    # ===== 4. 优化器与调度器 =====
    optimizer = AdamW(
        model.parameters(),
        lr=1e-5,
        weight_decay=0.01,
        betas=(0.9, 0.95),
    )
    
    total_steps = len(mixed_dataset) // (8 * 32)  # batch_size=8, grad_accum=32
    from transformers import get_cosine_schedule_with_warmup
    scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=2000,
        num_training_steps=total_steps,
    )
    
    # ===== 5. 训练参数 =====
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=1,
        per_device_train_batch_size=8,
        gradient_accumulation_steps=32,  # 有效 batch = 256
        learning_rate=1e-5,
        weight_decay=0.01,
        warmup_steps=2000,
        lr_scheduler_type="cosine",
        logging_steps=50,
        save_strategy="steps",
        save_steps=5000,
        save_total_limit=3,
        bf16=True,
        gradient_checkpointing=True,
        optim="adamw_torch",
        report_to="tensorboard",
        dataloader_num_workers=8,
        seed=42,
    )
    
    # ===== 6. 启动训练 =====
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=mixed_dataset,
        data_collator=data_collator,
        tokenizer=tokenizer,
        optimizers=(optimizer, scheduler),
    )
    
    # 保存初始参数（用于 EWC）
    initial_params = {name: param.data.clone() for name, param in model.named_parameters()}
    
    trainer.train()
    
    # ===== 7. 保存模型 =====
    trainer.save_model(output_dir)
    tokenizer.save_pretrained(output_dir)
    
    return model

class CPTEvaluator:
    """CPT 评估器"""
    
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer
    
    def evaluate_perplexity(self, dataset) -> float:
        """领域困惑度（越低越好）"""
        self.model.eval()
        total_loss, total_tokens = 0, 0
        
        with torch.no_grad():
            for batch in dataset:
                outputs = self.model(**batch)
                total_loss += outputs.loss.item() * batch["input_ids"].size(1)
                total_tokens += batch["input_ids"].size(1)
        
        avg_loss = total_loss / total_tokens
        return torch.exp(torch.tensor(avg_loss)).item()
    
    def evaluate_domain_benchmark(self, benchmark_name: str) -> dict:
        """领域 benchmark 评估"""
        benchmarks = {
            "medical": ["MedQA", "MedMCQA", "PubMedQA", "CMB"],
            "legal":   ["LawBench", "JEC-QA", "LegalBench"],
            "financial": ["FinEval", "CFLEB"],
        }
        # 运行 benchmark 评估...
        pass
    
    def evaluate_general(self) -> dict:
        """通用能力评估（检测遗忘）"""
        general_benchmarks = ["MMLU", "C-Eval", "BBH", "GSM8K"]
        # 运行通用 benchmark...
        pass

def check_forgetting(
    model,
    pre_cpt_scores: dict,   # CPT 前的 benchmark 分数
    post_cpt_scores: dict,  # CPT 后的 benchmark 分数
    threshold: float = 0.03,  # 允许的最大降幅
) -> dict:
    """
    检测灾难性遗忘
    """
    report = {}
    for benchmark, (pre, post) in zip(
        pre_cpt_scores.keys(),
        zip(pre_cpt_scores.values(), post_cpt_scores.values())
    ):
        change = post - pre
        report[benchmark] = {
            "pre_cpt": pre,
            "post_cpt": post,
            "change": change,
            "status": "⚠️ 遗忘" if change < -threshold else "✅ 正常",
        }
    return report

{
  "bf16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "steps_per_print": 2000
}

# 单节点 8 卡
deepspeed --num_gpus=8 train_cpt.py \
    --deepspeed ds_config_zero3.json \
    --model meta-llama/Meta-Llama-3-8B \
    --output_dir ./output/cpt

# 多节点
deepspeed --num_nodes=2 --num_gpus=8 \
    --hostfile hostfile \
    train_cpt.py \
    --deepspeed ds_config_zero3.json

checklist_before = [
    "✅ 基座模型选择（开源 license 允许商用？）",
    "✅ 领域数据量 >= 10B tokens（太少不如直接 SFT）",
    "✅ 数据质量审核（去重、去噪、质量评分）",
    "✅ 通用数据采样（至少 30% 混合比例）",
    "✅ GPU 资源评估（7B 至少 4×A100 40GB）",
    "✅ 训练监控配置（TensorBoard/WandB）",
    "✅ Benchmark 评估基线（CPT 前跑一遍）",
    "✅ Checkpoint 保存策略",
]

checklist_after = [
    "✅ 领域 benchmark 提升达标",
    "✅ 通用 benchmark 降幅 < 3%",
    "✅ 人工抽检生成质量",
    "✅ 对比 CPT 前后输出（同 prompt）",
    "✅ 后续 SFT 流程就绪",
]

# 完整流程示例
pipeline = """
# Step 1: CPT（领域知识注入）
python train_cpt.py --model llama-3-8b --domain medical --data ./data/medical

# Step 2: SFT（指令跟随能力）
python train_sft.py --model ./output/cpt-model --data ./data/medical-sft

# Step 3: DPO（偏好对齐）
python train_dpo.py --model ./output/sft-model --data ./data/medical-dpo

# Step 4: 评估
python evaluate.py --model ./output/dpo-model --benchmarks MedQA,MMLU,GSM8K

# Step 5: 量化部署
python quantize.py --model ./output/dpo-model --method awq --bits 4
"""

# 可以，但效果有限
# LoRA CPT 的效果介于 SFT 和全量 CPT 之间
# 适合预算受限的场景

# 推荐：LoRA rank >= 128，target_modules 覆盖所有线性层
config = LoraConfig(
    r=256,
    lora_alpha=512,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.0,
)
# 注意：LoRA CPT 的知识注入深度不如全量 CPT