第七章：训练循环与迭代优化

    ┌─────────────┐
    │   数据收集   │ ← 用户反馈/合成生成
    └──────┬──────┘
           │
    ┌──────▼──────┐
    │   数据标注   │ ← 人工/模型辅助
    └──────┬──────┘
           │
    ┌──────▼──────┐
    │   模型训练   │ ← SFT/RLHF/DPO
    └──────┬──────┘
           │
    ┌──────▼──────┐
    │   模型评估   │ → 指标监控
    └──────┬──────┘
           │
    ┌──────▼──────┐
    │  部署/迭代   │
    └─────────────┘

输入源 → 数据验证 → 去重清洗 → 格式标准化 → 数据池
  ↓         ↓          ↓           ↓          ↓
监控      质量报告   清洗日志    schema检查  版本控制

# 数据版本配置示例
data_version = {
    "version": "v2.3.1",
    "base_dataset": "v2.3.0",
    "incremental": {
        "user_feedback": 50000,
        "synthetic": 100000,
        "hard_negatives": 20000
    },
    "filters_applied": [
        "deduplication",
        "quality_threshold_0.8",
        "safety_filter"
    ],
    "split_ratio": {
        "train": 0.9,
        "val": 0.05,
        "test": 0.05
    }
}

原始数据 → 模型预标注 → 人工审核/修正 → 质量采样检查 → 最终标注
             ↓              ↓                ↓
         置信度评分    标注者一致性    随机质检(5-10%)

def should_trigger_training():
    conditions = [
        data_accumulated > min_data_threshold,
        days_since_last_training > max_wait_days,
        performance_degradation > alert_threshold,
        critical_bug_fixes_pending
    ]
    return any(conditions) or all(conditions[:2])

评估结果 ──┬── 数据收集：指导困难样本收集
          ├── 标注优化：调整标注规范
          ├── 训练策略：调整损失权重
          └── 模型选择：决定部署版本

def evaluate_and_decide(model_metrics):
    decisions = {
        "deploy": all([
            model_metrics["accuracy"] > baseline["accuracy"] * 1.02,
            model_metrics["safety_score"] > 0.95,
            model_metrics["latency_p99"] < baseline["latency_p99"] * 1.1
        ]),
        "collect_more_data": model_metrics["accuracy"] < target_accuracy,
        "adjust_training": model_metrics["loss_variance"] > threshold,
        "rollback": model_metrics["safety_score"] < 0.9
    }
    return decisions

def compute_llm_uncertainty(model, prompt, num_samples=10):
    """通过多次采样估计 LLM 的不确定性"""
    responses = []
    log_probs = []

    for _ in range(num_samples):
        response, lp = model.generate(
            prompt, 
            temperature=0.7,
            return_log_probs=True
        )
        responses.append(response)
        log_probs.append(lp)

    # 计算响应多样性
    unique_responses = len(set(responses))
    diversity_score = unique_responses / num_samples

    # 计算平均对数概率方差
    avg_log_prob = np.mean([np.mean(lp) for lp in log_probs])
    log_prob_variance = np.var([np.mean(lp) for lp in log_probs])

    uncertainty = {
        "diversity": diversity_score,
        "avg_confidence": np.exp(avg_log_prob),
        "confidence_variance": log_prob_variance
    }

    return uncertainty

def diversity_sampling(embeddings, n_samples, n_clusters=100):
    # 先聚类
    kmeans = KMeans(n_clusters=n_clusters)
    cluster_labels = kmeans.fit_predict(embeddings)

    selected_indices = []
    samples_per_cluster = n_samples // n_clusters

    for cluster_id in range(n_clusters):
        cluster_indices = np.where(cluster_labels == cluster_id)[0]
        # 从每个簇中选择最接近中心的样本
        cluster_center = kmeans.cluster_centers_[cluster_id]
        distances = np.linalg.norm(
            embeddings[cluster_indices] - cluster_center, 
            axis=1
        )
        selected = cluster_indices[np.argsort(distances)[:samples_per_cluster]]
        selected_indices.extend(selected)

    return selected_indices

def find_hard_examples(model, dataset, percentile=95):
    losses = []
    for batch in dataset:
        with torch.no_grad():
            loss = model.compute_loss(batch)
            losses.extend(loss.cpu().numpy())

    threshold = np.percentile(losses, percentile)
    hard_indices = np.where(losses > threshold)[0]
    return hard_indices

