# moe_router_simulator.py
# Astra-scale MoE Router Simulation with Load Balancing
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import math
from typing import List, Tuple, Optional
import time

class MoEConfig:
    """MoE Configuration for Astra-scale simulation"""
    def __init__(
        self,
        num_experts: int = 256,
        top_k: int = 12,
        d_model: int = 40960,  # hidden dimension ~40K
        d_ff: int = 81920,     # FFN hidden dimension
        capacity_factor: float = 1.25,
        use_aux_loss: bool = True,
        aux_loss_coef: float = 0.01,
        z_loss_coef: float = 0.001,
    ):
        self.num_experts = num_experts
        self.top_k = top_k
        self.d_model = d_model
        self.d_ff = d_ff
        self.capacity_factor = capacity_factor
        self.use_aux_loss = use_aux_loss
        self.aux_loss_coef = aux_loss_coef
        self.z_loss_coef = z_loss_coef

    @property
    def total_params_per_expert(self) -> int:
        """Total params in one expert FFN (gate + up + down projections)"""
        return 3 * self.d_model * self.d_ff

    @property
    def total_params_gating(self) -> int:
        """Gating network params"""
        return self.d_model * self.num_experts

    @property
    def total_params_single_layer(self) -> int:
        return self.num_experts * self.total_params_per_expert + self.total_params_gating

    def __repr__(self) -> str:
        return (
            f"MoEConfig(num_experts={self.num_experts}, top_k={self.top_k}, "
            f"d_model={self.d_model}, d_ff={self.d_ff}, "
            f"capacity_factor={self.capacity_factor})"
        )


class TopKRouter:
    """Top-K routing with load balancing and auxiliary loss"""

    def __init__(self, config: MoEConfig):
        self.config = config
        # Simulate gating weights
        self.gate_weights = np.random.randn(config.d_model, config.num_experts).astype(np.float32) * 0.02
        self.gate_bias = np.zeros(config.num_experts, dtype=np.float32)
        self.rng = np.random.default_rng(42)

    def forward(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray, dict]:
        """
        Forward pass with routing.
        Args:
            x: (batch_size, seq_len, d_model) or (num_tokens, d_model)
        Returns:
            routing_weights: (num_tokens, top_k)
            expert_indices: (num_tokens, top_k)
            aux_info: dict with auxiliary metrics
        """
        orig_shape = x.shape
        if len(orig_shape) == 3:
            batch, seq, d = orig_shape
            x_flat = x.reshape(-1, d)
        else:
            x_flat = x
            batch, seq = 1, len(x)

        num_tokens = x_flat.shape[0]

        # Compute logits: (num_tokens, num_experts)
        logits = x_flat @ self.gate_weights + self.gate_bias

        # Add noise for training stability (not used in inference)
        if self.rng.random() < 0.3:
            noise = self.rng.normal(0, 0.01, logits.shape).astype(np.float32)
            logits = logits + noise

        # Top-K selection
        top_k = min(self.config.top_k, self.config.num_experts)
        # Use partition-based selection for efficiency
        # Simulate: find top-k values and indices
        indices = np.argpartition(-logits, top_k, axis=1)[:, :top_k]
        values = np.take_along_axis(logits, indices, axis=1)

        # Softmax over selected experts
        values_exp = np.exp(values - np.max(values, axis=1, keepdims=True))
        routing_weights = values_exp / np.sum(values_exp, axis=1, keepdims=True)

        # Load balancing metrics
        expert_counts = np.zeros(self.config.num_experts, dtype=np.float32)
        for i in range(num_tokens):
            for j in range(top_k):
                expert_counts[indices[i, j]] += routing_weights[i, j]

        # Importance (sum of routing weights per expert)
        importance = expert_counts.copy()
        # Load (number of tokens routed to each expert)
        load = np.zeros(self.config.num_experts, dtype=np.float32)
        for i in range(num_tokens):
            for j in range(top_k):
                load[indices[i, j]] += 1.0

        # Auxiliary loss (load balancing loss)
        # CV = std(load) / mean(load)
        cv = float(np.std(load) / (np.mean(load) + 1e-8))

        aux_loss = 0.0
        if self.config.use_aux_loss:
            # z-loss: prevent logits from growing too large
            z_loss = np.mean(np.log(np.sum(np.exp(logits - np.max(logits, axis=1, keepdims=True)), axis=1)) ** 2)
            # Load balancing loss (simplified)
            bal_loss = cv * 0.1
            aux_loss = self.config.aux_loss_coef * bal_loss + self.config.z_loss_coef * float(z_loss)

        aux_info = {
            "expert_importance": importance,
            "expert_load": load,
            "cv": cv,
            "aux_loss": aux_loss,
            "num_tokens": num_tokens,
            "top_k_used": top_k,
            "capacity_utilization": np.mean(load) / (num_tokens * top_k / self.config.num_experts + 1e-8),
        }

        return routing_weights, indices, aux_info


