llm_parallel_simulator.py

"""
LLM并行计算模拟器


支持的并行策略:
- 数据并行 (Data Parallel): DDP, FSDP
- 张量并行 (Tensor Parallel): Megatron-LM风格
- 流水线并行 (Pipeline Parallel): GPipe, PipeDream
- 序列并行 (Sequence Parallel): 长序列处理
- 专家并行 (Expert Parallel): MoE并行
- 3D并行 (3D Parallel): 混合并行策略
"""

import time
import random
import math
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, field
from enum import Enum
from abc import ABC, abstractmethod

try:
    import matplotlib.pyplot as plt
    import matplotlib.patches as patches
    import numpy as np
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False
    print("警告: matplotlib/numpy未安装，将跳过可视化功能")

try:
    from tqdm import tqdm
    HAS_TQDM = True
except ImportError:
    HAS_TQDM = False
    print("警告: tqdm未安装，将使用简单进度显示")

# ======================
# LLM模型架构定义
# ======================

@dataclass
class TransformerConfig:
    """Transformer模型配置"""
    # 模型结构参数
    vocab_size: int = 50000
    hidden_size: int = 4096
    num_layers: int = 32
    num_attention_heads: int = 32
    intermediate_size: int = 11008  # FFN中间层大小
    max_sequence_length: int = 2048
    
    # 训练参数
    batch_size: int = 8
    micro_batch_size: int = 2
    gradient_accumulation_steps: int = 4
    
    # 内存和计算参数
    dtype_bytes: int = 2  # FP16
    activation_checkpointing: bool = True
    
    def __post_init__(self):
        """计算派生参数"""
        self.head_dim = self.hidden_size // self.num_attention_heads
        self.total_params = self.estimate_parameters()
        self.model_memory_gb = self.estimate_model_memory()
    
    def estimate_parameters(self) -> int:
        """估算模型参数量"""
        # Embedding层
        embedding_params = self.vocab_size * self.hidden_size
        
        # Transformer层参数
        # Attention: Q,K,V投影 + 输出投影
        attention_params = 4 * self.hidden_size * self.hidden_size
        # FFN: 上投影 + 下投影
        ffn_params = 2 * self.hidden_size * self.intermediate_size
        # LayerNorm: 2个LayerNorm层
        layernorm_params = 2 * self.hidden_size
        
        layer_params = attention_params + ffn_params + layernorm_params
        total_layer_params = self.num_layers * layer_params
        
        # 输出层
        output_params = self.hidden_size * self.vocab_size
        
        return embedding_params + total_layer_params + output_params
    
    def estimate_model_memory(self) -> float:
        """估算模型内存使用(GB)"""
        return self.total_params * self.dtype_bytes / (1024**3)

@dataclass
class LLMWorkload:
    """LLM工作负载定义"""
    model_config: TransformerConfig
    sequence_length: int
    batch_size: int
    
    def __post_init__(self):
        """计算工作负载特征"""
        self.tokens_per_batch = self.batch_size * self.sequence_length
        self.activation_memory_gb = self.estimate_activation_memory()
        self.compute_flops = self.estimate_compute_flops()
    
    def estimate_activation_memory(self) -> float:
        """估算激活内存使用(GB)"""
        # 每个token的激活内存
        hidden_states = self.model_config.hidden_size
        attention_scores = self.model_config.num_attention_heads * self.sequence_length
        
        activation_per_token = (hidden_states + attention_scores) * self.model_config.dtype_bytes
        total_activation = activation_per_token * self.tokens_per_batch * self.model_config.num_layers
        
        # 如果使用激活检查点，减少内存使用
        if self.model_config.activation_checkpointing:
            total_activation *= 0.5
            
        return total_activation / (1024**3)
    
    def estimate_compute_flops(self) -> int:
        """估算计算FLOPs"""
        # 前向传播FLOPs估算
        # Attention: O(n²d + nd²)
        attention_flops = (
            self.sequence_length ** 2 * self.model_config.hidden_size +
            self.sequence_length * self.model_config.hidden_size ** 2
        ) * self.model_config.num_layers
        
        # FFN: O(nd_ff)
        ffn_flops = (
            self.sequence_length * self.model_config.hidden_size * self.model_config.intermediate_size * 2
        ) * self.model_config.num_layers
        
        return (attention_flops + ffn_flops) * self.batch_size

# ======================
# 设备和通信模拟
# ======================

@dataclass
class DeviceSpec:
    """设备规格"""
    device_id: int
    memory_gb: float = 80.0  # A100 80GB
    compute_tflops: float = 312.0  # A100 BF16
    memory_bandwidth_gbps: float = 2000.0  # HBM带宽
    
class CommunicationTopology:
    """通信拓扑模拟"""
    
    def __init__(self, num_devices: int, topology_type: str = "dgx"):
        self.num_devices = num_devices
        self.topology_type = topology_type
        self.devices = [DeviceSpec(i) for i in range(num_devices)]
        self.bandwidth_matrix = self._create_bandwidth_matrix()
    
    def _create_bandwidth_matrix(self) -> List[List[float]]:
        """创建设备间带宽矩阵(GB/s)"""
        matrix = [[0.0] * self.num_devices for _ in range(self.num_devices)]
        
        if self.topology_type == "dgx":
            # DGX A100拓扑: NVLink连接
            nvlink_bandwidth = 600.0  # GB/s
            infiniband_bandwidth = 200.0  # GB/s
            
            for i in range(self.num_devices):
                for j in range(self.num_devices):
                    if i == j:
                        matrix[i][j] = float('inf')  # 本地内存
                    elif abs(i - j) <= 2:  # 相邻设备用NVLink
                        matrix[i][j] = nvlink_bandwidth
                    else:  # 远程设备用InfiniBand
                        matrix[i][j] = infiniband_bandwidth
        else:
            # 默认全连接拓扑
            default_bandwidth = 100.0
            for i in range(self.num_devices):
                for j in range(self.num_devices):
                    if i == j:
                        matrix[i][j] = float('inf')
                    else:
                        matrix[i][j] = default_bandwidth
        
        return matrix
    
    def get_communication_time(self, src: int, dst: int, data_size_gb: float) -> float:
        """计算通信时间(秒)"""
        if src == dst:
            return 0.0
        bandwidth = self.bandwidth_matrix[src][dst]
        return data_size_gb / bandwidth

@dataclass
class PerformanceMetrics:
    """性能指标"""
    strategy_name: str
    total_time: float
    compute_time: float
    communication_time: float
    memory_usage_per_device: float
    peak_memory_usage: float
    throughput_tokens_per_sec: float
    model_flops_utilization: float
    communication_efficiency: float
    scalability_efficiency: float
    details: Dict[str, Any] = field(default_factory=dict)
    
    @property
    def speedup(self) -> float:
        """相对于基准的加速比"""
        baseline_time = self.details.get('baseline_time', self.total_time)
        return baseline_time / self.total_time if self.total_time > 0 else 1.0

# ======================
# 内存管理器
# ======================

class MemoryManager:
    """内存管理器"""
    
    def __init__(self, devices: List[DeviceSpec]):
        self.devices = devices
        self.memory_usage = {dev.device_id: 0.0 for dev in devices}
        self.peak_usage = {dev.device_id: 0.0 for dev in devices}
    
    def allocate(self, device_id: int, size_gb: float) -> bool:
        """分配内存"""
        if device_id >= len(self.devices):
            return False
        
        device = self.devices[device_id]
        if self.memory_usage[device_id] + size_gb <= device.memory_gb:
            self.memory_usage[device_id] += size_gb
            self.peak_usage[device_id] = max(
                self.peak_usage[device_id], 
                self.memory_usage[device_id]
            )
            return True
        return False
    
    def deallocate(self, device_id: int, size_gb: float):
        """释放内存"""
        if device_id < len(self.devices):
            self.memory_usage[device_id] = max(0, self.memory_usage[device_id] - size_gb)
    
    def get_memory_utilization(self, device_id: int) -> float:
        """获取内存利用率"""
        if device_id >= len(self.devices):
            return 0.0
        return self.memory_usage[device_id] / self.devices[device_id].memory_gb
    
    def reset(self):
        """重置内存使用"""
        for device_id in self.memory_usage:
            self.memory_usage[device_id] = 0.0
            self.peak_usage[device_id] = 0.0

# ======================
# 并行策略抽象基类
# ======================