def generate_adversarial_prompts(model, base_prompt, epsilon=0.1):
    """生成对抗性提示"""
    embedding = model.encode(base_prompt)
    embedding.requires_grad = True

    output = model(embedding)
    loss = compute_target_loss(output)
    loss.backward()

    # 对嵌入添加扰动
    perturbation = epsilon * embedding.grad.sign()
    adversarial_embedding = embedding + perturbation

    # 解码回文本（近似）
    adversarial_prompt = model.decode(adversarial_embedding)
    return adversarial_prompt

class CurriculumScheduler:
    def __init__(self, dataset, difficulty_scores):
        self.dataset = dataset
        self.difficulty_scores = difficulty_scores
        self.current_difficulty = 0.0

    def get_batch(self, epoch, max_epoch):
        # 线性增加难度阈值
        difficulty_threshold = (epoch / max_epoch) * 1.0

        # 选择难度低于阈值的样本
        valid_indices = np.where(
            self.difficulty_scores <= difficulty_threshold
        )[0]

        # 加权采样，优先选择接近阈值的样本
        weights = 1.0 - np.abs(
            self.difficulty_scores[valid_indices] - difficulty_threshold * 0.8
        )
        weights = weights / weights.sum()

        batch_indices = np.random.choice(
            valid_indices, 
            size=batch_size, 
            p=weights
        )

        return self.dataset[batch_indices]

def adaptive_curriculum(model, dataset, window_size=1000):
    performance_history = deque(maxlen=window_size)

    for batch in dataset:
        loss = model.train_step(batch)
        performance_history.append(loss)

        # 如果性能停滞，增加难度
        if len(performance_history) == window_size:
            recent_improvement = abs(
                np.mean(list(performance_history)[:500]) - 
                np.mean(list(performance_history)[500:])
            )

            if recent_improvement < threshold:
                # 引入更难的样本
                increase_difficulty_level()

def compute_data_shapley(model, dataset, test_set, n_iterations=100):
    n_samples = len(dataset)
    shapley_values = np.zeros(n_samples)

    for _ in range(n_iterations):
        # 随机排列
        perm = np.random.permutation(n_samples)

        prev_score = evaluate(model, test_set)

        for i, idx in enumerate(perm):
            # 增量训练
            model.update(dataset[idx])
            new_score = evaluate(model, test_set)

            # 边际贡献
            shapley_values[idx] += (new_score - prev_score)
            prev_score = new_score

    return shapley_values / n_iterations

def slerp(theta_A, theta_B, alpha):
    """球面线性插值，保持参数向量的范数"""
    # 归一化参数向量
    theta_A_norm = theta_A / np.linalg.norm(theta_A)
    theta_B_norm = theta_B / np.linalg.norm(theta_B)

    # 计算夹角
    dot_product = np.dot(theta_A_norm, theta_B_norm)
    omega = np.arccos(np.clip(dot_product, -1, 1))

    # 球面插值
    if np.abs(omega) < 1e-10:
        return alpha * theta_A + (1 - alpha) * theta_B

    sin_omega = np.sin(omega)
    theta_merged = (np.sin((1 - alpha) * omega) / sin_omega) * theta_A + \
                  (np.sin(alpha * omega) / sin_omega) * theta_B

    return theta_merged

def fisher_weighted_merge(models, dataset):
    """基于 Fisher 信息的参数合并"""
    fisher_matrices = []

    for model in models:
        # 计算 Fisher 信息矩阵（对角近似）
        fisher = compute_fisher_diagonal(model, dataset)
        fisher_matrices.append(fisher)

    # 归一化 Fisher 权重
    total_fisher = sum(fisher_matrices)
    weights = [f / total_fisher for f in fisher_matrices]

    # 加权合并
    merged_params = {}
    for param_name in models[0].state_dict():
        merged_params[param_name] = sum(
            w[param_name] * m.state_dict()[param_name] 
            for w, m in zip(weights, models)
        )

    return merged_params

def add_capabilities(base_model, task_vectors, scaling_factors=None):
    """向基础模型添加多个任务能力"""
    if scaling_factors is None:
        scaling_factors = [1.0] * len(task_vectors)

    merged_params = base_model.state_dict().copy()

    for param_name in merged_params:
        # 累加任务向量
        delta = sum(
            scale * tv[param_name] 
            for scale, tv in zip(scaling_factors, task_vectors)
        )
        merged_params[param_name] += delta

    return merged_params