def simulate_astra_moe_routing():
    """Full-scale simulation of Astra MoE routing behavior"""
    print("=" * 70)
    print("Astra (10T params) MoE Router Simulation")
    print("=" * 70)

    # Astra-scale configuration
    config = MoEConfig(
        num_experts=256,
        top_k=12,
        d_model=40960,
        d_ff=81920,
        capacity_factor=1.25,
        use_aux_loss=True,
    )
    print(f"Config: {config}")
    print(f"  Total params per MoE layer: {config.total_params_single_layer / 1e12:.2f}T")
    print(f"  Gating params: {config.total_params_gating / 1e9:.2f}B")

    # Simulate multiple steps with varying token distributions
    router = TopKRouter(config)
    token_counts = [4096, 8192, 16384, 32768, 65536, 131072]
    results = []

    for n_tokens in token_counts:
        # Generate random input
        x = np.random.randn(n_tokens, config.d_model).astype(np.float32) * 0.1
        t0 = time.time()
        weights, indices, info = router.forward(x)
        elapsed = time.time() - t0

        results.append({
            "n_tokens": n_tokens,
            "cv": info["cv"],
            "aux_loss": info["aux_loss"],
            "capacity_util": info["capacity_utilization"],
            "time_ms": elapsed * 1000,
        })

        print(f"\n  Tokens: {n_tokens:>8d} | CV: {info['cv']:.4f} | "
              f"CapUtil: {info['capacity_utilization']:.2%} | Time: {elapsed*1000:.2f}ms")

    # Analyze expert load distribution
    print("\n" + "=" * 70)
    print("Expert Load Distribution Analysis")
    print("=" * 70)
    x_large = np.random.randn(65536, config.d_model).astype(np.float32) * 0.1
    _, _, info = router.forward(x_large)
    load = info["expert_load"]
    importance = info["expert_importance"]

    top_loaded = np.argsort(-load)[:10]
    bottom_loaded = np.argsort(load)[:10]

    print(f"  Top-10 most loaded experts: {top_loaded}")
    print(f"  Top-10 load values: {load[top_loaded]}")
    print(f"  Bottom-10 least loaded experts: {bottom_loaded}")
    print(f"  Bottom-10 load values: {load[bottom_loaded]}")
    print(f"  Load CV (coefficient of variation): {info['cv']:.4f}")
    print(f"  Ideal CV (uniform): {1.0 / math.sqrt(65536 * 12 / 256):.4f}")

    # Summary
    print("\n" + "=" * 70)
    print("Simulation Summary")
    print("=" * 70)
    print(f"  Astra parameter estimate: ~10T total, ~{config.top_k * config.total_params_per_expert / 1e12:.1f}T activated")
    print(f"  Activation ratio: {config.top_k / config.num_experts:.2%}")
    print(f"  Load balancing quality: {'EXCELLENT' if info['cv'] < 0.3 else 'GOOD' if info['cv'] < 0.5 else 'NEEDS IMPROVEMENT'}")

    return results


if __name__ == "__main__":
    simulate_astra_moe_routing()

# scaling_law_analysis.py
# Scaling Law fitting and analysis for Astra-scale models
import numpy as np
from scipy.optimize import curve_fit
from dataclasses import dataclass
from typing import List, Tuple, Optional
import json

@dataclass
class ScalingLawParams:
    """Chinchilla-style scaling law parameters"""
    A: float   # data scaling exponent
    B: float   # parameter scaling exponent
    E: float   # irreducible loss (entropy of data)
    alpha: float  # exponent for compute-optimal allocation

    def loss_from_params(self, N: float, D: float) -> float:
        """Compute loss given params N and data D"""
        return self.E + self.A / (N ** self.alpha) + self.B / (D ** self.alpha)

    def compute_optimal_allocation(self, C: float) -> Tuple[float, float]:
        """Compute optimal N and D for given compute budget C"""
        # C ≈ 6 * N * D (FLOPs)
        # Optimal: N_opt ∝ C^(1/(1+alpha)), D_opt ∝ C^(alpha/(1+alpha))
        N_opt = (C / 6) ** (1 / (1 + self.alpha))
        D_opt = (C / 6) ** (self.alpha / (1 + self.alpha))
        return N_opt, D_opt


def generate_scaling_data(
    param_range: Tuple[float, float] = (1e8, 1e13),
    data_range: Tuple[float, float] = (1e8, 1e13),
    noise_std: float = 0.02,
    n_samples: int = 50,
    seed: int = 42,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, ScalingLawParams]:
    """Generate synthetic scaling law data"""
    rng = np.random.default_rng(seed)

    # True underlying parameters (inspired by Chinchilla)
    true_params = ScalingLawParams(
        A=406.4,
        B=410.7,
        E=1.69,
        alpha=0.34,
    )

    log_N_min, log_N_max = np.log10(param_range[0]), np.log10(param_range[1])
    log_D_min, log_D_max = np.log10(data_range[0]), np.log10(data_range[1])

    N_vals = 10 ** rng.uniform(log_N_min, log_N_max, n_samples)
    D_vals = 10 ** rng.uniform(log_D_min, log_D_max, n_samples)

    losses = np.array([
        true_params.loss_from_params(N, D) + rng.normal(0, noise_std)
        for N, D in zip(N_vals, D_vals)
    ])

    return N_vals, D_vals, losses, true_params


def fit_scaling_law(
    N_vals: np.ndarray,
    D_vals: np.ndarray,
    losses: np.ndarray,
) -> ScalingLawParams:
    """Fit scaling law parameters to observed data"""

    def model_func(params_flat, N, D):
        E, A, B, alpha = params_flat
        return E + A / (N ** alpha) + B / (D ** alpha)

    def residuals(params_flat, N, D, losses):
        return model_func(params_flat, N, D) - losses

    # Initial guess
    initial = [1.5, 400.0, 400.0, 0.35]

    result = curve_fit(
        lambda x, E, A, B, alpha: E + A / (x[0] ** alpha) + B / (x[1] ** alpha),
        [N_vals, D_vals],
        losses,
        p0=initial,
        maxfev=10000,
    )

    E_fit, A_fit, B_fit, alpha_fit = result[0]
    return ScalingLawParams(A=A_fit, B=B_fit, E=E_fit, alpha=alpha_fit)


