"""
简化的Megatron风格张量并行实现

注意：这是教学示例，实际生产应使用Megatron-LM或DeepSpeed
"""

import torch
import torch.nn as nn
import torch.distributed as dist

class TensorParallelLinear(nn.Module):
    """
    张量并行的线性层
    
    支持列并行和行并行
    """
    
    def __init__(self, in_features, out_features, parallel_mode='column'):
        super().__init__()
        self.parallel_mode = parallel_mode
        self.world_size = dist.get_world_size() if dist.is_initialized() else 1
        self.rank = dist.get_rank() if dist.is_initialized() else 0
        
        assert out_features % self.world_size == 0
        
        if parallel_mode == 'column':
            # 列并行：分割输出维度
            self.local_out = out_features // self.world_size
            self.weight = nn.Parameter(torch.randn(in_features, self.local_out))
        elif parallel_mode == 'row':
            # 行并行：分割输入维度
            assert in_features % self.world_size == 0
            local_in = in_features // self.world_size
            self.weight = nn.Parameter(torch.randn(local_in, out_features))
        
        self.reset_parameters()
    
    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
    
    def forward(self, x):
        if self.parallel_mode == 'column':
            # 列并行：局部计算，然后AllGather
            local_output = x @ self.weight
            # 收集所有部分的输出
            outputs = [torch.empty_like(local_output) for _ in range(self.world_size)]
            dist.all_gather(outputs, local_output)
            return torch.cat(outputs, dim=-1)
        
        elif self.parallel_mode == 'row':
            # 行并行：先Scatter输入，然后局部计算并Reduce
            input_chunks = torch.chunk(x, self.world_size, dim=-1)
            local_input = input_chunks[self.rank]
            local_output = local_input @ self.weight
            # Reduce散射的结果
            dist.all_reduce(local_output)
            return local_output


class TensorParallelTransformerLayer(nn.Module):
    """
    简化的张量并行Transformer层
    """
    
    def __init__(self, hidden_size, num_heads):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        
        # QKV投影（列并行）
        self.qkv_proj = TensorParallelLinear(hidden_size, hidden_size * 3, 'column')
        
        # 输出投影（行并行）
        self.out_proj = TensorParallelLinear(hidden_size, hidden_size, 'row')
        
        # MLP（列+行并行）
        self.mlp_fc1 = TensorParallelLinear(hidden_size, hidden_size * 4, 'column')
        self.mlp_fc2 = TensorParallelLinear(hidden_size * 4, hidden_size, 'row')
        
        self.ln1 = nn.LayerNorm(hidden_size)
        self.ln2 = nn.LayerNorm(hidden_size)
        self.activation = nn.GELU()
    
    def forward(self, x):
        # 自注意力
        residual = x
        x = self.ln1(x)
        
        # QKV
        qkv = self.qkv_proj(x)
        q, k, v = torch.chunk(qkv, 3, dim=-1)
        
        # 简化注意力计算
        # 实际实现需要更复杂的逻辑
        attn_output = self.attention(q, k, v)
        
        # 输出投影
        x = self.out_proj(attn_output)
        x = x + residual
        
        # MLP
        residual = x
        x = self.ln2(x)
        x = self.mlp_fc1(x)
        x = self.activation(x)
        x = self.mlp_fc2(x)
        x = x + residual
        
        return x
    
    def attention(self, q, k, v):
        """简化的注意力计算"""
        scale = 1 / math.sqrt(q.size(-1))
        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
        weights = torch.softmax(scores, dim=-1)
        return torch.matmul(weights, v)

"""
简化的ZeRO优化器实现

注意：这是教学示例，实际应使用DeepSpeed的ZeRO
"""

import torch

class ZeroOptimizer:
    """
    ZeRO-1优化器的简化实现
    
    将优化器状态分片到不同进程
    """
    
    def __init__(self, model, optimizer_class=torch.optim.Adam, lr=0.001):
        self.model = model
        self.rank = dist.get_rank()
        self.world_size = dist.get_world_size()
        
        # 将参数分组
        self.param_groups = self._partition_parameters()
        
        # 每个进程只优化自己的参数组
        local_params = self.param_groups[self.rank]
        self.optimizer = optimizer_class(local_params, lr=lr)
    
    def _partition_parameters(self):
        """将参数均匀分配到各进程"""
        all_params = list(self.model.parameters())
        partitions = [[] for _ in range(self.world_size)]
        
        for i, param in enumerate(all_params):
            partitions[i % self.world_size].append(param)
        
        return partitions
    
    def step(self):
        """执行优化步骤"""
        # 本地更新
        self.optimizer.step()
        
        # 广播更新后的参数到其他进程
        for rank in range(self.world_size):
            if rank != self.rank:
                # 从拥有该参数组的进程广播
                pass  # 简化处理
    
    def zero_grad(self):
        """清零梯度"""
        self.optimizer.zero_grad()