def remove_capability(model, task_vector, scale=1.0):
    """从模型中移除特定能力"""
    updated_params = model.state_dict().copy()

    for param_name in updated_params:
        updated_params[param_name] -= scale * task_vector[param_name]

    return updated_params

def detect_task_conflicts(task_vectors):
    """检测任务向量之间的冲突"""
    conflicts = {}

    for i, tv1 in enumerate(task_vectors):
        for j, tv2 in enumerate(task_vectors[i+1:], i+1):
            # 计算余弦相似度
            similarity = cosine_similarity(
                flatten(tv1), flatten(tv2)
            )

            # 负相关表示潜在冲突
            if similarity < -0.5:
                conflicts[(i, j)] = similarity

    return conflicts

def compute_layer_importance(model, dataset, layer_names):
    """计算各层对任务的重要性"""
    importances = {}

    for layer_name in layer_names:
        # 临时移除层
        original_layer = getattr(model, layer_name)
        setattr(model, layer_name, nn.Identity())

        # 评估性能下降
        degraded_performance = evaluate(model, dataset)

        # 恢复层
        setattr(model, layer_name, original_layer)
        baseline_performance = evaluate(model, dataset)

        importance = baseline_performance - degraded_performance
        importances[layer_name] = importance

    return importances

def hierarchical_merge(models, split_layer=12):
    """低层参数平均，高层保持特化"""
    merged_model = models[0].copy()

    for layer_idx in range(len(merged_model.layers)):
        if layer_idx < split_layer:
            # 低层：平均合并
            merged_model.layers[layer_idx] = average_layers(
                [m.layers[layer_idx] for m in models]
            )
        else:
            # 高层：选择最佳或保持独立
            best_model_idx = select_best_for_layer(
                models, layer_idx
            )
            merged_model.layers[layer_idx] = \
                models[best_model_idx].layers[layer_idx]

    return merged_model

def merge_attention_heads(models, importance_threshold=0.1):
    """根据重要性选择性合并注意力头"""
    merged_attention = []

    n_heads = models[0].n_heads
    for head_idx in range(n_heads):
        head_candidates = [
            m.attention.heads[head_idx] for m in models
        ]

        # 评估每个头的重要性
        importances = [
            evaluate_head_importance(h) for h in head_candidates
        ]

        if max(importances) > importance_threshold:
            # 选择最重要的头
            best_idx = np.argmax(importances)
            merged_attention.append(head_candidates[best_idx])
        else:
            # 低重要性头可以平均
            merged_attention.append(average_heads(head_candidates))

    return merged_attention

class WeightedEnsemble:
    def __init__(self, models, weights=None):
        self.models = models
        self.weights = weights or [1/len(models)] * len(models)

    def predict(self, input_data):
        predictions = []
        for model, weight in zip(self.models, self.weights):
            pred = model.predict(input_data)
            predictions.append(weight * pred)

        # 加权平均
        ensemble_pred = sum(predictions)

        # 对于分类任务，可能需要重新归一化
        if self.task_type == 'classification':
            ensemble_pred = softmax(ensemble_pred)

        return ensemble_pred

class MixtureOfExperts:
    def __init__(self, experts, gating_network):
        self.experts = experts
        self.gating = gating_network

    def forward(self, x):
        # 计算门控权重
        gate_weights = self.gating(x)  # [batch_size, n_experts]

        # 获取每个专家的输出
        expert_outputs = torch.stack([
            expert(x) for expert in self.experts
        ], dim=1)  # [batch_size, n_experts, output_dim]

        # 加权组合
        output = torch.sum(
            gate_weights.unsqueeze(-1) * expert_outputs, 
            dim=1
        )

        return output

def cascade_ensemble(models, input_data, confidence_threshold=0.8):
    """按置信度级联使用模型"""
    for i, model in enumerate(models):
        prediction, confidence = model.predict_with_confidence(input_data)

        if confidence > confidence_threshold or i == len(models) - 1:
            return prediction

    # 如果所有模型置信度都低，可以组合预测
    return weighted_ensemble(models, input_data)

def resolve_sign_conflicts(param_A, param_B, method='magnitude'):
    """解决参数符号相反的冲突"""
    sign_conflict = (param_A * param_B) < 0

    if method == 'magnitude':
        # 选择绝对值较大的
        merged = torch.where(
            torch.abs(param_A) > torch.abs(param_B),
            param_A, param_B
        )
    elif method == 'zero':
        # 冲突位置置零
        merged = torch.where(sign_conflict, 0, (param_A + param_B) / 2)
    elif method == 'interpolate':
        # 根据重要性插值
        importance_A = compute_param_importance(param_A)
        importance_B = compute_param_importance(param_B)
        alpha = importance_A / (importance_A + importance_B)
        merged = alpha * param_A + (1 - alpha) * param_B

    return merged