def analyze_astra_scaling():
    """Analyze scaling law implications for Astra"""
    print("=" * 70)
    print("Scaling Law Analysis for GPT-6 Astra (10T params)")
    print("=" * 70)

    # Generate scaling data
    N_vals, D_vals, losses, true_params = generate_scaling_data(
        param_range=(1e8, 5e12),
        data_range=(1e8, 5e12),
        noise_std=0.015,
        n_samples=80,
    )

    # Fit scaling law
    fitted_params = fit_scaling_law(N_vals, D_vals, losses)
    print(f"\nFitted Scaling Law Parameters:")
    print(f"  E (irreducible loss): {fitted_params.E:.4f}")
    print(f"  A (param scaling):    {fitted_params.A:.2f}")
    print(f"  B (data scaling):     {fitted_params.B:.2f}")
    print(f"  alpha:                {fitted_params.alpha:.4f}")

    # Analyze Astra-scale predictions
    print("\n" + "-" * 50)
    print("Astra-Scale Predictions")
    print("-" * 50)

    astra_params = [2e12, 5e12, 1e13, 2e13]  # 2T, 5T, 10T, 20T
    astra_data = [5e12, 1e13, 2e13, 5e13]  # 5T, 10T, 20T, 50T tokens

    for N, D in zip(astra_params, astra_data):
        loss = fitted_params.loss_from_params(N, D)
        print(f"  N={N:.1e}, D={D:.1e} -> Loss={loss:.4f}")

    # Compute-optimal allocation
    print("\n" + "-" * 50)
    print("Compute-Optimal Allocation Analysis")
    print("-" * 50)

    # Astra estimated compute: 10T params * 10T tokens * 6 FLOPs/token/param
    # But MoE changes the effective compute
    for compute_mult in [1, 2, 5, 10, 20]:
        C_base = 6 * 1e12 * 1e12  # 6 * 1T * 1T
        C = C_base * compute_mult
        N_opt, D_opt = fitted_params.compute_optimal_allocation(C)
        loss_opt = fitted_params.loss_from_params(N_opt, D_opt)
        print(f"  Compute={C:.2e}: N_opt={N_opt:.2e}, D_opt={D_opt:.2e}, Loss={loss_opt:.4f}")

    # Performance degradation bottleneck analysis
    print("\n" + "-" * 50)
    print("Performance Degradation Bottleneck Analysis")
    print("-" * 50)

    # Simulate the effect of increasing N while D is fixed
    D_fixed = 1e13  # 10T tokens
    N_range = np.logspace(10, 14, 100)  # 10B to 100T
    losses_N = [fitted_params.loss_from_params(N, D_fixed) for N in N_range]

    # Find the point of diminishing returns
    marginal_gains = -np.diff(losses_N) / np.diff(np.log10(N_range))
    diminishing_threshold = np.where(marginal_gains < 0.01 * marginal_gains[0])[0]
    if len(diminishing_threshold) > 0:
        N_threshold = N_range[diminishing_threshold[0]]
        print(f"  Diminishing returns threshold (data=10T): N > {N_threshold:.2e}")
        print(f"  At this point, loss reduction per 10x params < 1% of initial")

    # Compare dense vs MoE scaling
    print("\n" + "-" * 50)
    print("Dense vs MoE Scaling Comparison")
    print("-" * 50)

    # Dense model: all params active
    dense_N = 1.8e12  # GPT-4
    moe_total_N = 1e13  # Astra total
    moe_active_N = 5e11  # Astra active (~5%)

    dense_loss = fitted_params.loss_from_params(dense_N, 1e13)
    moe_loss = fitted_params.loss_from_params(moe_total_N, 1e13)
    active_loss = fitted_params.loss_from_params(moe_active_N, 1e13)

    print(f"  Dense (GPT-4, 1.8T):        Loss={dense_loss:.4f}")
    print(f"  MoE Total (Astra, 10T):     Loss={moe_loss:.4f}")
    print(f"  MoE Active (Astra, 0.5T):   Loss={active_loss:.4f}")
    print(f"  MoE advantage (total vs active): {moe_loss - active_loss:.4f}")
    print(f"  MoE vs Dense improvement: {dense_loss - moe_loss:.4f}")

    return fitted_params


if __name__ == "__main__":
    analyze_astra_scaling()

# all_reduce_simulator.py
# Distributed All-Reduce Communication Simulation for 100K GPU Clusters
import numpy as np
from dataclasses import dataclass
from enum import Enum
from typing import List, Dict, Tuple, Optional
import math
import time

class Topology(Enum):
    RING = "ring"
    TREE = "tree"
    TORUS_2D = "torus_2d"
    TORUS_3D = "torus_3d"
    DRAGONFLY = "dragonfly"
    FAT_TREE = "fat_tree"


@dataclass
class NetworkConfig:
    """Network configuration for cluster"""
    topology: Topology
    num_gpus: int
    bw_nvlink: float = 3600.0  # GB/s, NVLink 6
    bw_ib: float = 800.0       # GB/s, InfiniBand NDR 800
    bw_eth: float = 200.0      # GB/s, Ethernet
    latency_nvlink: float = 0.5  # us
    latency_ib: float = 2.0      # us
    latency_eth: float = 10.0    # us
    num_nodes: int = 100000      # number of nodes
    gpus_per_node: int = 8       # GPUs per node


@dataclass
class AllReduceResult:
    """All-Reduce operation result"""
    algorithm: str
    topology: Topology
    message_size: int  # bytes
    theoretical_time: float  # seconds
    bottleneck_bandwidth: float  # GB/s
    total_bytes_transferred: int  # bytes
    bus_bandwidth_utilization: float  # percentage
    steps: int
    description: str