class ZeroStage3Optimizer:
    """
    ZeRO-3优化器的简化概念
    
    参数也进行分片，使用时动态获取
    """
    
    def __init__(self, model, lr=0.001):
        self.model = model
        self.rank = dist.get_rank()
        self.world_size = dist.get_world_size()
        
        # 参数分片
        self.param_shards = self._shard_parameters()
        
        # 本地优化器
        self.optimizer = torch.optim.Adam(self.param_shards, lr=lr)
    
    def _shard_parameters(self):
        """将参数分片"""
        all_params = list(self.model.parameters())
        shards = []
        
        for i, param in enumerate(all_params):
            if i % self.world_size == self.rank:
                shards.append(param)
        
        return shards
    
    def gather_params_for_forward(self, layer):
        """在前向传播前收集所需参数"""
        # 从其他进程获取缺失的参数
        pass
    
    def release_unused_params(self):
        """释放不再需要的参数以节省显存"""
        pass
    
    def step(self):
        self.optimizer.step()
        self.release_unused_params()

"""
大模型训练稳定性监控工具
"""

import torch
import numpy as np
from collections import deque

class TrainingMonitor:
    """
    训练监控器
    
    跟踪关键指标并检测异常
    """
    
    def __init__(self, window_size=100):
        self.loss_history = deque(maxlen=window_size)
        self.grad_norm_history = deque(maxlen=window_size)
        self.param_norm_history = deque(maxlen=window_size)
        
        self.alert_thresholds = {
            'loss_spike': 3.0,  # loss突增倍数
            'grad_explosion': 100.0,  # 梯度范数阈值
            'nan_detected': True,
        }
    
    def record_step(self, loss, grad_norm=None, param_norm=None):
        """记录一步训练的数据"""
        self.loss_history.append(loss.item())
        
        if grad_norm is not None:
            self.grad_norm_history.append(grad_norm)
        
        if param_norm is not None:
            self.param_norm_history.append(param_norm)
    
    def check_anomalies(self):
        """检查是否有异常"""
        alerts = []
        
        # 检查loss spike
        if len(self.loss_history) > 10:
            recent_avg = np.mean(list(self.loss_history)[-10:])
            previous_avg = np.mean(list(self.loss_history)[:-10])
            
            if previous_avg > 0 and recent_avg / previous_avg > self.alert_thresholds['loss_spike']:
                alerts.append(f"Loss spike detected: {previous_avg:.4f} -> {recent_avg:.4f}")
        
        # 检查梯度爆炸
        if self.grad_norm_history:
            latest_grad = self.grad_norm_history[-1]
            if latest_grad > self.alert_thresholds['grad_explosion']:
                alerts.append(f"Gradient explosion: norm={latest_grad:.2f}")
        
        # 检查NaN
        if self.loss_history and np.isnan(list(self.loss_history)[-1]):
            alerts.append("NaN detected in loss!")
        
        return alerts
    
    def get_statistics(self):
        """获取统计信息"""
        stats = {}
        
        if self.loss_history:
            losses = list(self.loss_history)
            stats['loss_mean'] = np.mean(losses)
            stats['loss_std'] = np.std(losses)
            stats['loss_min'] = np.min(losses)
            stats['loss_max'] = np.max(losses)
        
        if self.grad_norm_history:
            grads = list(self.grad_norm_history)
            stats['grad_norm_mean'] = np.mean(grads)
            stats['grad_norm_max'] = np.max(grads)
        
        return stats
    
    def plot_history(self):
        """绘制历史曲线"""
        try:
            import matplotlib.pyplot as plt
            
            fig, axes = plt.subplots(1, 3, figsize=(15, 4))
            
            if self.loss_history:
                axes[0].plot(list(self.loss_history))
                axes[0].set_title('Loss History')
                axes[0].set_xlabel('Step')
                axes[0].set_ylabel('Loss')
            
            if self.grad_norm_history:
                axes[1].plot(list(self.grad_norm_history))
                axes[1].set_title('Gradient Norm History')
                axes[1].set_xlabel('Step')
                axes[1].set_ylabel('Grad Norm')
            
            if self.param_norm_history:
                axes[2].plot(list(self.param_norm_history))
                axes[2].set_title('Parameter Norm History')
                axes[2].set_xlabel('Step')
                axes[2].set_ylabel('Param Norm')
            
            plt.tight_layout()
            plt.savefig('training_monitor.png')
            print("Plot saved to training_monitor.png")
            
        except ImportError:
            print("matplotlib not available")


def gradient_clipping(model, max_norm=1.0):
    """
    梯度裁剪
    
    Args:
        model: 模型
        max_norm: 最大梯度范数
    """
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_norm)


def mixed_precision_stability_check(loss_scaler, scaler_factor=2.0):
    """
    检查混合精度训练的稳定性
    
    如果频繁出现overflow，可能需要调整loss scale
    """
    if loss_scaler.get_scale() < 1e-10:
        print("Warning: Loss scale is very low, consider adjusting")
        return False
    return True


# 使用示例
def training_loop_with_monitoring():
    """带监控的训练循环示例"""
    
    monitor = TrainingMonitor()
    
    # 模拟训练
    for step in range(1000):
        # 模拟损失和梯度
        loss = torch.tensor(1.0 / (step + 1) + np.random.randn() * 0.1)
        grad_norm = torch.randn(1).item() * 10
        
        monitor.record_step(loss, grad_norm)
        
        # 检查异常
        if step % 100 == 0:
            alerts = monitor.check_anomalies()
            if alerts:
                for alert in alerts:
                    print(f"ALERT: {alert}")
            
            stats = monitor.get_statistics()
            print(f"Step {step}: {stats}")
    
    monitor.plot_history()


if __name__ == "__main__":
    training_loop_with_monitoring()