def gradient_surgery(gradients, threshold=0.0):
    """梯度手术：投影冲突梯度到正交空间"""
    modified_grads = []

    for i, g_i in enumerate(gradients):
        g_modified = g_i.clone()

        for j, g_j in enumerate(gradients):
            if i != j:
                # 计算余弦相似度
                cos_sim = F.cosine_similarity(g_i, g_j, dim=0)

                if cos_sim < threshold:  # 负相关，存在冲突
                    # 投影到正交空间
                    projection = (g_i @ g_j) / (g_j @ g_j + 1e-8) * g_j
                    g_modified = g_modified - projection

        modified_grads.append(g_modified)

    return modified_grads

def validate_merge_quality(original_models, merged_model, test_sets):
    """验证合并后模型的质量"""
    metrics = {
        'performance_retention': [],
        'capability_preservation': [],
        'emergence_score': 0
    }

    for original, test_set in zip(original_models, test_sets):
        # 性能保持率
        original_score = evaluate(original, test_set)
        merged_score = evaluate(merged_model, test_set)
        retention = merged_score / original_score
        metrics['performance_retention'].append(retention)

        # 能力保留检查
        capability_tests = generate_capability_tests(original)
        preservation = evaluate_capabilities(merged_model, capability_tests)
        metrics['capability_preservation'].append(preservation)

    # 涌现能力（合并后出现的新能力）
    combined_test = combine_test_sets(test_sets)
    emergence = evaluate_emergence(merged_model, combined_test, original_models)
    metrics['emergence_score'] = emergence

    return metrics

hyperparameter_space = {
    # 学习率相关
    "learning_rate": {
        "type": "log_uniform",
        "low": 1e-6,
        "high": 1e-3,
        "description": "峰值学习率"
    },
    "warmup_ratio": {
        "type": "uniform",
        "low": 0.0,
        "high": 0.1,
        "description": "预热步数比例"
    },

    # 训练配置
    "batch_size": {
        "type": "choice",
        "values": [8, 16, 32, 64],
        "description": "有效批次大小"
    },
    "gradient_accumulation_steps": {
        "type": "choice",
        "values": [1, 2, 4, 8],
        "description": "梯度累积步数"
    },

    # 正则化
    "weight_decay": {
        "type": "log_uniform",
        "low": 1e-4,
        "high": 1e-1
    },
    "dropout": {
        "type": "uniform",
        "low": 0.0,
        "high": 0.3
    },

    # RLHF 特定参数
    "kl_coefficient": {
        "type": "log_uniform",
        "low": 0.001,
        "high": 0.1,
        "description": "KL 散度惩罚系数"
    },
    "clip_range": {
        "type": "uniform",
        "low": 0.1,
        "high": 0.3,
        "description": "PPO 裁剪范围"
    }
}

def conditional_hyperparameters(config):
    """处理超参数间的条件依赖"""
    # 如果使用 LoRA，添加相关参数
    if config.get("use_lora", False):
        config["lora_rank"] = sample_from(4, 64, log=True)
        config["lora_alpha"] = config["lora_rank"] * 2
        config["lora_dropout"] = sample_from(0.0, 0.1)

    # 学习率与批次大小的关系
    if config["batch_size"] > 32:
        config["learning_rate"] *= np.sqrt(config["batch_size"] / 32)

    # 梯度累积与实际批次大小
    config["effective_batch_size"] = (
        config["batch_size"] * config["gradient_accumulation_steps"]
    )

    return config

def staged_search():
    # 第一阶段：粗粒度搜索
    stage1_space = {
        "learning_rate": [1e-5, 1e-4, 1e-3],
        "batch_size": [16, 32, 64]
    }
    best_config_stage1 = grid_search(stage1_space)

    # 第二阶段：细粒度搜索
    stage2_space = create_refined_space(best_config_stage1)
    best_config = bayesian_optimization(stage2_space)

    return best_config