class AllReduceSimulator:
    """Simulate all-reduce operations on various cluster topologies"""

    def __init__(self, config: NetworkConfig):
        self.config = config

    def simulate_ring_allreduce(self, message_size: int) -> AllReduceResult:
        """
        Ring All-Reduce: 2 * (N-1)/N * message_size / bandwidth
        Cost = 2 * (N-1) / N * (message_size / bandwidth)
        For large N, approximates 2 * message_size / bandwidth
        """
        N = self.config.num_gpus
        bw = self.config.bw_nvlink  # Use NVLink for intra-node
        latency = self.config.latency_nvlink

        # Ring all-reduce time: 2 * (N-1) * (message_size / (N * bw)) + latency overhead
        transfer_time = 2 * (N - 1) / N * message_size / (bw * 1e9)
        total_time = transfer_time + 2 * (N - 1) * latency * 1e-6

        total_bytes = 2 * (N - 1) / N * message_size * N  # total bytes transferred in network
        bottleneck_bw = message_size / total_time if total_time > 0 else float('inf')
        bus_util = (2 * message_size / total_time) / (bw * 1e9) if total_time > 0 else 0

        return AllReduceResult(
            algorithm="Ring",
            topology=self.config.topology,
            message_size=message_size,
            theoretical_time=total_time,
            bottleneck_bandwidth=bottleneck_bw / 1e9,
            total_bytes_transferred=int(total_bytes),
            bus_bandwidth_utilization=bus_util * 100,
            steps=2 * (N - 1),
            description=f"Ring all-reduce on {N} GPUs with {bw} GB/s NVLink"
        )

    def simulate_tree_allreduce(self, message_size: int) -> AllReduceResult:
        """
        Tree All-Reduce: log2(N) * 2 * message_size / bandwidth
        """
        N = self.config.num_gpus
        bw = self.config.bw_ib  # Use InfiniBand for tree
        latency = self.config.latency_ib
        logN = math.log2(N)

        transfer_time = 2 * logN * message_size / (bw * 1e9)
        total_time = transfer_time + 2 * logN * latency * 1e-6

        total_bytes = 2 * logN * message_size * N
        bottleneck_bw = message_size / total_time if total_time > 0 else float('inf')

        return AllReduceResult(
            algorithm="Tree",
            topology=self.config.topology,
            message_size=message_size,
            theoretical_time=total_time,
            bottleneck_bandwidth=bottleneck_bw / 1e9,
            total_bytes_transferred=int(total_bytes),
            bus_bandwidth_utilization=(2 * message_size / total_time) / (bw * 1e9) * 100 if total_time > 0 else 0,
            steps=int(2 * logN),
            description=f"Tree all-reduce on {N} GPUs, log2(N)={logN:.1f} steps"
        )

    def simulate_hierarchical_allreduce(self, message_size: int) -> AllReduceResult:
        """
        Hierarchical All-Reduce:
        Intra-node (NVLink Ring) -> Inter-node (IB Tree) -> Intra-node broadcast
        """
        gpu_per_node = self.config.gpus_per_node
        num_nodes = self.config.num_nodes
        bw_local = self.config.bw_nvlink
        bw_global = self.config.bw_ib
        lat_local = self.config.latency_nvlink
        lat_global = self.config.latency_ib

        # Phase 1: Intra-node reduce-scatter
        t1 = 2 * (gpu_per_node - 1) / gpu_per_node * message_size / (bw_local * 1e9)
        t1 += 2 * (gpu_per_node - 1) * lat_local * 1e-6

        # Phase 2: Inter-node all-reduce (on reduced data)
        reduced_size = message_size / gpu_per_node
        logN = math.log2(num_nodes)
        t2 = 2 * logN * reduced_size / (bw_global * 1e9)
        t2 += 2 * logN * lat_global * 1e-6

        # Phase 3: Intra-node broadcast
        t3 = (gpu_per_node - 1) / gpu_per_node * message_size / (bw_local * 1e9)
        t3 += (gpu_per_node - 1) * lat_local * 1e-6

        total_time = t1 + t2 + t3

        total_bytes = int(
            (2 * (gpu_per_node - 1) / gpu_per_node * message_size * gpu_per_node) +  # phase 1
            (2 * logN * reduced_size * num_nodes) +  # phase 2
            ((gpu_per_node - 1) / gpu_per_node * message_size * gpu_per_node)  # phase 3
        )

        return AllReduceResult(
            algorithm="Hierarchical",
            topology=self.config.topology,
            message_size=message_size,
            theoretical_time=total_time,
            bottleneck_bandwidth=message_size / total_time / 1e9 if total_time > 0 else 0,
            total_bytes_transferred=total_bytes,
            bus_bandwidth_utilization=0,
            steps=int(2 * (gpu_per_node - 1) + 2 * logN + (gpu_per_node - 1)),
            description=f"Hierarchical: intra-node ring + inter-node tree across {num_nodes} nodes"
        )

    def simulate_all_techniques(self, message_sizes: List[int]) -> Dict[str, List[AllReduceResult]]:
        """Simulate all techniques for various message sizes"""
        results = {
            "ring": [],
            "tree": [],
            "hierarchical": [],
        }

        for msg_size in message_sizes:
            results["ring"].append(self.simulate_ring_allreduce(msg_size))
            results["tree"].append(self.simulate_tree_allreduce(msg_size))
            results["hierarchical"].append(self.simulate_hierarchical_allreduce(msg_size))

        return results


