# 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,
        d_ff: int = 81920,
        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:
        return 3 * self.d_model * self.d_ff

    @property
    def total_params_gating(self) -> int:
        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
        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]:
        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]
        logits = x_flat @ self.gate_weights + self.gate_bias

        if self.rng.random() < 0.3:
            noise = self.rng.normal(0, 0.01, logits.shape).astype(np.float32)
            logits = logits + noise

        top_k = min(self.config.top_k, self.config.num_experts)
        indices = np.argpartition(-logits, top_k, axis=1)[:, :top_k]
        values = np.take_along_axis(logits, indices, axis=1)

        values_exp = np.exp(values - np.max(values, axis=1, keepdims=True))
        routing_weights = values_exp / np.sum(values_exp, axis=1, keepdims=True)

        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 = expert_counts.copy()
        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

        cv = float(np.std(load) / (np.mean(load) + 1e-8))

        aux_loss = 0.0
        if self.config.use_aux_loss:
            z_loss = np.mean(np.log(np.sum(np.exp(logits - np.max(logits, axis=1, keepdims=True)), axis=1)) ** 2)
            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)

    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")

    router = TopKRouter(config)
    token_counts = [4096, 8192, 16384, 32768, 65536, 131072]
    results = []

    for n_tokens in token_counts:
        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")

    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"]

    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: {bottom_loaded}")
    print(f"  Load CV: {info['cv']:.4f}")
    print(f"  Ideal CV: {1.0 / math.sqrt(65536 * 12 / 256):.4f}")

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

    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"""
        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_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"""

    result = curve_fit(
        lambda x, E, A, B, alpha: E + A / (x[0] ** alpha) + B / (x[1] ** alpha),
        [N_vals, D_vals],
        losses,
        p0=[1.5, 400.0, 400.0, 0.35],
        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)

    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,
    )

    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}")

    print("\n" + "-" * 50)
    print("Astra-Scale Predictions")
    print("-" * 50)

    astra_params = [2e12, 5e12, 1e13, 2e13]
    astra_data = [5e12, 1e13, 2e13, 5e13]

    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}")

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

    for compute_mult in [1, 2, 5, 10, 20]:
        C_base = 6 * 1e12 * 1e12
        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}")

    print("\n" + "-" * 50)
    print("Performance Degradation Bottleneck Analysis")
    print("-" * 50)

    D_fixed = 1e13
    N_range = np.logspace(10, 14, 100)
    losses_N = [fitted_params.loss_from_params(N, D_fixed) for N in N_range]

    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("\n" + "-" * 50)
    print("Dense vs MoE Scaling Comparison")
    print("-" * 50)

    dense_N = 1.8e12
    moe_total_N = 1e13
    moe_active_N = 5e11

    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
    gpus_per_node: int = 8


@dataclass
class AllReduceResult:
    """All-Reduce operation result"""
    algorithm: str
    topology: Topology
    message_size: int
    theoretical_time: float
    bottleneck_bandwidth: float
    total_bytes_transferred: int
    bus_bandwidth_utilization: float
    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
        """
        N = self.config.num_gpus
        bw = self.config.bw_nvlink
        latency = self.config.latency_nvlink

        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
        bottleneck_bw = message_size / total_time if total_time > 0 else float('inf')

        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=(2 * message_size / total_time) / (bw * 1e9) * 100 if total_time > 0 else 0,
            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
        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

        return AllReduceResult(
            algorithm="Tree",
            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=int(total_bytes),
            bus_bandwidth_utilization=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
        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) +
            (2 * logN * reduced_size * num_nodes) +
            ((gpu_per_node - 1) / gpu_per_node * message_size * gpu_per_node)
        )

        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]]:
        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,
        gpus_per_node=8,
        bw_nvlink=3600.0,
        bw_ib=800.0,
    )

    sim = AllReduceSimulator(config)

    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

    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")

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

    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 [("Ring (NVLink)", Topology.RING), ("Tree (IB)", Topology.TREE), ("Hierarchical", Topology.HIERARCHICAL)]:
        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)
}

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

func (mp *MemoryPlanner) ModelMemory() float64 {
	// Embedding: vocab_size * hidden_dim
	embeddingMem := float64(mp.VocabSize * mp.HiddenDim * mp.BytesPerParam())
	// 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)
	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)
	outputMem := float64(mp.VocabSize * mp.HiddenDim * mp.BytesPerParam())

	return (embeddingMem + transformerMem + outputMem) / 1e12
}

func (mp *MemoryPlanner) ActivationMemory() float64 {
	batchTokens := mp.MicroBatchSize * mp.SeqLen
	kvPerLayer := 2.0 * float64(batchTokens*mp.HiddenDim) * float64(mp.BytesPerParam())
	attnScores := float64(batchTokens*mp.NumHeads*mp.SeqLen) * float64(mp.BytesPerParam())
	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 {
	states := float64(mp.OptimizerStates)
	return mp.ModelMemory() * states / float64(mp.DP)
}

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
	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 {
		fmt.Printf("  ❌ Exceeds GPU memory by %.1f GB\n", total*1024-gpuMem)
		fmt.Printf("  Recommended: enable activation checkpointing + ZeRO-3\n")
	}

	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))

	peakFlops := 1979e12 * float64(mp.GPUCount)
	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() {
	planner := MemoryPlanner{
		ModelParams:     10_000_000_000_000,
		ActiveParams:    500_000_000_000,
		HiddenDim:       40960,
		NumLayers:       192,
		NumHeads:        256,
		NumExperts:      256,
		TopK:            12,
		VocabSize:       200000,
		SeqLen:          1_500_000,
		GlobalBatchSize: 4096,
		MicroBatchSize:  1,
		TP:              8,
		PP:              64,
		DP:              96,
		EP:              8,
		GPUCount:        400000,
		Precision:       2,
		OptimizerStates: 3,
	}

	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()

	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
    checkpoint_time_minutes: float
    checkpoint_size_tb: float
    restore_time_minutes: float
    network_bw_gbps: float
    training_duration_days: float
    steps_per_day: int
    loss_spike_probability: float

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

    @property
    def expected_failures_per_day(self) -> float:
        return 24 * 60 / self.cluster_mtbf_minutes


@dataclass
class TrainingSimulationResult:
    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:
    rng = np.random.default_rng(seed)
    total_minutes = config.training_duration_days * 24 * 60
    cluster_mtbf = config.cluster_mtbf_minutes

    ckpt_overhead = config.checkpoint_time_minutes
    restore_overhead = config.restore_time_minutes

    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_to_next_ckpt = checkpoint_interval_minutes - (current_time - last_checkpoint)

        if time_to_next_failure < time_to_next_ckpt:
            if current_time + time_to_next_failure <= total_minutes:
                lost = current_time + time_to_next_failure - last_checkpoint
                total_lost += lost
                failures += 1
                current_time = current_time + time_to_next_failure + restore_overhead
                last_checkpoint = current_time
            else:
                break
        else:
            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:
                current_time = total_minutes
                break

        if failures > 100_000:
            break

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

    gpu_hours = config.num_gpus * config.training_duration_days * 24
    total_cost = gpu_hours * 2.5 / 1e6

    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():
    print("=" * 70)
    print("Training Stability Analysis for Astra (400K GPU Cluster)")
    print("=" * 70)

    base_config = 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=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"\n{'=' * 60}")
    print("Monte Carlo Simulation Results")
    print(f"{'=' * 60}")

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

        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}")

    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%}")

    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,
        )
        rollback_steps = c.steps_per_day * spike_prob * c.steps_per_day
        time_lost_per_day = rollback_steps * 10 / c.steps_per_day * 24
        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()