class ParallelStrategy(ABC):
    """并行策略抽象基类"""

    def __init__(self,
                 workload: LLMWorkload,
                 topology: CommunicationTopology,
                 memory_manager: MemoryManager):
        self.workload = workload
        self.topology = topology
        self.memory_manager = memory_manager
        self.strategy_name = self.__class__.__name__

    @abstractmethod
    def setup(self) -> bool:
        """设置并行策略，返回是否成功"""
        pass

    @abstractmethod
    def forward_pass(self) -> Tuple[float, float]:
        """执行前向传播，返回(计算时间, 通信时间)"""
        pass

    @abstractmethod
    def backward_pass(self) -> Tuple[float, float]:
        """执行反向传播，返回(计算时间, 通信时间)"""
        pass

    @abstractmethod
    def get_memory_usage(self) -> Dict[int, float]:
        """获取各设备内存使用情况"""
        pass

    def simulate_compute(self, flops: int, device_id: int) -> float:
        """模拟计算时间"""
        device = self.topology.devices[device_id]
        # 考虑设备利用率和随机性
        utilization = 0.8 + 0.2 * random.random()
        compute_time = flops / (device.compute_tflops * 1e12) / utilization

        # 添加少量实际等待时间用于演示
        time.sleep(compute_time * 0.001)
        return compute_time

    def simulate_communication(self, src: int, dst: int, data_size_gb: float) -> float:
        """模拟通信时间"""
        comm_time = self.topology.get_communication_time(src, dst, data_size_gb)

        # 添加通信开销和随机性
        overhead = 1.1 + 0.1 * random.random()
        actual_time = comm_time * overhead

        # 添加少量实际等待时间用于演示
        time.sleep(actual_time * 0.001)
        return actual_time

    def run_simulation(self) -> PerformanceMetrics:
        """运行完整的训练模拟"""
        print(f"\n🔄 运行 {self.strategy_name} 模拟...")

        # 设置策略
        if not self.setup():
            raise RuntimeError(f"Failed to setup {self.strategy_name}")

        start_time = time.time()

        # 前向传播
        forward_compute, forward_comm = self.forward_pass()

        # 反向传播
        backward_compute, backward_comm = self.backward_pass()

        total_time = time.time() - start_time
        total_compute = forward_compute + backward_compute
        total_comm = forward_comm + backward_comm

        # 计算性能指标
        memory_usage = self.get_memory_usage()
        peak_memory = max(memory_usage.values()) if memory_usage else 0.0
        avg_memory = sum(memory_usage.values()) / len(memory_usage) if memory_usage else 0.0

        throughput = self.workload.tokens_per_batch / total_time
        flops_utilization = self.workload.compute_flops / (total_compute * sum(
            dev.compute_tflops * 1e12 for dev in self.topology.devices
        )) if total_compute > 0 else 0.0

        comm_efficiency = total_compute / (total_compute + total_comm) if (total_compute + total_comm) > 0 else 1.0

        # 可扩展性效率（相对于理想线性扩展）
        ideal_time = total_time * len(self.topology.devices)
        scalability_efficiency = ideal_time / (total_time * len(self.topology.devices)) if total_time > 0 else 1.0

        metrics = PerformanceMetrics(
            strategy_name=self.strategy_name,
            total_time=total_time,
            compute_time=total_compute,
            communication_time=total_comm,
            memory_usage_per_device=avg_memory,
            peak_memory_usage=peak_memory,
            throughput_tokens_per_sec=throughput,
            model_flops_utilization=flops_utilization,
            communication_efficiency=comm_efficiency,
            scalability_efficiency=scalability_efficiency,
            details={
                'memory_per_device': memory_usage,
                'forward_compute': forward_compute,
                'forward_comm': forward_comm,
                'backward_compute': backward_compute,
                'backward_comm': backward_comm,
                'num_devices': len(self.topology.devices)
            }
        )

        print(f"✅ {self.strategy_name} 完成: {total_time:.3f}秒, "
              f"吞吐量: {throughput:.1f} tokens/s, "
              f"通信效率: {comm_efficiency:.2%}")

        return metrics

# ======================
# 基准策略（单设备）
# ======================

class BaselineStrategy(ParallelStrategy):
    """基准策略：单设备训练"""

    def setup(self) -> bool:
        """设置单设备训练"""
        # 检查单设备是否能容纳模型和激活
        model_memory = self.workload.model_config.model_memory_gb
        activation_memory = self.workload.activation_memory_gb
        total_memory = model_memory + activation_memory

        if not self.memory_manager.allocate(0, total_memory):
            print(f"❌ 单设备内存不足: 需要 {total_memory:.1f}GB")
            return False

        print(f"📊 单设备配置: 模型 {model_memory:.1f}GB + 激活 {activation_memory:.1f}GB")
        return True

    def forward_pass(self) -> Tuple[float, float]:
        """前向传播"""
        # 计算前向传播FLOPs
        forward_flops = self.workload.compute_flops // 2  # 前向约占一半
        compute_time = self.simulate_compute(forward_flops, 0)
        return compute_time, 0.0  # 单设备无通信

    def backward_pass(self) -> Tuple[float, float]:
        """反向传播"""
        # 计算反向传播FLOPs（约为前向的2倍）
        backward_flops = self.workload.compute_flops
        compute_time = self.simulate_compute(backward_flops, 0)
        return compute_time, 0.0  # 单设备无通信

    def get_memory_usage(self) -> Dict[int, float]:
        """获取内存使用"""
        return {0: self.memory_manager.memory_usage[0]}

# ======================
# 数据并行策略
# ======================