def analyze_astra_cluster_communication():
    """Analyze communication patterns for Astra's training cluster"""
    print("=" * 70)
    print("Astra Training Cluster: All-Reduce Communication Analysis")
    print("Stargate Nevada: ~400,000 Blackwell-Ultra GPUs")
    print("=" * 70)

    config = NetworkConfig(
        topology=Topology.HIERARCHICAL,
        num_gpus=400000,
        num_nodes=50000,  # 50K nodes, 8 GPUs each
        gpus_per_node=8,
        bw_nvlink=3600.0,  # NVLink 6
        bw_ib=800.0,       # NDR 800
    )

    sim = AllReduceSimulator(config)

    # Typical gradient sizes during training
    # Model parallelism: each GPU handles a shard of the model
    # Typical gradient tensor sizes: 1MB to 1GB
    message_sizes = [1 * 1024 * 1024, 10 * 1024 * 1024, 100 * 1024 * 1024,
                     512 * 1024 * 1024, 1024 * 1024 * 1024]

    print(f"\nNetwork Configuration:")
    print(f"  GPUs: {config.num_gpus:,}")
    print(f"  Nodes: {config.num_nodes:,}")
    print(f"  GPUs/node: {config.gpus_per_node}")
    print(f"  NVLink 6: {config.bw_nvlink} GB/s")
    print(f"  InfiniBand NDR 800: {config.bw_ib} GB/s")
    print()

    results = sim.simulate_all_techniques(message_sizes)

    for algo_name, algo_results in results.items():
        print(f"\n{'=' * 60}")
        print(f"Algorithm: {algo_name.upper()}")
        print(f"{'=' * 60}")
        print(f"{'Msg Size':>15} | {'Time (s)':>12} | {'BW (GB/s)':>12} | {'Steps':>8}")
        print(f"{'-' * 15} | {'-' * 12} | {'-' * 12} | {'-' * 8}")

        for r in algo_results:
            msg_mb = r.message_size / (1024 * 1024)
            print(f"{msg_mb:>10.0f} MB | {r.theoretical_time:>10.6f} | {r.bottleneck_bandwidth:>10.2f} | {r.steps:>8}")

    # Gradient compression analysis
    print("\n" + "=" * 70)
    print("Gradient Compression Impact Analysis")
    print("=" * 70)

    compression_ratios = [1.0, 0.5, 0.2, 0.1, 0.05, 0.01]
    base_msg = 512 * 1024 * 1024  # 512MB

    print(f"\nBase message size: {base_msg / (1024*1024):.0f} MB")
    print(f"{'Ratio':>8} | {'Compressed':>12} | {'Ring Time':>12} | {'Hierarchical Time':>15} | {'Speedup':>8}")
    print(f"{'-' * 8} | {'-' * 12} | {'-' * 12} | {'-' * 15} | {'-' * 8}")

    for ratio in compression_ratios:
        compressed = base_msg * ratio
        ring = sim.simulate_ring_allreduce(int(compressed))
        hier = sim.simulate_hierarchical_allreduce(int(compressed))
        speedup = 1.0 / ratio
        print(f"{ratio:>7.0%} | {compressed / (1024*1024):>8.0f} MB | {ring.theoretical_time:>10.6f} | {hier.theoretical_time:>13.6f} | {speedup:>7.2f}x")

    # Communication topology comparison
    print("\n" + "=" * 70)
    print("Topology Comparison for 100K GPU Cluster")
    print("=" * 70)

    topologies = [
        ("Ring (NVLink)", Topology.RING),
        ("Tree (IB)", Topology.TREE),
        ("Hierarchical", Topology.HIERARCHICAL),
    ]

    msg_512mb = 512 * 1024 * 1024
    print(f"\nMessage size: 512 MB")
    print(f"{'Topology':>20} | {'Time (s)':>12} | {'Effective BW':>15}")
    print(f"{'-' * 20} | {'-' * 12} | {'-' * 15}")

    for name, topo in topologies:
        c = NetworkConfig(
            topology=topo,
            num_gpus=100000,
            num_nodes=12500,
            gpus_per_node=8,
            bw_nvlink=3600.0,
            bw_ib=800.0,
        )
        s = AllReduceSimulator(c)
        if topo == Topology.RING:
            r = s.simulate_ring_allreduce(msg_512mb)
        elif topo == Topology.TREE:
            r = s.simulate_tree_allreduce(msg_512mb)
        else:
            r = s.simulate_hierarchical_allreduce(msg_512mb)
        bw_eff = msg_512mb / r.theoretical_time / 1e9
        print(f"{name:>20} | {r.theoretical_time:>10.6f} | {bw_eff:>13.2f} GB/s")

    return results


if __name__ == "__main__":
    analyze_astra_cluster_communication()

package main

import (
	"fmt"
	"math"
)

// MemoryPlanner plans GPU memory for Astra-scale model training
type MemoryPlanner struct {
	ModelParams         int64   // total parameters
	ActiveParams        int64   // active parameters per forward
	HiddenDim           int64   // hidden dimension
	NumLayers           int64   // number of transformer layers
	NumHeads            int64   // number of attention heads
	NumExperts          int64   // number of MoE experts
	TopK                int64   // top-K experts selected
	VocabSize           int64   // vocabulary size
	SeqLen              int64   // sequence length
	GlobalBatchSize     int64   // global batch size
	MicroBatchSize      int64   // micro batch size for pipeline parallelism
	TP                  int64   // tensor parallelism degree
	PP                  int64   // pipeline parallelism degree
	DP                  int64   // data parallelism degree
	EP                  int64   // expert parallelism degree
	GPUCount            int64   // total GPUs
	Precision           int     // bytes per parameter (2 for FP16/BF16, 4 for FP32)
	OptimizerStates     int     // optimizer states (typically 3 for Adam: mom, var, grad)
}

func (mp *MemoryPlanner) BytesPerParam() int64 {
	return int64(mp.Precision)
}

func (mp *MemoryPlanner) ModelMemory() float64 {
	// Model weights
	// Embedding: vocab_size * hidden_dim
	embeddingMem := float64(mp.VocabSize * mp.HiddenDim * mp.BytesPerParam())
	// Transformer layers: attention + MLP
	// Attention: 4 * hidden_dim^2 (Q, K, V, O)
	attnPerLayer := 4.0 * float64(mp.HiddenDim*mp.HiddenDim) * float64(mp.BytesPerParam())
	// MoE FFN: num_experts * top_k/total * 3 * hidden_dim * (4 * hidden_dim)
	// Each expert has gate/up/down projections with 4x expansion
	moeFFNPerLayer := float64(mp.NumExperts) * float64(mp.TopK) / float64(mp.NumExperts) *
		3.0 * float64(mp.HiddenDim) * float64(4*mp.HiddenDim) * float64(mp.BytesPerParam())
	// Layer norm: 2 * hidden_dim per layer
	lnPerLayer := 2.0 * float64(mp.HiddenDim) * float64(mp.BytesPerParam())

	transformerMem := float64(mp.NumLayers) * (attnPerLayer + moeFFNPerLayer + lnPerLayer)
	// LM head (tied with embedding typically)
	outputMem := float64(mp.VocabSize * mp.HiddenDim * mp.BytesPerParam())

	return (embeddingMem + transformerMem + outputMem) / 1e12 // convert to TB
}