class BayesianOptimizer:
    def __init__(self, space, acquisition_func="ei"):
        self.space = space
        self.gp = GaussianProcessRegressor()
        self.acquisition_func = acquisition_func
        self.observations = []

    def suggest_next(self):
        if len(self.observations) < 5:
            # 初始随机探索
            return self.random_sample()

        # 训练高斯过程
        X = [obs["config"] for obs in self.observations]
        y = [obs["score"] for obs in self.observations]
        self.gp.fit(X, y)

        # 优化获取函数
        next_point = self.optimize_acquisition()
        return next_point

    def optimize_acquisition(self):
        """优化获取函数找到下一个采样点"""
        if self.acquisition_func == "ei":
            return self.expected_improvement()
        elif self.acquisition_func == "ucb":
            return self.upper_confidence_bound()
        elif self.acquisition_func == "pi":
            return self.probability_improvement()

    def expected_improvement(self):
        """期望改进获取函数"""
        best_y = max([obs["score"] for obs in self.observations])

        def ei(x):
            mean, std = self.gp.predict(x, return_std=True)
            z = (mean - best_y) / (std + 1e-9)
            ei_value = (mean - best_y) * norm.cdf(z) + std * norm.pdf(z)
            return -ei_value  # 负值用于最小化

        # 多起点优化
        best_x = None
        best_ei = float("-inf")

        for _ in range(10):
            x0 = self.random_sample()
            result = minimize(ei, x0, bounds=self.space.bounds)
            if -result.fun > best_ei:
                best_ei = -result.fun
                best_x = result.x

        return best_x