class DataParallelStrategy(ParallelStrategy):
    """数据并行策略：DDP风格"""

    def __init__(self, workload: LLMWorkload, topology: CommunicationTopology,
                 memory_manager: MemoryManager, dp_type: str = "ddp"):
        super().__init__(workload, topology, memory_manager)
        self.dp_type = dp_type  # "ddp" or "fsdp"
        self.num_devices = len(topology.devices)
        self.micro_batch_per_device = workload.batch_size // self.num_devices

    def setup(self) -> bool:
        """设置数据并行"""
        model_memory = self.workload.model_config.model_memory_gb

        if self.dp_type == "fsdp":
            # FSDP: 模型参数分片
            model_memory_per_device = model_memory / self.num_devices
            print(f"📊 FSDP配置: 模型分片 {model_memory_per_device:.1f}GB/设备")
        else:
            # DDP: 每个设备复制完整模型
            model_memory_per_device = model_memory
            print(f"📊 DDP配置: 完整模型 {model_memory_per_device:.1f}GB/设备")

        # 激活内存按设备分配
        activation_memory_per_device = self.workload.activation_memory_gb / self.num_devices
        total_memory_per_device = model_memory_per_device + activation_memory_per_device

        # 检查每个设备内存
        for device_id in range(self.num_devices):
            if not self.memory_manager.allocate(device_id, total_memory_per_device):
                print(f"❌ 设备 {device_id} 内存不足: 需要 {total_memory_per_device:.1f}GB")
                return False

        print(f"✅ 数据并行设置完成: {self.num_devices}设备, "
              f"微批次大小: {self.micro_batch_per_device}")
        return True

    def forward_pass(self) -> Tuple[float, float]:
        """前向传播"""
        # 每个设备处理自己的微批次
        forward_flops_per_device = (self.workload.compute_flops // 2) // self.num_devices

        # 并行计算（取最慢设备的时间）
        compute_times = []
        for device_id in range(self.num_devices):
            compute_time = self.simulate_compute(forward_flops_per_device, device_id)
            compute_times.append(compute_time)

        max_compute_time = max(compute_times)

        # 前向传播无需通信
        return max_compute_time, 0.0

    def backward_pass(self) -> Tuple[float, float]:
        """反向传播"""
        # 反向传播计算
        backward_flops_per_device = self.workload.compute_flops // self.num_devices

        compute_times = []
        for device_id in range(self.num_devices):
            compute_time = self.simulate_compute(backward_flops_per_device, device_id)
            compute_times.append(compute_time)

        max_compute_time = max(compute_times)

        # 梯度同步通信
        if self.dp_type == "fsdp":
            # FSDP: 分片梯度收集和重分布
            comm_time = self._simulate_fsdp_communication()
        else:
            # DDP: All-Reduce梯度同步
            comm_time = self._simulate_allreduce_communication()

        return max_compute_time, comm_time

    def _simulate_allreduce_communication(self) -> float:
        """模拟All-Reduce通信"""
        # 梯度大小等于模型参数大小
        gradient_size_gb = self.workload.model_config.model_memory_gb

        # All-Reduce算法：Ring All-Reduce
        # 通信量 = 2 * (N-1) / N * 数据大小，其中N是设备数
        comm_volume = 2 * (self.num_devices - 1) / self.num_devices * gradient_size_gb

        # 使用最慢链路的带宽
        min_bandwidth = min(
            self.topology.bandwidth_matrix[i][j]
            for i in range(self.num_devices)
            for j in range(self.num_devices)
            if i != j
        )

        comm_time = comm_volume / min_bandwidth

        # 模拟实际通信
        time.sleep(comm_time * 0.001)

        return comm_time

    def _simulate_fsdp_communication(self) -> float:
        """模拟FSDP通信"""
        # FSDP通信包括：
        # 1. All-Gather参数分片
        # 2. Reduce-Scatter梯度

        param_shard_size = self.workload.model_config.model_memory_gb / self.num_devices

        # All-Gather: 每个设备收集其他设备的参数分片
        allgather_volume = (self.num_devices - 1) * param_shard_size

        # Reduce-Scatter: 梯度归约并分散
        reduce_scatter_volume = (self.num_devices - 1) / self.num_devices * param_shard_size

        total_volume = allgather_volume + reduce_scatter_volume

        # 使用平均带宽
        avg_bandwidth = sum(
            self.topology.bandwidth_matrix[i][j]
            for i in range(self.num_devices)
            for j in range(self.num_devices)
            if i != j
        ) / (self.num_devices * (self.num_devices - 1))

        comm_time = total_volume / avg_bandwidth

        # 模拟实际通信
        time.sleep(comm_time * 0.001)

        return comm_time

    def get_memory_usage(self) -> Dict[int, float]:
        """获取内存使用"""
        return {i: self.memory_manager.memory_usage[i] for i in range(self.num_devices)}

# ======================
# 张量并行策略
# ======================

class TensorParallelStrategy(ParallelStrategy):
    """张量并行策略：Megatron-LM风格"""

    def __init__(self, workload: LLMWorkload, topology: CommunicationTopology,
                 memory_manager: MemoryManager, tp_degree: int = None):
        super().__init__(workload, topology, memory_manager)
        self.tp_degree = tp_degree or len(topology.devices)
        self.num_devices = len(topology.devices)

        if self.num_devices % self.tp_degree != 0:
            raise ValueError(f"设备数 {self.num_devices} 必须能被张量并行度 {self.tp_degree} 整除")

    def setup(self) -> bool:
        """设置张量并行"""
        # 张量并行：模型参数按张量维度分片
        model_memory_per_device = self.workload.model_config.model_memory_gb / self.tp_degree

        # 激活内存：每个设备处理完整批次，但中间激活分片
        activation_memory_per_device = self.workload.activation_memory_gb / self.tp_degree

        total_memory_per_device = model_memory_per_device + activation_memory_per_device

        # 检查每个设备内存
        for device_id in range(self.num_devices):
            if not self.memory_manager.allocate(device_id, total_memory_per_device):
                print(f"❌ 设备 {device_id} 内存不足: 需要 {total_memory_per_device:.1f}GB")
                return False

        print(f"📊 张量并行配置: {self.tp_degree}路并行, "
              f"模型分片 {model_memory_per_device:.1f}GB/设备")
        return True

    def forward_pass(self) -> Tuple[float, float]:
        """前向传播"""
        # 张量并行的前向传播包括多个通信点
        total_compute_time = 0.0
        total_comm_time = 0.0

        # 模拟每一层的计算和通信
        layers_per_group = self.workload.model_config.num_layers

        for layer_idx in range(layers_per_group):
            # Attention层
            attn_compute, attn_comm = self._simulate_attention_layer()

            # FFN层
            ffn_compute, ffn_comm = self._simulate_ffn_layer()

            total_compute_time += attn_compute + ffn_compute
            total_comm_time += attn_comm + ffn_comm

        return total_compute_time, total_comm_time

    def backward_pass(self) -> Tuple[float, float]:
        """反向传播"""
        # 反向传播的计算量约为前向的2倍
        forward_compute, forward_comm = self.forward_pass()

        # 反向传播有额外的梯度All-Reduce
        gradient_comm = self._simulate_gradient_allreduce()

        return forward_compute * 2, forward_comm + gradient_comm

    def _simulate_attention_layer(self) -> Tuple[float, float]:
        """模拟Attention层的张量并行"""
        # Attention计算：Q, K, V投影并行，输出投影需要All-Reduce

        # 计算时间：每个设备处理1/tp_degree的头
        seq_len = self.workload.sequence_length
        hidden_size = self.workload.model_config.hidden_size
        batch_size = self.workload.batch_size

        # QKV投影计算（并行）
        qkv_flops = 3 * seq_len * hidden_size * hidden_size * batch_size // self.tp_degree
        qkv_compute_time = self.simulate_compute(qkv_flops, 0)

        # Attention计算（并行）
        attn_flops = seq_len * seq_len * hidden_size * batch_size // self.tp_degree
        attn_compute_time = self.simulate_compute(attn_flops, 0)

        # 输出投影（需要All-Reduce）
        output_flops = seq_len * hidden_size * hidden_size * batch_size // self.tp_degree
        output_compute_time = self.simulate_compute(output_flops, 0)

        total_compute = qkv_compute_time + attn_compute_time + output_compute_time

        # 通信：输出投影后的All-Reduce
        output_size_gb = seq_len * hidden_size * batch_size * self.workload.model_config.dtype_bytes / (1024**3)
        comm_time = self._simulate_allreduce(output_size_gb)

        return total_compute, comm_time

    def _simulate_ffn_layer(self) -> Tuple[float, float]:
        """模拟FFN层的张量并行"""
        # FFN: 上投影并行，下投影需要All-Reduce

        seq_len = self.workload.sequence_length
        hidden_size = self.workload.model_config.hidden_size
        intermediate_size = self.workload.model_config.intermediate_size
        batch_size = self.workload.batch_size

        # 上投影（并行）
        up_flops = seq_len * hidden_size * intermediate_size * batch_size // self.tp_degree
        up_compute_time = self.simulate_compute(up_flops, 0)

        # 激活函数（并行）
        activation_flops = seq_len * intermediate_size * batch_size // self.tp_degree
        activation_compute_time = self.simulate_compute(activation_flops, 0)

        # 下投影（需要All-Reduce）
        down_flops = seq_len * intermediate_size * hidden_size * batch_size // self.tp_degree
        down_compute_time = self.simulate_compute(down_flops, 0)

        total_compute = up_compute_time + activation_compute_time + down_compute_time

        # 通信：下投影后的All-Reduce
        output_size_gb = seq_len * hidden_size * batch_size * self.workload.model_config.dtype_bytes / (1024**3)
        comm_time = self._simulate_allreduce(output_size_gb)

        return total_compute, comm_time

    def _simulate_allreduce(self, data_size_gb: float) -> float:
        """模拟张量并行组内的All-Reduce"""
        # Ring All-Reduce算法
        comm_volume = 2 * (self.tp_degree - 1) / self.tp_degree * data_size_gb

        # 使用张量并行组内的最小带宽
        min_bandwidth = min(
            self.topology.bandwidth_matrix[i][j]
            for i in range(self.tp_degree)
            for j in range(self.tp_degree)
            if i != j
        )

        comm_time = comm_volume / min_bandwidth
        time.sleep(comm_time * 0.001)

        return comm_time

    def _simulate_gradient_allreduce(self) -> float:
        """模拟梯度All-Reduce"""
        # 梯度大小等于模型参数分片大小
        gradient_size_gb = self.workload.model_config.model_memory_gb / self.tp_degree
        return self._simulate_allreduce(gradient_size_gb)

    def get_memory_usage(self) -> Dict[int, float]:
        """获取内存使用"""
        return {i: self.memory_manager.memory_usage[i] for i in range(self.num_devices)}

# ======================
# 流水线并行策略
# ======================

class PipelineParallelStrategy(ParallelStrategy):
    """流水线并行策略：GPipe/PipeDream风格"""

    def __init__(self, workload: LLMWorkload, topology: CommunicationTopology,
                 memory_manager: MemoryManager, pp_degree: int = None,
                 schedule: str = "gpipe"):
        super().__init__(workload, topology, memory_manager)
        self.pp_degree = pp_degree or len(topology.devices)
        self.num_devices = len(topology.devices)
        self.schedule = schedule  # "gpipe" or "pipedream"

        if self.num_devices % self.pp_degree != 0:
            raise ValueError(f"设备数 {self.num_devices} 必须能被流水线并行度 {self.pp_degree} 整除")

        # 计算微批次数量
        self.num_microbatches = max(4, self.pp_degree * 2)  # 通常是流水线深度的2-4倍
        self.microbatch_size = workload.batch_size // self.num_microbatches

        # 每个阶段的层数
        self.layers_per_stage = workload.model_config.num_layers // self.pp_degree

    def setup(self) -> bool:
        """设置流水线并行"""
        # 每个阶段只需要部分模型参数
        model_memory_per_stage = self.workload.model_config.model_memory_gb / self.pp_degree

        # 激活内存：需要存储多个微批次的激活
        activation_memory_per_stage = (
            self.workload.activation_memory_gb / self.pp_degree *
            self.num_microbatches / self.workload.batch_size * self.microbatch_size
        )

        total_memory_per_stage = model_memory_per_stage + activation_memory_per_stage

        # 检查每个设备内存
        for device_id in range(self.num_devices):
            if not self.memory_manager.allocate(device_id, total_memory_per_stage):
                print(f"❌ 设备 {device_id} 内存不足: 需要 {total_memory_per_stage:.1f}GB")
                return False

        print(f"📊 流水线并行配置: {self.pp_degree}阶段, {self.num_microbatches}微批次, "
              f"调度: {self.schedule}")
        return True

    def forward_pass(self) -> Tuple[float, float]:
        """前向传播"""
        if self.schedule == "gpipe":
            return self._simulate_gpipe_forward()
        else:
            return self._simulate_pipedream_forward()

    def backward_pass(self) -> Tuple[float, float]:
        """反向传播"""
        if self.schedule == "gpipe":
            return self._simulate_gpipe_backward()
        else:
            return self._simulate_pipedream_backward()

    def _simulate_gpipe_forward(self) -> Tuple[float, float]:
        """模拟GPipe前向传播"""
        # GPipe: 严格的前向-反向分离
        stage_compute_time = self._get_stage_compute_time() / 2  # 前向约占一半

        # 流水线执行时间线
        pipeline_timeline = []

        # 前向传播：填充流水线
        for microbatch_id in range(self.num_microbatches):
            for stage_id in range(self.pp_degree):
                start_time = microbatch_id * stage_compute_time + stage_id * stage_compute_time
                pipeline_timeline.append({
                    'stage': stage_id,
                    'microbatch': microbatch_id,
                    'phase': 'forward',
                    'start_time': start_time,
                    'duration': stage_compute_time
                })

        # 计算总的前向时间
        total_forward_time = (self.num_microbatches + self.pp_degree - 1) * stage_compute_time

        # 通信时间：阶段间激活传递
        activation_size_per_microbatch = (
            self.workload.sequence_length * self.workload.model_config.hidden_size *
            self.microbatch_size * self.workload.model_config.dtype_bytes / (1024**3)
        )

        total_comm_time = 0.0
        for stage_id in range(self.pp_degree - 1):
            comm_time = self.simulate_communication(
                stage_id, stage_id + 1,
                activation_size_per_microbatch * self.num_microbatches
            )
            total_comm_time += comm_time

        return total_forward_time, total_comm_time

    def _simulate_gpipe_backward(self) -> Tuple[float, float]:
        """模拟GPipe反向传播"""
        # 反向传播计算时间
        stage_compute_time = self._get_stage_compute_time()  # 反向约为前向的2倍

        # 反向传播：排空流水线
        total_backward_time = (self.num_microbatches + self.pp_degree - 1) * stage_compute_time

        # 通信时间：梯度传递
        gradient_size_per_microbatch = (
            self.workload.sequence_length * self.workload.model_config.hidden_size *
            self.microbatch_size * self.workload.model_config.dtype_bytes / (1024**3)
        )

        total_comm_time = 0.0
        for stage_id in range(1, self.pp_degree):
            comm_time = self.simulate_communication(
                stage_id, stage_id - 1,
                gradient_size_per_microbatch * self.num_microbatches
            )
            total_comm_time += comm_time

        return total_backward_time, total_comm_time

    def _simulate_pipedream_forward(self) -> Tuple[float, float]:
        """模拟PipeDream前向传播（1F1B调度）"""
        # PipeDream: 前向和反向交替执行
        stage_compute_time = self._get_stage_compute_time() / 2

        # 1F1B调度：每个阶段在前向后立即执行反向
        # 这里简化为前向时间，反向在backward_pass中计算
        warmup_time = self.pp_degree * stage_compute_time
        steady_time = (self.num_microbatches - self.pp_degree) * stage_compute_time

        total_time = warmup_time + steady_time

        # 通信时间类似GPipe但重叠更好
        activation_size = (
            self.workload.sequence_length * self.workload.model_config.hidden_size *
            self.microbatch_size * self.workload.model_config.dtype_bytes / (1024**3)
        )

        # 通信与计算重叠，实际通信时间更少
        comm_time = activation_size * self.num_microbatches / self.topology.bandwidth_matrix[0][1] * 0.5

        return total_time, comm_time

    def _simulate_pipedream_backward(self) -> Tuple[float, float]:
        """模拟PipeDream反向传播"""
        # 1F1B调度中反向与前向交替，总时间已在前向中计算
        stage_compute_time = self._get_stage_compute_time()

        # 反向传播时间
        backward_time = stage_compute_time * self.num_microbatches

        # 梯度通信
        gradient_size = (
            self.workload.model_config.model_memory_gb / self.pp_degree / self.num_microbatches
        )

        comm_time = gradient_size * self.num_microbatches / self.topology.bandwidth_matrix[0][1] * 0.3

        return backward_time, comm_time

    def _get_stage_compute_time(self) -> float:
        """计算每个阶段的计算时间"""
        # 每个阶段处理部分层和一个微批次
        stage_flops = (
            self.workload.compute_flops / self.workload.batch_size * self.microbatch_size / self.pp_degree
        )

        return self.simulate_compute(stage_flops, 0)

    def get_memory_usage(self) -> Dict[int, float]:
        """获取内存使用"""
        return {i: self.memory_manager.memory_usage[i] for i in range(self.num_devices)}

# ======================
# 序列并行策略
# ======================

class SequenceParallelStrategy(ParallelStrategy):
    """序列并行策略：长序列处理"""

    def __init__(self, workload: LLMWorkload, topology: CommunicationTopology,
                 memory_manager: MemoryManager, sp_degree: int = None):
        super().__init__(workload, topology, memory_manager)
        self.sp_degree = sp_degree or len(topology.devices)
        self.num_devices = len(topology.devices)

        # 序列长度必须能被并行度整除
        if workload.sequence_length % self.sp_degree != 0:
            raise ValueError(f"序列长度 {workload.sequence_length} 必须能被序列并行度 {self.sp_degree} 整除")

        self.seq_len_per_device = workload.sequence_length // self.sp_degree

    def setup(self) -> bool:
        """设置序列并行"""
        # 每个设备存储完整模型但只处理部分序列
        model_memory_per_device = self.workload.model_config.model_memory_gb

        # 激活内存按序列长度分片
        activation_memory_per_device = self.workload.activation_memory_gb / self.sp_degree

        total_memory_per_device = model_memory_per_device + activation_memory_per_device

        # 检查每个设备内存
        for device_id in range(self.num_devices):
            if not self.memory_manager.allocate(device_id, total_memory_per_device):
                print(f"❌ 设备 {device_id} 内存不足: 需要 {total_memory_per_device:.1f}GB")
                return False

        print(f"📊 序列并行配置: {self.sp_degree}路并行, "
              f"每设备序列长度: {self.seq_len_per_device}")
        return True

    def forward_pass(self) -> Tuple[float, float]:
        """前向传播"""
        total_compute_time = 0.0
        total_comm_time = 0.0

        # 模拟每一层的序列并行计算
        for layer_idx in range(self.workload.model_config.num_layers):
            layer_compute, layer_comm = self._simulate_sequence_parallel_layer()
            total_compute_time += layer_compute
            total_comm_time += layer_comm

        return total_compute_time, total_comm_time

    def backward_pass(self) -> Tuple[float, float]:
        """反向传播"""
        # 反向传播的计算量约为前向的2倍
        forward_compute, forward_comm = self.forward_pass()

        # 梯度All-Reduce
        gradient_comm = self._simulate_gradient_allreduce()

        return forward_compute * 2, forward_comm + gradient_comm

    def _simulate_sequence_parallel_layer(self) -> Tuple[float, float]:
        """模拟序列并行层"""
        # 每个设备处理部分序列长度
        hidden_size = self.workload.model_config.hidden_size
        batch_size = self.workload.batch_size

        # Attention计算：需要All-Gather序列维度
        # 1. All-Gather输入
        input_size_gb = (
            self.seq_len_per_device * hidden_size * batch_size *
            self.workload.model_config.dtype_bytes / (1024**3)
        )
        allgather_time = self._simulate_allgather(input_size_gb)

        # 2. 本地Attention计算（处理完整序列）
        attn_flops = (
            self.workload.sequence_length * self.workload.sequence_length *
            hidden_size * batch_size // self.sp_degree
        )
        attn_compute_time = self.simulate_compute(attn_flops, 0)

        # 3. Reduce-Scatter输出
        output_size_gb = input_size_gb
        reduce_scatter_time = self._simulate_reduce_scatter(output_size_gb)

        # FFN计算（本地，无通信）
        ffn_flops = (
            self.seq_len_per_device * hidden_size *
            self.workload.model_config.intermediate_size * batch_size * 2
        )
        ffn_compute_time = self.simulate_compute(ffn_flops, 0)

        total_compute = attn_compute_time + ffn_compute_time
        total_comm = allgather_time + reduce_scatter_time

        return total_compute, total_comm

    def _simulate_allgather(self, data_size_gb: float) -> float:
        """模拟All-Gather通信"""
        # All-Gather: 每个设备收集其他设备的数据
        comm_volume = (self.sp_degree - 1) * data_size_gb

        avg_bandwidth = sum(
            self.topology.bandwidth_matrix[i][j]
            for i in range(self.sp_degree)
            for j in range(self.sp_degree)
            if i != j
        ) / (self.sp_degree * (self.sp_degree - 1))

        comm_time = comm_volume / avg_bandwidth
        time.sleep(comm_time * 0.001)

        return comm_time

    def _simulate_reduce_scatter(self, data_size_gb: float) -> float:
        """模拟Reduce-Scatter通信"""
        # Reduce-Scatter: 归约并分散结果
        comm_volume = (self.sp_degree - 1) / self.sp_degree * data_size_gb

        avg_bandwidth = sum(
            self.topology.bandwidth_matrix[i][j]
            for i in range(self.sp_degree)
            for j in range(self.sp_degree)
            if i != j
        ) / (self.sp_degree * (self.sp_degree - 1))

        comm_time = comm_volume / avg_bandwidth
        time.sleep(comm_time * 0.001)

        return comm_time

    def _simulate_gradient_allreduce(self) -> float:
        """模拟梯度All-Reduce"""
        gradient_size_gb = self.workload.model_config.model_memory_gb

        # Ring All-Reduce
        comm_volume = 2 * (self.sp_degree - 1) / self.sp_degree * gradient_size_gb

        min_bandwidth = min(
            self.topology.bandwidth_matrix[i][j]
            for i in range(self.sp_degree)
            for j in range(self.sp_degree)
            if i != j
        )

        comm_time = comm_volume / min_bandwidth
        time.sleep(comm_time * 0.001)

        return comm_time

    def get_memory_usage(self) -> Dict[int, float]:
        """获取内存使用"""
        return {i: self.memory_manager.memory_usage[i] for i in range(self.num_devices)}

# ======================
# 专家并行策略 (MoE)
# ======================

class ExpertParallelStrategy(ParallelStrategy):
    """专家并行策略：MoE (Mixture of Experts)"""

    def __init__(self, workload: LLMWorkload, topology: CommunicationTopology,
                 memory_manager: MemoryManager, num_experts: int = 8,
                 experts_per_device: int = 2, top_k: int = 2):
        super().__init__(workload, topology, memory_manager)
        self.num_experts = num_experts
        self.experts_per_device = experts_per_device
        self.top_k = top_k
        self.num_devices = len(topology.devices)

        if num_experts % experts_per_device != 0:
            raise ValueError(f"专家数量 {num_experts} 必须能被每设备专家数 {experts_per_device} 整除")

        self.expert_devices = num_experts // experts_per_device

        # MoE层配置（假设每4层有一个MoE层）
        self.moe_layers = workload.model_config.num_layers // 4
        self.regular_layers = workload.model_config.num_layers - self.moe_layers

    def setup(self) -> bool:
        """设置专家并行"""
        # 基础模型参数（非专家部分）
        base_model_memory = self.workload.model_config.model_memory_gb * 0.7  # 假设70%是非专家参数

        # 专家参数：每个设备存储部分专家
        expert_memory_per_device = (
            self.workload.model_config.model_memory_gb * 0.3 / self.expert_devices
        )

        # 激活内存：需要存储路由和专家计算的中间结果
        activation_memory_per_device = self.workload.activation_memory_gb * 1.2  # MoE需要更多激活内存

        total_memory_per_device = base_model_memory + expert_memory_per_device + activation_memory_per_device

        # 检查每个设备内存
        for device_id in range(self.num_devices):
            if not self.memory_manager.allocate(device_id, total_memory_per_device):
                print(f"❌ 设备 {device_id} 内存不足: 需要 {total_memory_per_device:.1f}GB")
                return False

        print(f"📊 专家并行配置: {self.num_experts}专家, "
              f"{self.experts_per_device}专家/设备, Top-{self.top_k}")
        return True

    def forward_pass(self) -> Tuple[float, float]:
        """前向传播"""
        total_compute_time = 0.0
        total_comm_time = 0.0

        # 常规Transformer层
        for layer_idx in range(self.regular_layers):
            layer_compute = self._simulate_regular_layer()
            total_compute_time += layer_compute

        # MoE层
        for moe_layer_idx in range(self.moe_layers):
            moe_compute, moe_comm = self._simulate_moe_layer()
            total_compute_time += moe_compute
            total_comm_time += moe_comm

        return total_compute_time, total_comm_time

    def backward_pass(self) -> Tuple[float, float]:
        """反向传播"""
        # 反向传播的计算量约为前向的2倍
        forward_compute, forward_comm = self.forward_pass()

        # 专家梯度All-Reduce（只在专家设备间进行）
        expert_gradient_comm = self._simulate_expert_gradient_allreduce()

        return forward_compute * 2, forward_comm + expert_gradient_comm

    def _simulate_regular_layer(self) -> float:
        """模拟常规Transformer层"""
        # 常规层在所有设备上复制执行
        hidden_size = self.workload.model_config.hidden_size
        seq_len = self.workload.sequence_length
        batch_size = self.workload.batch_size

        # Attention + FFN计算
        layer_flops = (
            # Attention
            seq_len * seq_len * hidden_size * batch_size +
            seq_len * hidden_size * hidden_size * batch_size * 4 +
            # 常规FFN
            seq_len * hidden_size * self.workload.model_config.intermediate_size * batch_size * 2
        )

        return self.simulate_compute(layer_flops, 0)

    def _simulate_moe_layer(self) -> Tuple[float, float]:
        """模拟MoE层"""
        hidden_size = self.workload.model_config.hidden_size
        seq_len = self.workload.sequence_length
        batch_size = self.workload.batch_size

        # 1. 门控网络计算（本地）
        gate_flops = seq_len * hidden_size * self.num_experts * batch_size
        gate_compute_time = self.simulate_compute(gate_flops, 0)

        # 2. Token路由和All-to-All通信
        # 假设每个token平均路由到top_k个专家
        tokens_per_expert = seq_len * batch_size * self.top_k / self.num_experts
        token_size_gb = (
            tokens_per_expert * hidden_size * self.workload.model_config.dtype_bytes / (1024**3)
        )

        # All-to-All通信：将token发送到对应的专家设备
        alltoall_time = self._simulate_alltoall_communication(token_size_gb)

        # 3. 专家计算（并行）
        expert_flops = (
            tokens_per_expert * hidden_size *
            self.workload.model_config.intermediate_size * 2 * self.experts_per_device
        )
        expert_compute_time = self.simulate_compute(expert_flops, 0)

        # 4. All-to-All通信返回结果
        alltoall_return_time = self._simulate_alltoall_communication(token_size_gb)

        total_compute = gate_compute_time + expert_compute_time
        total_comm = alltoall_time + alltoall_return_time

        return total_compute, total_comm

    def _simulate_alltoall_communication(self, data_size_gb: float) -> float:
        """模拟All-to-All通信"""
        # All-to-All: 每个设备向每个其他设备发送数据
        total_comm_volume = data_size_gb * (self.expert_devices - 1)

        # 使用平均带宽
        avg_bandwidth = sum(
            self.topology.bandwidth_matrix[i][j]
            for i in range(self.expert_devices)
            for j in range(self.expert_devices)
            if i != j
        ) / (self.expert_devices * (self.expert_devices - 1))

        comm_time = total_comm_volume / avg_bandwidth
        time.sleep(comm_time * 0.001)

        return comm_time

    def _simulate_expert_gradient_allreduce(self) -> float:
        """模拟专家梯度All-Reduce"""
        # 只有专家参数需要在专家设备间同步
        expert_gradient_size_gb = self.workload.model_config.model_memory_gb * 0.3 / self.expert_devices

        # Ring All-Reduce在专家设备间
        comm_volume = 2 * (self.expert_devices - 1) / self.expert_devices * expert_gradient_size_gb

        min_bandwidth = min(
            self.topology.bandwidth_matrix[i][j]
            for i in range(self.expert_devices)
            for j in range(self.expert_devices)
            if i != j
        )

        comm_time = comm_volume / min_bandwidth
        time.sleep(comm_time * 0.001)

        return comm_time

    def get_memory_usage(self) -> Dict[int, float]:
        """获取内存使用"""
        return {i: self.memory_manager.memory_usage[i] for i in range(self.num_devices)}

# ======================
# 3D并行策略 (混合并行)
# ======================

class HybridParallelStrategy(ParallelStrategy):
    """3D并行策略：数据并行 + 张量并行 + 流水线并行"""

    def __init__(self, workload: LLMWorkload, topology: CommunicationTopology,
                 memory_manager: MemoryManager, dp_degree: int = 2,
                 tp_degree: int = 2, pp_degree: int = 2):
        super().__init__(workload, topology, memory_manager)
        self.dp_degree = dp_degree
        self.tp_degree = tp_degree
        self.pp_degree = pp_degree
        self.num_devices = len(topology.devices)

        # 验证并行度配置
        total_parallelism = dp_degree * tp_degree * pp_degree
        if total_parallelism != self.num_devices:
            raise ValueError(f"并行度乘积 {total_parallelism} 必须等于设备数 {self.num_devices}")

        # 计算各种分组
        self.tp_group_size = tp_degree
        self.pp_group_size = pp_degree
        self.dp_group_size = dp_degree

        # 微批次配置
        self.num_microbatches = max(4, pp_degree * 2)
        self.microbatch_size = workload.batch_size // (dp_degree * self.num_microbatches)

        # 每个流水线阶段的层数
        self.layers_per_stage = workload.model_config.num_layers // pp_degree

    def setup(self) -> bool:
        """设置3D并行"""
        # 模型内存：按张量并行和流水线并行分片
        model_memory_per_device = (
            self.workload.model_config.model_memory_gb / (self.tp_degree * self.pp_degree)
        )

        # 激活内存：按所有维度分片
        activation_memory_per_device = (
            self.workload.activation_memory_gb / (self.dp_degree * self.tp_degree * self.pp_degree) *
            self.num_microbatches
        )

        total_memory_per_device = model_memory_per_device + activation_memory_per_device

        # 检查每个设备内存
        for device_id in range(self.num_devices):
            if not self.memory_manager.allocate(device_id, total_memory_per_device):
                print(f"❌ 设备 {device_id} 内存不足: 需要 {total_memory_per_device:.1f}GB")
                return False

        print(f"📊 3D并行配置: DP={self.dp_degree}, TP={self.tp_degree}, PP={self.pp_degree}")
        print(f"   微批次: {self.num_microbatches}, 微批次大小: {self.microbatch_size}")
        return True

    def forward_pass(self) -> Tuple[float, float]:
        """前向传播"""
        # 3D并行的前向传播结合了所有并行策略的特点
        total_compute_time = 0.0
        total_comm_time = 0.0

        # 流水线并行：处理多个微批次
        stage_compute_time = self._get_stage_compute_time()

        # 流水线填充时间
        pipeline_fill_time = self.pp_degree * stage_compute_time

        # 稳态执行时间
        steady_state_time = (self.num_microbatches - self.pp_degree) * stage_compute_time

        total_compute_time = pipeline_fill_time + steady_state_time

        # 通信时间包括：
        # 1. 张量并行通信（每层内）
        tp_comm_time = self._simulate_tensor_parallel_communication()

        # 2. 流水线并行通信（阶段间）
        pp_comm_time = self._simulate_pipeline_communication()

        total_comm_time = tp_comm_time + pp_comm_time

        return total_compute_time, total_comm_time

    def backward_pass(self) -> Tuple[float, float]:
        """反向传播"""
        # 反向传播计算时间
        forward_compute, forward_comm = self.forward_pass()
        backward_compute = forward_compute * 2  # 反向约为前向的2倍

        # 反向传播通信包括：
        # 1. 张量并行梯度通信
        tp_grad_comm = self._simulate_tensor_parallel_communication()

        # 2. 流水线并行梯度通信
        pp_grad_comm = self._simulate_pipeline_communication()

        # 3. 数据并行梯度All-Reduce
        dp_grad_comm = self._simulate_data_parallel_allreduce()

        total_comm = forward_comm + tp_grad_comm + pp_grad_comm + dp_grad_comm

        return backward_compute, total_comm

    def _get_stage_compute_time(self) -> float:
        """计算每个流水线阶段的计算时间"""
        # 每个阶段处理部分层，每个张量并行组处理部分计算
        stage_flops = (
            self.workload.compute_flops / self.workload.batch_size * self.microbatch_size /
            (self.pp_degree * self.tp_degree)
        )

        return self.simulate_compute(stage_flops, 0)

    def _simulate_tensor_parallel_communication(self) -> float:
        """模拟张量并行通信"""
        # 每层的All-Reduce通信
        activation_size_per_layer = (
            self.workload.sequence_length * self.workload.model_config.hidden_size *
            self.microbatch_size * self.workload.model_config.dtype_bytes / (1024**3)
        )

        # 每层有2次All-Reduce（Attention输出 + FFN输出）
        comm_per_layer = activation_size_per_layer * 2
        total_layers = self.layers_per_stage

        total_comm_volume = comm_per_layer * total_layers

        # 张量并行组内通信
        tp_bandwidth = self.topology.bandwidth_matrix[0][1]  # 假设组内带宽相同
        comm_time = total_comm_volume / tp_bandwidth

        time.sleep(comm_time * 0.001)
        return comm_time

    def _simulate_pipeline_communication(self) -> float:
        """模拟流水线并行通信"""
        # 阶段间激活传递
        activation_size_per_microbatch = (
            self.workload.sequence_length * self.workload.model_config.hidden_size *
            self.microbatch_size * self.workload.model_config.dtype_bytes / (1024**3)
        )

        # 每个微批次在阶段间传递
        total_comm_volume = activation_size_per_microbatch * self.num_microbatches

        # 流水线阶段间通信
        pp_bandwidth = self.topology.bandwidth_matrix[0][self.tp_degree]  # 跨阶段带宽
        comm_time = total_comm_volume / pp_bandwidth

        time.sleep(comm_time * 0.001)
        return comm_time

    def _simulate_data_parallel_allreduce(self) -> float:
        """模拟数据并行All-Reduce"""
        # 梯度大小等于每个数据并行副本的模型参数
        gradient_size_gb = self.workload.model_config.model_memory_gb / (self.tp_degree * self.pp_degree)

        # Ring All-Reduce在数据并行组内
        comm_volume = 2 * (self.dp_degree - 1) / self.dp_degree * gradient_size_gb

        # 数据并行组间通信带宽
        dp_bandwidth = self.topology.bandwidth_matrix[0][self.tp_degree * self.pp_degree]
        comm_time = comm_volume / dp_bandwidth

        time.sleep(comm_time * 0.001)
        return comm_time

    def get_memory_usage(self) -> Dict[int, float]:
        """获取内存使用"""
        return {i: self.memory_manager.memory_usage[i] for i in range(self.num_devices)}

# ======================
# LLM专用可视化器
# ======================

class LLMVisualizer:
    """LLM专用可视化器"""

    def __init__(self):
        self.has_matplotlib = HAS_MATPLOTLIB

    def print_progress(self, current: int, total: int, prefix: str = "进度"):
        """打印进度条"""
        if HAS_TQDM:
            return

        percent = (current / total) * 100
        bar_length = 30
        filled_length = int(bar_length * current // total)
        bar = '█' * filled_length + '-' * (bar_length - filled_length)
        print(f'\r{prefix}: |{bar}| {percent:.1f}% ({current}/{total})', end='', flush=True)
        if current == total:
            print()

    def show_model_architecture(self, config: TransformerConfig):
        """显示模型架构信息"""
        print("\n🏗️  LLM模型架构")
        print("="*60)
        print(f"模型参数量: {config.total_params/1e9:.1f}B")
        print(f"隐藏层大小: {config.hidden_size}")
        print(f"层数: {config.num_layers}")
        print(f"注意力头数: {config.num_attention_heads}")
        print(f"FFN中间层大小: {config.intermediate_size}")
        print(f"最大序列长度: {config.max_sequence_length}")
        print(f"模型内存: {config.model_memory_gb:.1f}GB")
        print("="*60)

    def show_workload_info(self, workload: LLMWorkload):
        """显示工作负载信息"""
        print("\n📊 训练工作负载")
        print("="*60)
        print(f"批次大小: {workload.batch_size}")
        print(f"序列长度: {workload.sequence_length}")
        print(f"每批次Token数: {workload.tokens_per_batch:,}")
        print(f"激活内存: {workload.activation_memory_gb:.1f}GB")
        print(f"计算FLOPs: {workload.compute_flops/1e12:.1f}T")
        print("="*60)

    def show_performance_comparison(self, metrics_list: List[PerformanceMetrics]):
        """显示性能对比"""
        if not metrics_list:
            return

        if self.has_matplotlib:
            self._plot_performance_charts(metrics_list)
        else:
            self._print_performance_table(metrics_list)

    def _print_performance_table(self, metrics_list: List[PerformanceMetrics]):
        """打印性能对比表格"""
        print("\n📈 LLM并行策略性能对比")
        print("="*120)
        print(f"{'策略':<20} {'总时间(s)':<12} {'计算时间(s)':<12} {'通信时间(s)':<12} "
              f"{'吞吐量(tok/s)':<15} {'内存(GB)':<12} {'通信效率':<12} {'加速比':<10}")
        print("-"*120)

        for metric in metrics_list:
            print(f"{metric.strategy_name:<20} "
                  f"{metric.total_time:<12.3f} "
                  f"{metric.compute_time:<12.3f} "
                  f"{metric.communication_time:<12.3f} "
                  f"{metric.throughput_tokens_per_sec:<15.1f} "
                  f"{metric.memory_usage_per_device:<12.1f} "
                  f"{metric.communication_efficiency:<12.2%} "
                  f"{metric.speedup:<10.2f}")

        print("="*120)

        # 显示最佳策略
        best_throughput = max(metrics_list, key=lambda m: m.throughput_tokens_per_sec)
        best_efficiency = max(metrics_list, key=lambda m: m.communication_efficiency)

        print(f"\n🏆 最高吞吐量: {best_throughput.strategy_name} "
              f"({best_throughput.throughput_tokens_per_sec:.1f} tokens/s)")
        print(f"🏆 最高通信效率: {best_efficiency.strategy_name} "
              f"({best_efficiency.communication_efficiency:.2%})")

    def _plot_performance_charts(self, metrics_list: List[PerformanceMetrics]):
        """绘制性能图表"""
        if not self.has_matplotlib:
            return

        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(16, 12))
        fig.suptitle('LLM并行策略性能分析', fontsize=16, fontweight='bold')

        strategies = [m.strategy_name for m in metrics_list]

        # 1. 时间分解图
        compute_times = [m.compute_time for m in metrics_list]
        comm_times = [m.communication_time for m in metrics_list]

        x = np.arange(len(strategies))
        width = 0.35

        ax1.bar(x, compute_times, width, label='计算时间', color='skyblue')
        ax1.bar(x, comm_times, width, bottom=compute_times, label='通信时间', color='lightcoral')
        ax1.set_title('执行时间分解')
        ax1.set_ylabel('时间 (秒)')
        ax1.set_xticks(x)
        ax1.set_xticklabels(strategies, rotation=45, ha='right')
        ax1.legend()

        # 2. 吞吐量对比
        throughputs = [m.throughput_tokens_per_sec for m in metrics_list]
        bars2 = ax2.bar(strategies, throughputs, color='lightgreen')
        ax2.set_title('吞吐量对比')
        ax2.set_ylabel('Tokens/秒')
        ax2.tick_params(axis='x', rotation=45)

        # 添加数值标签
        for bar, throughput in zip(bars2, throughputs):
            height = bar.get_height()
            ax2.text(bar.get_x() + bar.get_width()/2., height + height*0.01,
                    f'{throughput:.0f}', ha='center', va='bottom')

        # 3. 内存使用对比
        memory_usage = [m.memory_usage_per_device for m in metrics_list]
        bars3 = ax3.bar(strategies, memory_usage, color='gold')
        ax3.set_title('内存使用对比')
        ax3.set_ylabel('内存 (GB)')
        ax3.tick_params(axis='x', rotation=45)

        # 4. 效率指标
        comm_efficiency = [m.communication_efficiency * 100 for m in metrics_list]
        scalability_efficiency = [m.scalability_efficiency * 100 for m in metrics_list]

        x = np.arange(len(strategies))
        width = 0.35

        ax4.bar(x - width/2, comm_efficiency, width, label='通信效率', color='mediumpurple')
        ax4.bar(x + width/2, scalability_efficiency, width, label='可扩展性效率', color='orange')
        ax4.set_title('效率指标对比')
        ax4.set_ylabel('效率 (%)')
        ax4.set_xticks(x)
        ax4.set_xticklabels(strategies, rotation=45, ha='right')
        ax4.legend()

        plt.tight_layout()
        plt.show()

    def show_memory_breakdown(self, metrics_list: List[PerformanceMetrics]):
        """显示内存使用分解"""
        print("\n💾 内存使用分析")
        print("="*80)

        for metric in metrics_list:
            print(f"\n{metric.strategy_name}:")
            print(f"  平均内存使用: {metric.memory_usage_per_device:.1f}GB/设备")
            print(f"  峰值内存使用: {metric.peak_memory_usage:.1f}GB")

            if 'memory_per_device' in metric.details:
                device_memory = metric.details['memory_per_device']
                print(f"  设备内存分布: {[f'{mem:.1f}GB' for mem in device_memory.values()]}")

        print("="*80)

# ======================
# 主要的LLM模拟器类
# ======================

class LLMParallelSimulator:
    """LLM并行训练模拟器"""

    def __init__(self, model_config: TransformerConfig, num_devices: int = 8,
                 topology_type: str = "dgx"):
        self.model_config = model_config
        self.num_devices = num_devices
        self.topology = CommunicationTopology(num_devices, topology_type)
        self.memory_manager = MemoryManager(self.topology.devices)
        self.visualizer = LLMVisualizer()

        # 支持的并行策略
        self.available_strategies = {
            'baseline': BaselineStrategy,
            'data_parallel': DataParallelStrategy,
            'tensor_parallel': TensorParallelStrategy,
            'pipeline_parallel': PipelineParallelStrategy,
            'sequence_parallel': SequenceParallelStrategy,
            'expert_parallel': ExpertParallelStrategy,
            'hybrid_3d': HybridParallelStrategy
        }

    def create_workload(self, batch_size: int, sequence_length: int) -> LLMWorkload:
        """创建工作负载"""
        return LLMWorkload(self.model_config, sequence_length, batch_size)

    def run_strategy(self, strategy_name: str, workload: LLMWorkload,
                    **strategy_kwargs) -> PerformanceMetrics:
        """运行特定并行策略"""
        if strategy_name not in self.available_strategies:
            raise ValueError(f"不支持的策略: {strategy_name}")

        # 重置内存管理器
        self.memory_manager.reset()

        # 创建策略实例
        strategy_class = self.available_strategies[strategy_name]

        try:
            if strategy_name == 'baseline':
                strategy = strategy_class(workload, self.topology, self.memory_manager)
            else:
                strategy = strategy_class(workload, self.topology, self.memory_manager, **strategy_kwargs)

            # 运行模拟
            metrics = strategy.run_simulation()
            return metrics

        except Exception as e:
            print(f"❌ 策略 {strategy_name} 运行失败: {e}")
            raise

    def run_comprehensive_comparison(self, workload: LLMWorkload) -> List[PerformanceMetrics]:
        """运行全面的并行策略对比"""
        print("🚀 开始LLM并行策略全面对比...")
        print("="*80)

        # 显示模型和工作负载信息
        self.visualizer.show_model_architecture(self.model_config)
        self.visualizer.show_workload_info(workload)

        all_metrics = []

        # 定义要测试的策略配置
        strategy_configs = [
            ('baseline', 'Baseline (单设备)', {}),
            ('data_parallel', 'Data Parallel (DDP)', {}),
            ('data_parallel', 'Data Parallel (FSDP)', {'dp_type': 'fsdp'}),
            ('tensor_parallel', 'Tensor Parallel', {'tp_degree': min(4, self.num_devices)}),
            ('pipeline_parallel', 'Pipeline Parallel (GPipe)', {'pp_degree': min(4, self.num_devices)}),
            ('pipeline_parallel', 'Pipeline Parallel (PipeDream)', {
                'pp_degree': min(4, self.num_devices), 'schedule': 'pipedream'
            }),
        ]

        # 如果设备足够多，添加更复杂的策略
        if self.num_devices >= 8:
            strategy_configs.extend([
                ('sequence_parallel', 'Sequence Parallel', {'sp_degree': 4}),
                ('expert_parallel', 'Expert Parallel (MoE)', {'num_experts': 8, 'experts_per_device': 2}),
                ('hybrid_3d', '3D Parallel', {'dp_degree': 2, 'tp_degree': 2, 'pp_degree': 2}),
            ])

        # 运行每个策略
        for strategy_name, display_name, kwargs in strategy_configs:
            try:
                print(f"\n🔄 运行 {display_name}...")
                metrics = self.run_strategy(strategy_name, workload, **kwargs)
                metrics.strategy_name = display_name  # 使用更友好的显示名称
                all_metrics.append(metrics)
                time.sleep(0.5)  # 短暂暂停以便观察
            except Exception as e:
                print(f"❌ {display_name} 跳过: {e}")
                continue

        print("\n" + "="*80)
        print("🎯 所有策略测试完成!")

        # 显示综合对比
        if all_metrics:
            # 设置基准时间用于计算加速比
            baseline_time = all_metrics[0].total_time if all_metrics else 1.0
            for metric in all_metrics:
                metric.details['baseline_time'] = baseline_time

            self.visualizer.show_performance_comparison(all_metrics)
            self.visualizer.show_memory_breakdown(all_metrics)

        return all_metrics

    def run_scalability_analysis(self, workload: LLMWorkload, strategy_name: str = 'data_parallel'):
        """运行可扩展性分析"""
        print(f"\n📊 {strategy_name} 可扩展性分析")
        print("="*60)

        device_counts = [1, 2, 4, 8] if self.num_devices >= 8 else [1, 2, 4]
        scalability_metrics = []

        for num_dev in device_counts:
            if num_dev > self.num_devices:
                continue

            print(f"\n测试 {num_dev} 设备配置...")

            # 创建临时拓扑
            temp_topology = CommunicationTopology(num_dev, self.topology.topology_type)
            temp_memory_manager = MemoryManager(temp_topology.devices)

            try:
                strategy_class = self.available_strategies[strategy_name]
                strategy = strategy_class(workload, temp_topology, temp_memory_manager)
                metrics = strategy.run_simulation()
                metrics.details['num_devices'] = num_dev
                scalability_metrics.append(metrics)
            except Exception as e:
                print(f"❌ {num_dev}设备配置失败: {e}")

        # 分析可扩展性
        if len(scalability_metrics) > 1:
            print(f"\n📈 可扩展性分析结果:")
            print("-"*60)
            baseline = scalability_metrics[0]

            for metric in scalability_metrics:
                num_dev = metric.details['num_devices']
                ideal_speedup = num_dev
                actual_speedup = baseline.total_time / metric.total_time
                efficiency = actual_speedup / ideal_speedup * 100

                print(f"{num_dev}设备: 实际加速比 {actual_speedup:.2f}x, "
                      f"理想加速比 {ideal_speedup:.2f}x, 效率 {efficiency:.1f}%")

        return scalability_metrics

# ======================
# 交互式演示界面
# ======================

def create_preset_models() -> Dict[str, TransformerConfig]:
    """创建预设模型配置"""
    return {
        'GPT-Small': TransformerConfig(
            vocab_size=50000, hidden_size=768, num_layers=12, num_attention_heads=12,
            intermediate_size=3072, max_sequence_length=1024
        ),
        'GPT-Medium': TransformerConfig(
            vocab_size=50000, hidden_size=1024, num_layers=24, num_attention_heads=16,
            intermediate_size=4096, max_sequence_length=1024
        ),
        'GPT-Large': TransformerConfig(
            vocab_size=50000, hidden_size=1536, num_layers=48, num_attention_heads=24,
            intermediate_size=6144, max_sequence_length=1024
        ),
        'GPT-XL': TransformerConfig(
            vocab_size=50000, hidden_size=2048, num_layers=48, num_attention_heads=32,
            intermediate_size=8192, max_sequence_length=2048
        ),
        'LLaMA-7B': TransformerConfig(
            vocab_size=32000, hidden_size=4096, num_layers=32, num_attention_heads=32,
            intermediate_size=11008, max_sequence_length=2048
        ),
        'LLaMA-13B': TransformerConfig(
            vocab_size=32000, hidden_size=5120, num_layers=40, num_attention_heads=40,
            intermediate_size=13824, max_sequence_length=2048
        ),
    }

def interactive_llm_demo():
    """交互式LLM并行演示"""
    print("🤖 LLM并行训练模拟器")
    print("="*80)
    print("欢迎使用专业的大语言模型并行训练教学工具！")
    print("本工具模拟现代LLM的各种并行训练策略，帮助理解分布式训练原理。")

    preset_models = create_preset_models()

    while True:
        print("\n📋 主菜单:")
        print("1. 选择预设模型进行对比测试")
        print("2. 自定义模型配置")
        print("3. 可扩展性分析")
        print("4. 单个策略深入分析")
        print("5. 查看并行策略详细说明")
        print("6. 退出")

        try:
            choice = input("\n请输入选择 (1-6): ").strip()

            if choice == '1':
                # 预设模型对比测试
                print("\n🏗️  选择预设模型:")
                model_names = list(preset_models.keys())
                for i, name in enumerate(model_names, 1):
                    config = preset_models[name]
                    print(f"{i}. {name} ({config.total_params/1e9:.1f}B参数, {config.model_memory_gb:.1f}GB)")

                try:
                    model_idx = int(input("请选择模型 (1-{}): ".format(len(model_names)))) - 1
                    if 0 <= model_idx < len(model_names):
                        model_name = model_names[model_idx]
                        config = preset_models[model_name]

                        # 设备数量选择
                        num_devices = int(input("设备数量 (默认8): ") or "8")

                        # 训练配置
                        batch_size = int(input("批次大小 (默认32): ") or "32")
                        sequence_length = int(input(f"序列长度 (默认{config.max_sequence_length}): ") or str(config.max_sequence_length))

                        # 创建模拟器和工作负载
                        simulator = LLMParallelSimulator(config, num_devices)
                        workload = simulator.create_workload(batch_size, sequence_length)

                        # 运行对比测试
                        metrics = simulator.run_comprehensive_comparison(workload)

                    else:
                        print("❌ 无效选择")
                except ValueError:
                    print("❌ 输入无效，请输入数字")

            elif choice == '2':
                # 自定义模型配置
                print("\n🔧 自定义模型配置:")
                try:
                    hidden_size = int(input("隐藏层大小 (默认4096): ") or "4096")
                    num_layers = int(input("层数 (默认32): ") or "32")
                    num_heads = int(input("注意力头数 (默认32): ") or "32")
                    vocab_size = int(input("词汇表大小 (默认50000): ") or "50000")

                    config = TransformerConfig(
                        vocab_size=vocab_size,
                        hidden_size=hidden_size,
                        num_layers=num_layers,
                        num_attention_heads=num_heads,
                        intermediate_size=hidden_size * 4,
                        max_sequence_length=2048
                    )

                    num_devices = int(input("设备数量 (默认8): ") or "8")
                    batch_size = int(input("批次大小 (默认16): ") or "16")
                    sequence_length = int(input("序列长度 (默认2048): ") or "2048")

                    simulator = LLMParallelSimulator(config, num_devices)
                    workload = simulator.create_workload(batch_size, sequence_length)

                    metrics = simulator.run_comprehensive_comparison(workload)

                except ValueError:
                    print("❌ 输入无效，请输入数字")

            elif choice == '3':
                # 可扩展性分析
                print("\n📊 可扩展性分析")
                print("选择模型:")
                model_names = list(preset_models.keys())
                for i, name in enumerate(model_names, 1):
                    print(f"{i}. {name}")

                try:
                    model_idx = int(input("请选择模型: ")) - 1
                    if 0 <= model_idx < len(model_names):
                        config = preset_models[model_names[model_idx]]

                        print("选择并行策略:")
                        strategies = ['data_parallel', 'tensor_parallel', 'pipeline_parallel']
                        for i, strategy in enumerate(strategies, 1):
                            print(f"{i}. {strategy}")

                        strategy_idx = int(input("请选择策略: ")) - 1
                        if 0 <= strategy_idx < len(strategies):
                            strategy_name = strategies[strategy_idx]

                            simulator = LLMParallelSimulator(config, 8)
                            workload = simulator.create_workload(16, 1024)

                            simulator.run_scalability_analysis(workload, strategy_name)
                        else:
                            print("❌ 无效选择")
                    else:
                        print("❌ 无效选择")
                except ValueError:
                    print("❌ 输入无效，请输入数字")

            elif choice == '4':
                # 单个策略深入分析
                print("\n🔍 单个策略深入分析")
                strategies = {
                    '1': ('data_parallel', 'Data Parallel'),
                    '2': ('tensor_parallel', 'Tensor Parallel'),
                    '3': ('pipeline_parallel', 'Pipeline Parallel'),
                    '4': ('sequence_parallel', 'Sequence Parallel'),
                    '5': ('expert_parallel', 'Expert Parallel (MoE)'),
                    '6': ('hybrid_3d', '3D Parallel')
                }

                print("选择并行策略:")
                for key, (_, name) in strategies.items():
                    print(f"{key}. {name}")

                strategy_choice = input("请选择策略: ").strip()
                if strategy_choice in strategies:
                    strategy_name, display_name = strategies[strategy_choice]

                    # 使用默认配置进行演示
                    config = preset_models['GPT-Large']
                    simulator = LLMParallelSimulator(config, 8)
                    workload = simulator.create_workload(16, 1024)

                    print(f"\n🔄 运行 {display_name} 详细分析...")

                    try:
                        if strategy_name == 'expert_parallel':
                            metrics = simulator.run_strategy(strategy_name, workload,
                                                           num_experts=8, experts_per_device=2)
                        elif strategy_name == 'hybrid_3d':
                            metrics = simulator.run_strategy(strategy_name, workload,
                                                           dp_degree=2, tp_degree=2, pp_degree=2)
                        else:
                            metrics = simulator.run_strategy(strategy_name, workload)

                        # 显示详细信息
                        print(f"\n📊 {display_name} 详细分析结果:")
                        print("-"*60)
                        print(f"总执行时间: {metrics.total_time:.3f}秒")
                        print(f"计算时间: {metrics.compute_time:.3f}秒")
                        print(f"通信时间: {metrics.communication_time:.3f}秒")
                        print(f"吞吐量: {metrics.throughput_tokens_per_sec:.1f} tokens/s")
                        print(f"内存使用: {metrics.memory_usage_per_device:.1f}GB/设备")
                        print(f"通信效率: {metrics.communication_efficiency:.2%}")

                    except Exception as e:
                        print(f"❌ 策略运行失败: {e}")
                else:
                    print("❌ 无效选择")

            elif choice == '5':
                # 查看并行策略说明
                show_strategy_explanations()

            elif choice == '6':
                print("👋 感谢使用LLM并行训练模拟器！")
                break

            else:
                print("❌ 无效选择，请重新输入")

        except KeyboardInterrupt:
            print("\n\n👋 程序已退出")
            break
        except Exception as e:
            print(f"❌ 发生错误: {e}")

def show_strategy_explanations():
    """显示并行策略详细说明"""
    strategies_info = {
        "数据并行 (Data Parallel)": {
            "原理": "将训练数据分割到多个设备，每个设备维护完整模型副本",
            "优点": ["实现简单", "适合大批量训练", "通信开销相对较小"],
            "缺点": ["受限于单设备内存", "模型复制开销", "批次大小受限"],
            "适用场景": "中小型模型、大批量训练",
            "通信模式": "All-Reduce梯度同步",
            "变种": "DDP (DistributedDataParallel), FSDP (Fully Sharded Data Parallel)"
        },
        "张量并行 (Tensor Parallel)": {
            "原理": "将模型的张量运算分割到多个设备上并行计算",
            "优点": ["突破单设备内存限制", "细粒度并行", "支持超大模型"],
            "缺点": ["通信开销大", "实现复杂", "需要高带宽互连"],
            "适用场景": "超大模型、Transformer架构",
            "通信模式": "All-Reduce张量同步",
            "变种": "Megatron-LM, DeepSpeed-Inference"
        },
        "流水线并行 (Pipeline Parallel)": {
            "原理": "将模型按层分割成多个阶段，不同阶段并行处理不同微批次",
            "优点": ["内存效率高", "支持超深模型", "通信量相对较小"],
            "缺点": ["流水线气泡", "需要微批次调优", "负载均衡困难"],
            "适用场景": "深度模型、内存受限场景",
            "通信模式": "点对点激活传递",
            "变种": "GPipe, PipeDream, 1F1B调度"
        },
        "序列并行 (Sequence Parallel)": {
            "原理": "将长序列按序列维度分割到多个设备上处理",
            "优点": ["支持超长序列", "内存效率高", "适合注意力机制"],
            "缺点": ["通信开销大", "实现复杂", "序列长度限制"],
            "适用场景": "长序列任务、文档级处理",
            "通信模式": "All-Gather/Reduce-Scatter",
            "变种": "Sequence Parallel, Ring Attention"
        },
        "专家并行 (Expert Parallel)": {
            "原理": "MoE模型中将不同专家分布到不同设备上",
            "优点": ["模型容量大", "计算效率高", "专家特化"],
            "缺点": ["负载不均衡", "路由复杂", "通信不规则"],
            "适用场景": "MoE模型、多任务学习",
            "通信模式": "All-to-All专家路由",
            "变种": "Switch Transformer, GLaM, PaLM"
        },
        "3D并行 (3D Parallel)": {
            "原理": "组合数据并行、张量并行、流水线并行",
            "优点": ["最大化并行度", "灵活配置", "适应各种模型"],
            "缺点": ["配置复杂", "通信复杂", "调优困难"],
            "适用场景": "超大规模训练、千亿参数模型",
            "通信模式": "多层次通信拓扑",
            "变种": "Megatron-DeepSpeed, FairScale"
        }
    }

    print("\n📚 LLM并行策略详细说明")
    print("="*80)

    for strategy_name, info in strategies_info.items():
        print(f"\n🔹 {strategy_name}")
        print("-" * 60)
        print(f"原理: {info['原理']}")
        print(f"✅ 优点: {', '.join(info['优点'])}")
        print(f"❌ 缺点: {', '.join(info['缺点'])}")
        print(f"🎯 适用场景: {info['适用场景']}")
        print(f"🔄 通信模式: {info['通信模式']}")
        print(f"🔧 主要变种: {info['变种']}")

    print("\n" + "="*80)

# ======================
# 主程序入口
# ======================

if __name__ == "__main__":
    # 检查依赖
    missing_deps = []
    if not HAS_MATPLOTLIB:
        missing_deps.append("matplotlib numpy")
    if not HAS_TQDM:
        missing_deps.append("tqdm")

    if missing_deps:
        print("⚠️  建议安装以下依赖以获得最佳体验:")
        for dep in missing_deps:
            print(f"   pip install {dep}")
        print()

    # 启动交互式演示
    interactive_llm_demo()