func (mp *MemoryPlanner) ActivationMemory() float64 {
	// Per-micro-batch activation memory
	// This is a simplified estimate
	batchTokens := mp.MicroBatchSize * mp.SeqLen
	// Attention: each layer stores K,V cache + attention scores
	// K,V: 2 * batch * seq_len * hidden_dim * bytes
	// Attention scores: batch * num_heads * seq_len^2 * bytes
	kvPerLayer := 2.0 * float64(batchTokens*mp.HiddenDim) * float64(mp.BytesPerParam())
	attnScores := float64(batchTokens*mp.NumHeads*mp.SeqLen) * float64(mp.BytesPerParam())
	// MoE intermediate: batch * seq * top_k * 4 * hidden_dim * bytes
	moeIntermediate := float64(batchTokens*mp.TopK*4*mp.HiddenDim) * float64(mp.BytesPerParam())
	perLayer := kvPerLayer + attnScores + moeIntermediate

	return perLayer * float64(mp.NumLayers) / 1e12
}

func (mp *MemoryPlanner) OptimizerMemory() float64 {
	// Adam optimizer: 2 states (momentum, variance) + 1 gradient
	states := float64(mp.OptimizerStates)
	return mp.ModelMemory() * states / float64(mp.DP) // distributed across DP group
}

func (mp *MemoryPlanner) TotalMemoryPerGPU() float64 {
	modelMem := mp.ModelMemory() / float64(mp.TP*mp.PP*mp.EP)
	actMem := mp.ActivationMemory() / float64(mp.TP)
	optMem := mp.OptimizerMemory() / float64(mp.TP*mp.PP*mp.EP)

	return modelMem + actMem + optMem
}

func (mp *MemoryPlanner) Validate() {
	fmt.Println("=" + repeat("=", 69))
	fmt.Println("Astra (10T params) GPU Memory Planning")
	fmt.Println("=" + repeat("=", 69))
	fmt.Printf("\nModel Configuration:\n")
	fmt.Printf("  Total Parameters: %d (%.1fT)\n", mp.ModelParams, float64(mp.ModelParams)/1e12)
	fmt.Printf("  Active Parameters: %d (%.1fB)\n", mp.ActiveParams, float64(mp.ActiveParams)/1e9)
	fmt.Printf("  Hidden Dim: %d\n", mp.HiddenDim)
	fmt.Printf("  Layers: %d\n", mp.NumLayers)
	fmt.Printf("  Heads: %d\n", mp.NumHeads)
	fmt.Printf("  Experts: %d, Top-K: %d\n", mp.NumExperts, mp.TopK)
	fmt.Printf("  Vocabulary: %d\n", mp.VocabSize)
	fmt.Printf("  Sequence Length: %d\n", mp.SeqLen)
	fmt.Printf("  Global Batch: %d, Micro Batch: %d\n", mp.GlobalBatchSize, mp.MicroBatchSize)

	fmt.Printf("\nParallelism Strategy:\n")
	fmt.Printf("  Tensor Parallel (TP): %d\n", mp.TP)
	fmt.Printf("  Pipeline Parallel (PP): %d\n", mp.PP)
	fmt.Printf("  Data Parallel (DP): %d\n", mp.DP)
	fmt.Printf("  Expert Parallel (EP): %d\n", mp.EP)
	fmt.Printf("  Total GPUs: %d\n", mp.GPUCount)
	fmt.Printf("  Verification: TP*PP*DP*EP = %d\n", mp.TP*mp.PP*mp.DP*mp.EP)

	fmt.Printf("\nMemory Breakdown (per GPU):\n")
	modelMem := mp.ModelMemory() / float64(mp.TP*mp.PP*mp.EP)
	actMem := mp.ActivationMemory() / float64(mp.TP)
	optMem := mp.OptimizerMemory() / float64(mp.TP*mp.PP*mp.EP)
	total := mp.TotalMemoryPerGPU()

	fmt.Printf("  Model Weights:  %.2f TB\n", modelMem)
	fmt.Printf("  Activations:    %.2f TB\n", actMem)
	fmt.Printf("  Optimizer:      %.2f TB\n", optMem)
	fmt.Printf("  Total:          %.2f TB\n", total)

	gpuMem := 80.0 // GB, H100/H200 standard
	fmt.Printf("\n  GPU Memory:     %.0f GB\n", gpuMem)
	if total*1024 <= gpuMem {
		fmt.Printf("  ✅ Fits in GPU memory (%.1f%% utilization)\n", total*1024/gpuMem*100)
	} else {
		overflow := total*1024 - gpuMem
		fmt.Printf("  ❌ Exceeds GPU memory by %.1f GB\n", overflow)
		fmt.Printf("  Recommended: enable activation checkpointing + ZeRO-3\n")
	}

	// Compute efficiency
	fmt.Printf("\nCompute Efficiency:\n")
	flopsPerToken := 6.0 * float64(mp.ActiveParams) * float64(mp.SeqLen)
	totalFlops := flopsPerToken * float64(mp.GlobalBatchSize)
	fmt.Printf("  FLOPs per forward: %.2e\n", flopsPerToken)
	fmt.Printf("  Total FLOPs per step: %.2e\n", totalFlops)
	fmt.Printf("  FLOPs per GPU per step: %.2e\n", totalFlops/float64(mp.GPUCount))

	// Model FLOPs utilization (MFU) estimate
	peakFlops := 1979e12 * float64(mp.GPUCount) // H100 FP16 peak: 1979 TFLOPS
	theoreticalStepTime := totalFlops / peakFlops
	fmt.Printf("  Theoretical step time (100% MFU): %.3f s\n", theoreticalStepTime)
	fmt.Printf("  Estimated step time (45% MFU):    %.3f s\n", theoreticalStepTime/0.45)
	fmt.Printf("  Estimated training days:          %.1f\n", theoreticalStepTime/0.45*100000/86400)
}