def successive_halving(configs, budget, eta=3):
    """Successive Halving 算法"""
    n = len(configs)
    r = budget / n  # 初始资源分配

    while n > 1:
        # 训练所有配置 r 个 epoch
        scores = []
        for config in configs:
            score = train_and_evaluate(config, epochs=int(r))
            scores.append(score)

        # 保留前 1/eta
        k = max(1, n // eta)
        top_k_indices = np.argsort(scores)[-k:]
        configs = [configs[i] for i in top_k_indices]

        # 增加资源
        n = len(configs)
        r = r * eta

    return configs[0]

class PopulationBasedTraining:
    def __init__(self, population_size=16):
        self.population_size = population_size
        self.population = self.initialize_population()

    def initialize_population(self):
        """初始化种群"""
        population = []
        for _ in range(self.population_size):
            member = {
                "config": self.sample_hyperparameters(),
                "model": None,
                "score": 0,
                "age": 0
            }
            population.append(member)
        return population

    def evolve(self, steps=100):
        """进化过程"""
        for step in range(steps):
            # 并行训练一段时间
            self.train_population(epochs=5)

            # 评估性能
            self.evaluate_population()

            # 执行进化操作
            if step % 10 == 0:
                self.exploit_and_explore()

    def exploit_and_explore(self):
        """利用和探索"""
        # 排序种群
        self.population.sort(key=lambda x: x["score"], reverse=True)

        # 底部 25% 被顶部 25% 替换
        bottom_quartile = self.population_size // 4
        top_quartile = self.population_size // 4

        for i in range(bottom_quartile):
            # 复制高性能成员
            source_idx = i % top_quartile
            self.population[-(i+1)] = self.copy_member(
                self.population[source_idx]
            )

            # 扰动超参数（探索）
            self.perturb_hyperparameters(self.population[-(i+1)])

    def perturb_hyperparameters(self, member):
        """扰动超参数进行探索"""
        config = member["config"]
        for param, value in config.items():
            if random.random() < 0.2:  # 20% 概率扰动
                if isinstance(value, float):
                    # 乘性扰动
                    factor = random.choice([0.8, 1.2])
                    config[param] = value * factor
                elif isinstance(value, int):
                    # 加性扰动
                    delta = random.choice([-1, 1])
                    config[param] = max(1, value + delta)

class AdaptiveEarlyStopping:
    def __init__(self, patience=10, min_delta=0.001):
        self.patience = patience
        self.min_delta = min_delta
        self.best_score = float('-inf')
        self.counter = 0
        self.early_stop = False

    def __call__(self, current_score, epoch):
        # 相对改进
        relative_improvement = (
            (current_score - self.best_score) / (abs(self.best_score) + 1e-10)
        )

        if relative_improvement > self.min_delta:
            self.best_score = current_score
            self.counter = 0
        else:
            self.counter += 1

        # 动态调整 patience
        if epoch < 10:
            adjusted_patience = self.patience * 2  # 早期更宽容
        elif epoch > 50:
            adjusted_patience = self.patience // 2  # 后期更严格
        else:
            adjusted_patience = self.patience

        if self.counter >= adjusted_patience:
            self.early_stop = True

        return self.early_stop

def hyperband(max_iter=81, eta=3):
    """Hyperband 算法"""
    logeta = lambda x: np.log(x) / np.log(eta)
    s_max = int(logeta(max_iter))
    B = (s_max + 1) * max_iter

    results = []

    for s in reversed(range(s_max + 1)):
        n = int(np.ceil(B / max_iter / (s + 1) * eta ** s))
        r = max_iter * eta ** (-s)

        # Successive halving
        configs = [sample_configuration() for _ in range(n)]

        for i in range(s + 1):
            n_configs = n * eta ** (-i)
            n_iterations = r * eta ** i

            # 训练和评估
            scores = []
            for config in configs[:int(n_configs)]:
                score = train_and_evaluate(
                    config, 
                    iterations=int(n_iterations)
                )
                scores.append((score, config))

            # 选择最佳配置继续
            scores.sort(reverse=True)
            configs = [config for _, config in scores[:int(n_configs / eta)]]

        results.extend(scores[:1])  # 保存最佳结果

    return max(results, key=lambda x: x[0])

class HyperparameterMetaLearner:
    def __init__(self):
        self.task_embeddings = {}
        self.hyperparameter_history = []
        self.meta_model = self.build_meta_model()

    def build_meta_model(self):
        """构建元学习模型"""
        return RandomForestRegressor(n_estimators=100)

    def learn_from_task(self, task_features, best_hyperparams, performance):
        """从任务中学习"""
        self.hyperparameter_history.append({
            "task_features": task_features,
            "hyperparams": best_hyperparams,
            "performance": performance
        })

        # 更新元模型
        if len(self.hyperparameter_history) > 10:
            X = [h["task_features"] + list(h["hyperparams"].values()) 
                 for h in self.hyperparameter_history]
            y = [h["performance"] for h in self.hyperparameter_history]
            self.meta_model.fit(X, y)

    def suggest_initial_hyperparams(self, new_task_features, n_suggestions=5):
        """为新任务建议初始超参数"""
        if len(self.hyperparameter_history) < 10:
            # 历史数据不足，返回随机配置
            return [sample_random_config() for _ in range(n_suggestions)]

        # 找到相似任务
        similar_tasks = self.find_similar_tasks(new_task_features)

        suggestions = []
        for task in similar_tasks[:n_suggestions]:
            # 基于相似任务的超参数，添加小扰动
            base_config = task["hyperparams"].copy()
            perturbed_config = self.perturb_config(base_config)
            suggestions.append(perturbed_config)

        return suggestions

    def find_similar_tasks(self, task_features, k=5):
        """找到最相似的历史任务"""
        distances = []
        for hist in self.hyperparameter_history:
            dist = np.linalg.norm(
                np.array(task_features) - np.array(hist["task_features"])
            )
            distances.append((dist, hist))

        distances.sort(key=lambda x: x[0])
        return [hist for _, hist in distances[:k]]

class HybridParallelConfig:
    def __init__(self, world_size, model_config):
        self.world_size = world_size
        self.model_config = model_config
        self.strategy = self.determine_strategy()

    def determine_strategy(self):
        """根据模型和硬件自动确定并行策略"""
        model_params = self.model_config.num_parameters
        gpu_memory = torch.cuda.get_device_properties(0).total_memory

        # 估算单GPU能否容纳模型
        bytes_per_param = 4  # FP32
        model_memory = model_params * bytes_per_param
        activation_memory = self.estimate_activation_memory()

        if model_memory + activation_memory < gpu_memory * 0.8:
            # 纯数据并行
            return {"dp": self.world_size, "tp": 1, "pp": 1}

        # 需要模型并行
        if self.model_config.num_layers > 24:
            # 深层网络，使用流水线并行
            pp_size = min(8, self.model_config.num_layers // 12)
            remaining = self.world_size // pp_size

            if self.model_config.hidden_size > 8192:
                # 宽模型，添加张量并行
                tp_size = min(8, remaining)
                dp_size = remaining // tp_size
            else:
                tp_size = 1
                dp_size = remaining
        else:
            # 浅层宽模型，主要使用张量并行
            tp_size = min(8, self.world_size)
            dp_size = self.world_size // tp_size
            pp_size = 1

        return {"dp": dp_size, "tp": tp_size, "pp": pp_size}

def configure_zero_optimization(stage=2):
    """配置 ZeRO 优化"""
    zero_config = {
        "stage": stage,
        "allgather_partitions": True,
        "allgather_bucket_size": 2e8,
        "overlap_comm": True,
        "reduce_scatter": True,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": True
    }

    if stage >= 2:
        # ZeRO-2: 优化器状态分片
        zero_config["offload_optimizer"] = {
            "device": "cpu",
            "pin_memory": True
        }

    if stage >= 3:
        # ZeRO-3: 参数分片
        zero_config["offload_param"] = {
            "device": "cpu",
            "pin_memory": True
        }
        zero_config["param_persistence_threshold"] = 1e5

    return zero_config

class GradientCompressor:
    def __init__(self, compression_ratio=0.01):
        self.compression_ratio = compression_ratio
        self.residuals = {}

    def compress(self, tensor, name):
        """Top-k 稀疏化压缩"""
        # 添加残差
        if name in self.residuals:
            tensor = tensor + self.residuals[name]

        # 选择 top-k 元素
        numel = tensor.numel()
        k = max(1, int(numel * self.compression_ratio))

        values, indices = torch.topk(tensor.abs().view(-1), k)
        mask = torch.zeros_like(tensor.view(-1))
        mask[indices] = 1

        compressed = tensor.view(-1) * mask

        # 保存残差
        self.residuals[name] = tensor.view(-1) - compressed

        # 返回稀疏表示
        return indices, compressed[indices]

    def decompress(self, indices, values, shape):
        """解压缩"""
        tensor = torch.zeros(shape).view(-1)
        tensor[indices] = values
        return tensor.view(shape)

class CommunicationScheduler:
    def __init__(self, model):
        self.model = model
        self.comm_groups = self.create_comm_groups()

    def schedule_allreduce(self, gradients):
        """优化 AllReduce 调度"""
        # 按大小分组
        small_grads = []
        large_grads = []

        for name, grad in gradients.items():
            if grad.numel() < 1e6:
                small_grads.append((name, grad))
            else:
                large_grads.append((name, grad))

        # 小梯度合并通信
        if small_grads:
            merged = torch.cat([g.view(-1) for _, g in small_grads])
            handle = dist.all_reduce(merged, async_op=True)

        # 大梯度流水线通信
        handles = []
        for name, grad in large_grads:
            handle = dist.all_reduce(grad, async_op=True)
            handles.append(handle)

        return handles

class SelectiveCheckpointing:
    """选择性激活检查点"""

    def __init__(self, model, checkpoint_ratio=0.5):
        self.model = model
        self.checkpoint_ratio = checkpoint_ratio
        self.setup_checkpointing()

    def setup_checkpointing(self):
        """配置哪些层使用检查点"""
        total_layers = len(self.model.layers)
        checkpoint_layers = int(total_layers * self.checkpoint_ratio)

        # 选择内存占用大的层
        layer_memories = []
        for i, layer in enumerate(self.model.layers):
            memory = self.estimate_layer_memory(layer)
            layer_memories.append((i, memory))

        # 优先检查点内存占用大的层
        layer_memories.sort(key=lambda x: x[1], reverse=True)

        for i, _ in layer_memories[:checkpoint_layers]:
            self.model.layers[i].use_checkpoint = True

    def estimate_layer_memory(self, layer):
        """估算层的激活内存"""
        # 简化估算
        return sum(p.numel() for p in layer.parameters())

class MemoryPool:
    """预分配内存池减少碎片"""

    def __init__(self, pool_size=1024**3):  # 1GB
        self.pool = torch.cuda.ByteTensor(pool_size)
        self.allocated = 0
        self.allocations = {}

    def allocate(self, size, name):
        """从池中分配内存"""
        if self.allocated + size > len(self.pool):
            raise RuntimeError("Memory pool exhausted")

        start = self.allocated
        self.allocated += size
        self.allocations[name] = (start, size)

        return self.pool[start:start+size].view(size)

    def free(self, name):
        """释放内存（逻辑释放）"""
        if name in self.allocations:
            del self.allocations[name]
            # 实际实现中需要内存整理

class RobustCheckpointing:
    def __init__(self, checkpoint_dir, keep_last_n=3):
        self.checkpoint_dir = checkpoint_dir
        self.keep_last_n = keep_last_n
        self.checkpoint_history = []

    def save_checkpoint(self, model, optimizer, epoch, metrics):
        """保存检查点with原子操作"""
        checkpoint = {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "epoch": epoch,
            "metrics": metrics,
            "timestamp": time.time()
        }

        # 临时文件
        temp_path = f"{self.checkpoint_dir}/temp_ckpt_{epoch}.pt"
        final_path = f"{self.checkpoint_dir}/checkpoint_{epoch}.pt"

        # 原子写入
        torch.save(checkpoint, temp_path)

        # 验证检查点
        if self.verify_checkpoint(temp_path):
            os.rename(temp_path, final_path)
            self.checkpoint_history.append(final_path)

            # 清理旧检查点
            if len(self.checkpoint_history) > self.keep_last_n:
                old_ckpt = self.checkpoint_history.pop(0)
                os.remove(old_ckpt)
        else:
            os.remove(temp_path)
            raise RuntimeError("Checkpoint verification failed")

    def verify_checkpoint(self, path):
        """验证检查点完整性"""
        try:
            checkpoint = torch.load(path, map_location="cpu")
            required_keys = ["model_state_dict", "optimizer_state_dict", "epoch"]
            return all(k in checkpoint for k in required_keys)
        except:
            return False

class ElasticTrainer:
    """支持动态节点增减的弹性训练"""

    def __init__(self, min_nodes=1, max_nodes=8):
        self.min_nodes = min_nodes
        self.max_nodes = max_nodes
        self.current_nodes = self.detect_available_nodes()

    def handle_node_failure(self, failed_node):
        """处理节点故障"""
        self.current_nodes.remove(failed_node)

        if len(self.current_nodes) < self.min_nodes:
            # 等待新节点或终止
            self.wait_for_nodes()
        else:
            # 重新配置并继续
            self.reconfigure_training()

    def reconfigure_training(self):
        """重新配置训练"""
        # 重新初始化进程组
        dist.destroy_process_group()
        dist.init_process_group(
            backend="nccl",
            world_size=len(self.current_nodes)
        )

        # 调整批次大小保持全局批次不变
        global_batch_size = self.config.global_batch_size
        self.config.local_batch_size = (
            global_batch_size // len(self.current_nodes)
        )

        # 重新分配数据
        self.redistribute_data()

class PerformanceMonitor:
    def __init__(self):
        self.metrics = defaultdict(list)
        self.profiler = None

    def profile_iteration(self, iteration):
        """分析单次迭代"""
        with torch.profiler.profile(
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
            record_shapes=True,
            profile_memory=True,
            with_stack=True
        ) as prof:
            # 执行训练步骤
            yield

        # 分析结果
        self.analyze_profile(prof, iteration)

    def analyze_profile(self, prof, iteration):
        """分析性能瓶颈"""
        # GPU利用率
        cuda_time = sum([
            item.cuda_time_total for item in prof.key_averages()
        ])

        # 通信时间
        comm_time = sum([
            item.cuda_time_total for item in prof.key_averages()
            if "nccl" in item.key or "allreduce" in item.key
        ])

        # 内存峰值
        memory_peak = torch.cuda.max_memory_allocated()

        self.metrics["gpu_utilization"].append(cuda_time)
        self.metrics["comm_overhead"].append(comm_time / cuda_time)
        self.metrics["memory_peak"].append(memory_peak)

        # 识别瓶颈
        if comm_time / cuda_time > 0.3:
            print(f"Warning: High communication overhead ({comm_time/cuda_time:.2%})")

class AutoTuner:
    """自动调优训练配置"""

    def __init__(self, model, initial_config):
        self.model = model
        self.config = initial_config
        self.performance_history = []

    def auto_tune(self, num_trials=10):
        """自动调优"""
        best_throughput = 0
        best_config = self.config.copy()

        for trial in range(num_trials):
            # 生成新配置
            trial_config = self.generate_config_variant()

            # 测试性能
            throughput = self.benchmark_config(trial_config)

            if throughput > best_throughput:
                best_throughput = throughput
                best_config = trial_config

            self.performance_history.append({
                "config": trial_config,
                "throughput": throughput
            })

        return best_config

    def generate_config_variant(self):
        """生成配置变体"""
        config = self.config.copy()

        # 调整关键参数
        params_to_tune = [
            ("micro_batch_size", [1, 2, 4, 8]),
            ("gradient_accumulation_steps", [1, 2, 4, 8, 16]),
            ("num_workers", [0, 2, 4, 8]),
            ("pin_memory", [True, False]),
            ("prefetch_factor", [2, 4, 8])
        ]

        for param, values in params_to_tune:
            if random.random() < 0.3:  # 30% 概率修改
                config[param] = random.choice(values)

        return config