func repeat(s string, n int) string {
	result := ""
	for i := 0; i < n; i++ {
		result += s
	}
	return result
}

func main() {
	// Astra-scale configuration
	planner := MemoryPlanner{
		ModelParams:     10_000_000_000_000, // 10T
		ActiveParams:    500_000_000_000,    // 500B
		HiddenDim:       40960,
		NumLayers:       192,
		NumHeads:        256,
		NumExperts:      256,
		TopK:            12,
		VocabSize:       200000,
		SeqLen:          1_500_000, // 1.5M context
		GlobalBatchSize: 4096,
		MicroBatchSize:  1,
		TP:              8,
		PP:              64,
		DP:              96,
		EP:              8,
		GPUCount:        400000,
		Precision:       2, // BF16
		OptimizerStates: 3, // Adam
	}

	// Verify GPU count
	expected := planner.TP * planner.PP * planner.DP * planner.EP
	if expected != planner.GPUCount {
		fmt.Printf("WARNING: GPU count mismatch! TP*PP*DP*EP=%d != %d\n", expected, planner.GPUCount)
	}

	planner.Validate()

	// Sensitivity analysis
	fmt.Println("\n" + repeat("=", 70))
	fmt.Println("Sensitivity Analysis: Varying Parallelism Strategy")
	fmt.Println(repeat("=", 70))

	strategies := []struct {
		name string
		tp   int64
		pp   int64
		dp   int64
		ep   int64
	}{
		{"Balanced (8,64,96,8)", 8, 64, 96, 8},
		{"High TP (16,32,96,8)", 16, 32, 96, 8},
		{"High PP (8,128,48,8)", 8, 128, 48, 8},
		{"High DP (8,64,192,4)", 8, 64, 192, 4},
		{"Aggressive EP (8,64,48,16)", 8, 64, 48, 16},
	}

	for _, s := range strategies {
		p := planner
		p.TP = s.tp
		p.PP = s.pp
		p.DP = s.dp
		p.EP = s.ep
		p.GPUCount = s.tp * s.pp * s.dp * s.ep

		modelMem := p.ModelMemory() / float64(p.TP*p.PP*p.EP)
		actMem := p.ActivationMemory() / float64(p.TP)
		optMem := p.OptimizerMemory() / float64(p.TP*p.PP*p.EP)
		total := modelMem + actMem + optMem

		fmt.Printf("\n  %s:\n", s.name)
		fmt.Printf("    Model: %.2f TB, Act: %.2f TB, Opt: %.2f TB, Total: %.2f TB\n",
			modelMem, actMem, optMem, total)
		if total*1024 <= 80 {
			fmt.Printf("    ✅ Fits in 80GB GPU\n")
		} else {
			fmt.Printf("    ❌ Exceeds 80GB by %.1f GB\n", total*1024-80)
		}
	}
}

# training_stability_analyzer.py
# Training stability and fault tolerance analysis for 100K GPU clusters
import numpy as np
from dataclasses import dataclass
from typing import List, Tuple, Optional
import math

@dataclass
class ClusterConfig:
    """Cluster configuration for stability analysis"""
    num_gpus: int
    gpu_mtbf_hours: float  # Mean Time Between Failures per GPU
    checkpoint_time_minutes: float  # Time to save a checkpoint
    checkpoint_size_tb: float  # Checkpoint size
    restore_time_minutes: float  # Time to restore from checkpoint
    network_bw_gbps: float  # Network bandwidth for checkpoint
    training_duration_days: float  # Total training duration
    steps_per_day: int  # Training steps per day
    loss_spike_probability: float  # Probability of loss spike per step

    @property
    def cluster_mtbf_minutes(self) -> float:
        """Overall cluster MTBF"""
        return self.gpu_mtbf_hours * 60 / self.num_gpus

    @property
    def expected_failures_per_day(self) -> float:
        """Expected number of failures per day"""
        return 24 * 60 / self.cluster_mtbf_minutes


@dataclass
class TrainingSimulationResult:
    """Training simulation result"""
    total_days: float
    effective_days: float
    lost_days: float
    failure_count: int
    checkpoint_count: int
    utilization: float
    total_cost_millions: float
    details: str


def simulate_training_run(
    config: ClusterConfig,
    checkpoint_interval_minutes: float = 30,
    seed: int = 42,
) -> TrainingSimulationResult:
    """Monte Carlo simulation of training run with failures"""
    rng = np.random.default_rng(seed)
    total_minutes = config.training_duration_days * 24 * 60
    cluster_mtbf = config.cluster_mtbf_minutes

    # Checkpoint overhead
    ckpt_overhead = config.checkpoint_time_minutes
    restore_overhead = config.restore_time_minutes

    # We simulate in discrete steps
    current_time = 0.0
    total_lost = 0.0
    failures = 0
    checkpoints = 0
    last_checkpoint = 0.0

    while current_time < total_minutes:
        time_to_next_failure = rng.exponential(cluster_mtbf)

        # Time until next checkpoint
        time_to_next_ckpt = checkpoint_interval_minutes - (current_time - last_checkpoint)

        if time_to_next_failure < time_to_next_ckpt:
            # Failure occurs before checkpoint
            if current_time + time_to_next_failure <= total_minutes:
                # Loss since last checkpoint
                lost = current_time + time_to_next_failure - last_checkpoint
                total_lost += lost
                failures += 1

                # Restore
                current_time = current_time + time_to_next_failure + restore_overhead
                last_checkpoint = current_time  # restore from last checkpoint
            else:
                break
        else:
            # Checkpoint occurs before failure
            if current_time + time_to_next_ckpt + ckpt_overhead <= total_minutes:
                current_time += time_to_next_ckpt + ckpt_overhead
                checkpoints += 1
                last_checkpoint = current_time
            else:
                # Can't complete checkpoint before end
                current_time = total_minutes
                break

        # Safety: prevent infinite loop
        if failures > 100_000:
            break

    effective_days = (total_minutes - total_lost) / (24 * 60)
    utilization = effective_days / config.training_duration_days

    # Cost estimate: ~$2.5 per GPU-hour for H100 cluster
    gpu_hours = config.num_gpus * config.training_duration_days * 24
    total_cost = gpu_hours * 2.5 / 1e6  # in millions

    return TrainingSimulationResult(
        total_days=config.training_duration_days,
        effective_days=effective_days,
        lost_days=total_lost / (24 * 60),
        failure_count=failures,
        checkpoint_count=checkpoints,
        utilization=utilization,
        total_cost_millions=total_cost,
        details=(
            f"Cluster MTBF: {cluster_mtbf:.2f} min | "
            f"Failures: {failures} | "
            f"Effective: {effective_days:.1f}/{config.training_duration_days:.0f} days "
            f"({utilization:.1%})"
        )
    )


def analyze_stability():
    """Analyze training stability for Astra cluster"""
    print("=" * 70)
    print("Training Stability Analysis for Astra (400K GPU Cluster)")
    print("=" * 70)

    # Base configuration
    base_config = ClusterConfig(
        num_gpus=400000,
        gpu_mtbf_hours=5000,  # H100 typical MTBF
        checkpoint_time_minutes=5,
        checkpoint_size_tb=150,
        restore_time_minutes=10,
        network_bw_gbps=1600,
        training_duration_days=180,
        steps_per_day=5000,
        loss_spike_probability=0.001,
    )

    print(f"\nBase Configuration:")
    print(f"  GPUs: {base_config.num_gpus:,}")
    print(f"  GPU MTBF: {base_config.gpu_mtbf_hours:,} hours")
    print(f"  Cluster MTBF: {base_config.cluster_mtbf_minutes:.2f} minutes")
    print(f"  Expected failures/day: {base_config.expected_failures_per_day:.1f}")
    print(f"  Training duration: {base_config.training_duration_days} days")

    # Run simulation
    print(f"\n{'=' * 60}")
    print("Monte Carlo Simulation Results")
    print(f"{'=' * 60}")

    for ckpt_interval in [10, 20, 30, 60, 120]:
        results = []
        for _ in range(10):
            r = simulate_training_run(base_config, checkpoint_interval_minutes=ckpt_interval)
            results.append(r)

        avg_util = np.mean([r.utilization for r in results])
        avg_failures = np.mean([r.failure_count for r in results])
        avg_ckpts = np.mean([r.checkpoint_count for r in results])
        avg_lost = np.mean([r.lost_days for r in results])

        print(f"\n  Checkpoint interval: {ckpt_interval} min")
        print(f"    Avg utilization: {avg_util:.1%}")
        print(f"    Avg failures:   {avg_failures:.0f}")
        print(f"    Avg checkpoints:{avg_ckpts:.0f}")
        print(f"    Avg lost days:  {avg_lost:.1f}")

    # Sensitivity analysis
    print(f"\n{'=' * 60}")
    print("Sensitivity Analysis: GPU MTBF")
    print(f"{'=' * 60}")

    for mtbf in [1000, 2000, 5000, 10000, 20000]:
        c = ClusterConfig(
            num_gpus=400000,
            gpu_mtbf_hours=mtbf,
            checkpoint_time_minutes=5,
            checkpoint_size_tb=150,
            restore_time_minutes=10,
            network_bw_gbps=1600,
            training_duration_days=180,
            steps_per_day=5000,
            loss_spike_probability=0.001,
        )
        r = simulate_training_run(c, checkpoint_interval_minutes=30)

        print(f"\n  GPU MTBF: {mtbf:>6,} hours | "
              f"Cluster MTBF: {c.cluster_mtbf_minutes:>6.2f} min | "
              f"Utilization: {r.utilization:.1%}")

    # Loss spike analysis
    print(f"\n{'=' * 60}")
    print("Loss Spike Impact Analysis")
    print(f"{'=' * 60}")

    for spike_prob in [0.0, 0.0001, 0.0005, 0.001, 0.005, 0.01]:
        c = ClusterConfig(
            num_gpus=400000,
            gpu_mtbf_hours=5000,
            checkpoint_time_minutes=5,
            checkpoint_size_tb=150,
            restore_time_minutes=10,
            network_bw_gbps=1600,
            training_duration_days=180,
            steps_per_day=5000,
            loss_spike_probability=spike_prob,
        )
        # Expected rollback steps per day due to loss spikes
        rollback_steps = c.steps_per_day * spike_prob * c.steps_per_day
        # Time lost: each rollback costs ~10 steps (recompute + re-evaluate)
        time_lost_per_day = rollback_steps * 10 / c.steps_per_day * 24  # hours
        print(f"  Spike prob: {spike_prob:.4f} | "
              f"Rollback steps/day: {rollback_steps:.1f} | "
              f"Time lost/day: {time_lost_per_day:.2f} hours")

    return base_config


if __name__ == "__main__":
    analyze_stability()