필사 모드: 大规模模型训练完全指南：100B+ 参数 LLM 预训练策略

N_optimal ≈ 0.1174 × C^0.4999  (参数数量)
D_optimal ≈ 1.6972 × C^0.5001  (token 数量)

D_optimal ≈ 20 × N

def estimate_training_flops(
    num_params,      # 参数数量
    num_tokens,      # 训练 token 数量
    include_backward=True
):
    """
    C ≈ 6 × N × D (前向传播 + 反向传播)
    反向传播的成本约为前向传播的 2 倍
    """
    flops_per_token = 6 * num_params
    total_flops = flops_per_token * num_tokens
    return total_flops

# 示例: LLaMA-2 7B, 2T token
flops = estimate_training_flops(7e9, 2e12)
print(f"Total FLOPs: {flops:.2e}")
# Total FLOPs: 8.40e+22

# GPU 训练时间估算 (A100 312 TFLOPs, 假设 MFU 50%)
a100_flops = 312e12  # 312 TFLOPs
mfu = 0.5           # Model FLOP Utilization
training_seconds = flops / (a100_flops * mfu)
training_hours = training_seconds / 3600

# 使用 1000 张 A100 时
num_gpus = 1000
gpu_hours_per_gpu = training_hours / num_gpus
print(f"训练时间: {gpu_hours_per_gpu:.1f} GPU-hours per GPU")
print(f"总 GPU 时间: {training_hours:.0f} GPU-hours")

data_mixture = {
    "Common Crawl (filtered)": 0.67,   # 67%
    "Books": 0.14,                     # 14%
    "GitHub": 0.045,                   # 4.5%
    "Wikipedia": 0.045,                # 4.5%
    "Gutenberg": 0.025,                # 2.5%
    "ArXiv": 0.025,                    # 2.5%
    "StackExchange": 0.02,             # 2%
}

import re
from typing import List, Optional
from dataclasses import dataclass

@dataclass
class DocumentFilter:
    min_tokens: int = 50
    max_tokens: int = 100000
    min_avg_word_length: float = 3.0
    max_symbol_ratio: float = 0.1
    min_alpha_ratio: float = 0.7

def filter_document(text: str, config: DocumentFilter) -> Optional[str]:
    """基础质量过滤器"""
    tokens = text.split()
    token_count = len(tokens)

    # 长度过滤
    if not (config.min_tokens <= token_count <= config.max_tokens):
        return None

    # 平均单词长度
    avg_word_len = sum(len(t) for t in tokens) / token_count
    if avg_word_len < config.min_avg_word_length:
        return None

    # 字母占比
    alpha_chars = sum(1 for c in text if c.isalpha())
    if alpha_chars / len(text) < config.min_alpha_ratio:
        return None

    return text

def deduplicate_documents(texts: List[str], n_gram_size: int = 13) -> List[str]:
    """
    基于 MinHash LSH 的去重
    (实际实现使用 datasketch 库)
    """
    from datasketch import MinHash, MinHashLSH

    lsh = MinHashLSH(threshold=0.8, num_perm=128)
    unique_texts = []

    for i, text in enumerate(texts):
        minhash = MinHash(num_perm=128)
        # 按 n-gram 单位做哈希
        words = text.lower().split()
        for j in range(len(words) - n_gram_size + 1):
            ngram = " ".join(words[j:j+n_gram_size])
            minhash.update(ngram.encode("utf-8"))

        if not lsh.query(minhash):
            lsh.insert(str(i), minhash)
            unique_texts.append(text)

    return unique_texts

from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel

def train_tokenizer(
    corpus_files: List[str],
    vocab_size: int = 32000,
    output_path: str = "tokenizer.json"
):
    """BPE 分词器训练(SentencePiece 方式)"""
    tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)

    trainer = BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=2,
        special_tokens=["[UNK]", "[BOS]", "[EOS]", "[PAD]"],
        show_progress=True
    )

    tokenizer.train(files=corpus_files, trainer=trainer)
    tokenizer.save(output_path)
    return tokenizer

# 计算 token 数量
def count_tokens(file_path: str, tokenizer) -> int:
    total = 0
    with open(file_path, "r") as f:
        for line in f:
            tokens = tokenizer.encode(line.strip())
            total += len(tokens.ids)
    return total

Q、K、V 投影:  [d_model, d_head × n_heads]
→ 每张 GPU: [d_model, d_head × (n_heads/tp_size)]

第一层线性层:   [d_model, d_ffn]   → 按列方向切分
第二层线性层:  [d_ffn, d_model]   → 按行方向切分
→ 每张 GPU: [d_model, d_ffn/tp_size] + [d_ffn/tp_size, d_model]

# Megatron 的 ColumnParallelLinear 概念(简化版)
class ColumnParallelLinear(nn.Module):
    """将权重矩阵按列方向切分"""
    def __init__(self, in_features, out_features, tp_size):
        super().__init__()
        self.tp_size = tp_size
        assert out_features % tp_size == 0
        local_out = out_features // tp_size
        # 每张 GPU 只持有整体的 1/tp_size
        self.weight = nn.Parameter(torch.empty(local_out, in_features))

    def forward(self, x):
        # 输入是完整大小,输出是切分后的大小
        return torch.nn.functional.linear(x, self.weight)
        # 之后需要 All-reduce 或 All-gather

class RowParallelLinear(nn.Module):
    """将权重矩阵按行方向切分"""
    def __init__(self, in_features, out_features, tp_size):
        super().__init__()
        self.tp_size = tp_size
        assert in_features % tp_size == 0
        local_in = in_features // tp_size
        # 每张 GPU 处理输入维度的 1/tp_size
        self.weight = nn.Parameter(torch.empty(out_features, local_in))

    def forward(self, x):
        # x: [batch, seq, in_features/tp_size]
        local_output = torch.nn.functional.linear(x, self.weight)
        # 用 All-reduce 汇总各 GPU 的部分和
        dist.all_reduce(local_output)
        return local_output

Attention 输入: [batch, seq/tp_size, d_model] (每张 GPU)
→ All-gather: [batch, seq, d_model]
→ Self-Attention (按列切分)
→ Reduce-scatter: [batch, seq/tp_size, d_model]
→ FFN (按行切分)

#!/bin/bash
# 用 Megatron-LM 训练 GPT-3 175B (示例)

GPUS_PER_NODE=8
NNODES=64
TP_SIZE=8    # 张量并行
PP_SIZE=16   # 流水线并行
DP_SIZE=$((GPUS_PER_NODE * NNODES / TP_SIZE / PP_SIZE))
# DP_SIZE = 8 * 64 / 8 / 16 = 4

torchrun \
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    pretrain_gpt.py \
    --num-layers 96 \
    --hidden-size 12288 \
    --num-attention-heads 96 \
    --seq-length 2048 \
    --max-position-embeddings 2048 \
    --global-batch-size 1536 \
    --train-iters 300000 \
    --lr 6e-5 \
    --lr-decay-style cosine \
    --min-lr 6e-6 \
    --lr-warmup-fraction 0.001 \
    --weight-decay 0.1 \
    --tensor-model-parallel-size $TP_SIZE \
    --pipeline-model-parallel-size $PP_SIZE \
    --micro-batch-size 1 \
    --bf16 \
    --use-flash-attn \
    --sequence-parallel \
    --data-path /data/pile_text_document \
    --vocab-file /data/gpt2-vocab.json \
    --merge-file /data/gpt2-merges.txt \
    --save /checkpoints/gpt3-175b \
    --load /checkpoints/gpt3-175b

def find_optimal_parallelism(
    world_size: int,
    num_layers: int,
    hidden_size: int,
    gpu_memory_gb: float = 80.0
) -> dict:
    """
    3D 并行化最优配置搜索
    一般规则:
    - TP: 单节点内(利用 NVLink),通常为 4 或 8
    - PP: 跨节点分散,取决于层数
    - DP: 剩余的 GPU
    """
    configs = []
    for tp in [1, 2, 4, 8]:
        for pp in range(1, world_size + 1):
            dp = world_size // (tp * pp)
            if dp < 1:
                continue
            if world_size != dp * tp * pp:
                continue
            if num_layers % pp != 0:
                continue

            # 流水线气泡比率计算
            # 微批次数 m、pp 阶段数 p 时: bubble = (p-1)/(m+p-1)
            # 这里 m = global_batch / (dp * micro_batch)
            micro_batch = 4  # 假设值
            global_batch = 2048  # 假设值
            m = global_batch // (dp * micro_batch)
            if m <= 0:
                continue
            bubble_rate = (pp - 1) / (m + pp - 1)

            configs.append({
                "tp": tp, "pp": pp, "dp": dp,
                "bubble_rate": bubble_rate,
                "layers_per_stage": num_layers // pp
            })

    # 优先选择气泡比率低、TP 合理的配置
    configs.sort(key=lambda x: (x["bubble_rate"], x["tp"]))
    return configs[:5]

# 示例: 512 张 GPU, GPT-3 规模 (96 层)
best_configs = find_optimal_parallelism(512, 96, 12288)
for c in best_configs:
    print(f"TP={c['tp']}, PP={c['pp']}, DP={c['dp']}, "
          f"Bubble={c['bubble_rate']:.2%}, Layers/Stage={c['layers_per_stage']}")

# Megatron-DeepSpeed 集成配置
deepspeed_config = {
    "train_batch_size": 2048,
    "train_micro_batch_size_per_gpu": 1,
    "gradient_accumulation_steps": 256,   # 2048 / (8 DP * 1 micro)
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 1,   # 与 TP+PP 一起使用 ZeRO-1
    },
    "gradient_clipping": 1.0,
    "wall_clock_breakdown": True,
    "steps_per_print": 10
}

# 考虑 NVLink 与 InfiniBand 的带宽差异
# NVLink 3.0: 600 GB/s (双向)
# InfiniBand HDR: 200 Gb/s = 每端口 25 GB/s

# 通信组设计原则:
# TP 组: 同一节点内的 GPU (利用 NVLink)
# PP 组: 跨节点 (利用 InfiniBand)
# DP 组: 跨节点 (用 ZeRO 最小化通信)

def setup_process_groups(tp_size, pp_size, dp_size):
    """设置进程组"""
    world_size = tp_size * pp_size * dp_size
    rank = torch.distributed.get_rank()

    # 张量并行组 (推荐同一节点内)
    for dp_rank in range(dp_size):
        for pp_rank in range(pp_size):
            ranks = [
                dp_rank * tp_size * pp_size + pp_rank * tp_size + tp_rank
                for tp_rank in range(tp_size)
            ]
            group = torch.distributed.new_group(ranks)
            if rank in ranks:
                tp_group = group

    return tp_group

import torch
import wandb
from collections import deque

class TrainingMonitor:
    def __init__(self, spike_threshold: float = 3.0, window_size: int = 100):
        self.loss_history = deque(maxlen=window_size)
        self.grad_norm_history = deque(maxlen=window_size)
        self.spike_threshold = spike_threshold
        self.spike_count = 0

    def check_loss_spike(self, current_loss: float) -> bool:
        if len(self.loss_history) < 10:
            self.loss_history.append(current_loss)
            return False

        mean_loss = sum(self.loss_history) / len(self.loss_history)
        if current_loss > mean_loss * self.spike_threshold:
            self.spike_count += 1
            print(f"SPIKE 检测: current={current_loss:.4f}, mean={mean_loss:.4f}")
            return True

        self.loss_history.append(current_loss)
        return False

    def compute_grad_norm(self, model) -> float:
        total_norm = 0.0
        for p in model.parameters():
            if p.grad is not None:
                param_norm = p.grad.data.norm(2)
                total_norm += param_norm.item() ** 2
        return total_norm ** 0.5

    def log_step(self, step: int, loss: float, model, lr: float):
        grad_norm = self.compute_grad_norm(model)
        is_spike = self.check_loss_spike(loss)

        wandb.log({
            "train/loss": loss,
            "train/grad_norm": grad_norm,
            "train/lr": lr,
            "train/is_spike": int(is_spike),
            "train/spike_count": self.spike_count,
        }, step=step)

        return grad_norm, is_spike

def adaptive_gradient_clipping(
    model,
    optimizer,
    max_norm: float = 1.0,
    clip_coef: float = 0.01
):
    """
    自适应梯度裁剪
    AGC (Adaptive Gradient Clipping): 按层裁剪
    """
    for module in model.modules():
        for name, param in module.named_parameters(recurse=False):
            if param.grad is None:
                continue

            # 比较参数范数与梯度范数
            param_norm = param.data.norm(2)
            grad_norm = param.grad.data.norm(2)

            # 裁剪阈值: 参数范数的 clip_coef 倍
            max_grad_norm = clip_coef * param_norm

            if grad_norm > max_grad_norm and grad_norm > 0:
                param.grad.data.mul_(max_grad_norm / grad_norm)

def compute_gradient_noise_scale(model, world_size: int) -> float:
    """
    梯度噪声尺度 (Gradient Noise Scale) 计算
    GNS = tr(S) / ||g||^2
    GNS 越高 → 可以使用越大的批次大小
    """
    grads = [p.grad.data for p in model.parameters() if p.grad is not None]

    # 整体梯度范数
    g_norm_sq = sum(g.norm(2).item() ** 2 for g in grads)

    # 简化的方差估计
    variance = sum((g - g.mean()).norm(2).item() ** 2 for g in grads)

    gns = variance / g_norm_sq if g_norm_sq > 0 else 0
    return gns

import math

def cosine_schedule_with_warmup(
    current_step: int,
    warmup_steps: int,
    total_steps: int,
    max_lr: float,
    min_lr: float
) -> float:
    """
    余弦衰减 + 线性预热
    绝大多数 LLM 训练中作为标准使用
    """
    if current_step < warmup_steps:
        # 线性预热
        return max_lr * current_step / warmup_steps
    else:
        # 余弦衰减
        progress = (current_step - warmup_steps) / (total_steps - warmup_steps)
        cosine_val = 0.5 * (1 + math.cos(math.pi * progress))
        return min_lr + (max_lr - min_lr) * cosine_val

# WSD (Warmup-Stable-Decay) 调度 - 近来很受欢迎
def wsd_schedule(
    current_step: int,
    warmup_steps: int,
    stable_steps: int,
    decay_steps: int,
    max_lr: float,
    min_lr: float
) -> float:
    """
    Warmup → Stable → Decay 三阶段调度
    MiniCPM、LLaMA-3 等采用
    对持续学习(continual learning)有利
    """
    if current_step < warmup_steps:
        return max_lr * current_step / warmup_steps
    elif current_step < warmup_steps + stable_steps:
        return max_lr
    else:
        decay_progress = (current_step - warmup_steps - stable_steps) / decay_steps
        return max_lr - (max_lr - min_lr) * min(decay_progress, 1.0)

# OpenAI GPT-3 论文中提出的批次大小递增策略
def get_batch_size_schedule(
    current_tokens: int,
    final_batch_tokens: int = 4_000_000,  # 最终批次大小 (以 token 为单位)
    initial_batch_tokens: int = 32_000,   # 初始批次大小
    ramp_tokens: int = 1_200_000_000      # 递增区间 (token)
) -> int:
    """
    批次大小随 token 处理量线性增加
    初期: 用小批次快速收敛
    后期: 用大批次稳定训练
    """
    if current_tokens < ramp_tokens:
        progress = current_tokens / ramp_tokens
        batch_tokens = initial_batch_tokens + (
            final_batch_tokens - initial_batch_tokens
        ) * progress
        return int(batch_tokens)
    return final_batch_tokens

import os
import torch
import torch.distributed as dist
from pathlib import Path

class DistributedCheckpointer:
    def __init__(self, save_dir: str, max_checkpoints: int = 5):
        self.save_dir = Path(save_dir)
        self.max_checkpoints = max_checkpoints
        self.save_dir.mkdir(parents=True, exist_ok=True)

    def save(
        self,
        model,
        optimizer,
        scheduler,
        step: int,
        rank: int,
        world_size: int
    ):
        """保存分布式检查点"""
        checkpoint_dir = self.save_dir / f"step_{step:08d}"
        checkpoint_dir.mkdir(parents=True, exist_ok=True)

        # 各 rank 保存自己的分片
        rank_path = checkpoint_dir / f"rank_{rank:04d}_of_{world_size:04d}.pt"

        # 模型状态 (ZeRO 分片)
        model_state = {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
            "scheduler": scheduler.state_dict() if scheduler else None,
            "step": step,
            "rank": rank,
            "world_size": world_size,
        }

        torch.save(model_state, rank_path)

        # rank 0 保存元数据
        if rank == 0:
            meta = {
                "step": step,
                "world_size": world_size,
                "timestamp": __import__("time").time()
            }
            torch.save(meta, checkpoint_dir / "meta.pt")

        # 清理旧检查点
        if rank == 0:
            self._cleanup_old_checkpoints()

    def _cleanup_old_checkpoints(self):
        checkpoints = sorted(self.save_dir.glob("step_*"))
        while len(checkpoints) > self.max_checkpoints:
            oldest = checkpoints.pop(0)
            import shutil
            shutil.rmtree(oldest)

import threading
from queue import Queue

class AsyncCheckpointer:
    """在独立线程中保存检查点 (不中断训练)"""
    def __init__(self, save_dir: str):
        self.save_dir = save_dir
        self.queue = Queue(maxsize=2)
        self.worker = threading.Thread(target=self._worker, daemon=True)
        self.worker.start()

    def _worker(self):
        while True:
            item = self.queue.get()
            if item is None:
                break
            state_dict, path = item
            torch.save(state_dict, path)
            self.queue.task_done()

    def save_async(self, model, step: int, rank: int):
        """不阻塞主线程,加入保存队列"""
        # 复制到 CPU (释放 GPU 内存)
        state_dict = {
            k: v.cpu().clone() for k, v in model.state_dict().items()
        }
        path = os.path.join(self.save_dir, f"step_{step}_rank_{rank}.pt")

        if not self.queue.full():
            self.queue.put((state_dict, path))
        else:
            # 队列已满则同步保存
            torch.save(state_dict, path)

import wandb
import numpy as np

class LLMTrainingTracker:
    def __init__(self, project_name: str, run_name: str):
        wandb.init(project=project_name, name=run_name)
        self.step = 0
        self.token_count = 0

    def log_training_step(
        self,
        loss: float,
        learning_rate: float,
        grad_norm: float,
        batch_tokens: int,
        elapsed_seconds: float
    ):
        tokens_per_second = batch_tokens / elapsed_seconds
        # 用 token 处理量估算 MFU
        # MFU = actual_flops / theoretical_peak_flops

        self.token_count += batch_tokens
        self.step += 1

        wandb.log({
            # 训练指标
            "train/loss": loss,
            "train/perplexity": np.exp(min(loss, 20)),
            "train/grad_norm": grad_norm,
            "train/learning_rate": learning_rate,
            # 效率指标
            "throughput/tokens_per_second": tokens_per_second,
            "throughput/samples_per_second": tokens_per_second / 2048,
            # 进度
            "progress/total_tokens": self.token_count,
            "progress/step": self.step,
        }, step=self.step)

    def log_evaluation(self, eval_loss: float, perplexities: dict):
        wandb.log({
            "eval/loss": eval_loss,
            "eval/perplexity": np.exp(eval_loss),
            **{f"eval/{k}_ppl": v for k, v in perplexities.items()}
        }, step=self.step)

def analyze_loss_curve(losses: list, window: int = 100) -> dict:
    """损失曲线自动分析"""
    if len(losses) < window * 2:
        return {}

    recent = losses[-window:]
    previous = losses[-2*window:-window]

    recent_mean = np.mean(recent)
    previous_mean = np.mean(previous)
    improvement = (previous_mean - recent_mean) / previous_mean

    # 突增检测
    global_mean = np.mean(losses)
    global_std = np.std(losses)
    spikes = [l for l in losses if l > global_mean + 3 * global_std]

    return {
        "recent_loss": recent_mean,
        "improvement_rate": improvement,
        "spike_count": len(spikes),
        "is_stagnant": improvement < 0.001,
        "recommendation": "建议提高学习率" if improvement < 0.001 else "正常"
    }

# 安装并运行 GPT-NeoX
git clone https://github.com/EleutherAI/gpt-neox
cd gpt-neox
pip install -r requirements/requirements.txt

# 配置文件 (configs/20B.yml)
# 开始训练
python deepy.py train configs/20B.yml

# OLMo 训练配置 (简化版)
from olmo import TrainConfig, ModelConfig, OptimizerConfig

train_config = TrainConfig(
    model=ModelConfig(
        d_model=4096,
        n_heads=32,
        n_layers=32,
        mlp_ratio=8/3,
        vocab_size=50280,
        max_sequence_length=2048,
        attention_type="flash",
    ),
    optimizer=OptimizerConfig(
        name="adamw",
        learning_rate=3e-4,
        weight_decay=0.1,
        betas=(0.9, 0.95),
    ),
    max_duration="300000ba",    # 300,000 个批次
    global_train_batch_size=2048,
    device_train_microbatch_size=2,
    precision="bf16",
    fsdp_config={
        "wrapping_strategy": "by_block",
        "precision": "bf16",
        "sharding_strategy": "FULL_SHARD",
    },
)

# torchtitan - PyTorch 原生 LLM 训练
# Meta 开发的最新预训练框架

# torchtitan 配置 (toml 格式)
config = """
[model]
name = "llama3"
flavor = "8B"
tokenizer_path = "./original/tokenizer.model"

[optimizer]
name = "AdamW"
lr = 3e-4

[training]
batch_size = 8
seq_len = 8192
max_norm = 1.0
steps = 10000
data_parallel_replicate_degree = 1
data_parallel_shard_degree = -1  # FSDP2
tensor_parallel_degree = 1
compile = true   # 启用 torch.compile

[checkpoint]
enable_checkpoint = true
folder = "outputs/checkpoint"
interval_type = "steps"
interval = 500
"""

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

from transformers.models.llama.modeling_llama import LlamaDecoderLayer

# FSDP 包装策略
auto_wrap_policy = transformer_auto_wrap_policy(
    transformer_layer_cls={LlamaDecoderLayer}
)

# FSDP2 (torch 2.x 的新 API)
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy

mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
)

for layer in model.model.layers:
    fully_shard(
        layer,
        mp_policy=mp_policy,
        reshard_after_forward=True,  # 类似 ZeRO-3
    )

fully_shard(model, mp_policy=mp_policy)

def estimate_training_cost(
    model_params: float,  # 参数数量 (例如: 7e9)
    training_tokens: float,  # 训练 token 数 (例如: 2e12)
    num_gpus: int,
    gpu_type: str = "A100-80GB",
    cloud_provider: str = "AWS"
) -> dict:
    """训练成本估算"""
    # GPU 规格 (理论峰值, bf16)
    gpu_specs = {
        "A100-80GB": {"tflops": 312, "memory_gb": 80, "nvlink": True},
        "H100-80GB": {"tflops": 989, "memory_gb": 80, "nvlink": True},
        "A10G-24GB": {"tflops": 125, "memory_gb": 24, "nvlink": False},
        "V100-32GB": {"tflops": 130, "memory_gb": 32, "nvlink": True},
    }

    # 云服务价格 (美元/小时, 大致估算)
    cloud_prices = {
        "AWS": {"A100-80GB": 3.97, "H100-80GB": 8.0, "A10G-24GB": 1.006},
        "GCP": {"A100-80GB": 3.67, "H100-80GB": 7.0},
        "Azure": {"A100-80GB": 3.40, "H100-80GB": 6.0},
        "Lambda": {"A100-80GB": 1.99, "H100-80GB": 3.29},
    }

    spec = gpu_specs[gpu_type]
    price_per_hour = cloud_prices.get(cloud_provider, {}).get(gpu_type, 0)

    # FLOPs 计算
    total_flops = 6 * model_params * training_tokens

    # 实际 MFU 假设 (通常为 35-55%)
    mfu = 0.45
    effective_tflops = spec["tflops"] * 1e12 * mfu

    # 总训练时间
    total_seconds = total_flops / (effective_tflops * num_gpus)
    total_hours = total_seconds / 3600
    total_days = total_hours / 24

    # 成本计算
    total_cost = price_per_hour * num_gpus * total_hours

    return {
        "total_flops": f"{total_flops:.2e}",
        "training_hours": f"{total_hours:.1f}",
        "training_days": f"{total_days:.1f}",
        "gpu_hours": f"{total_hours * num_gpus:.0f}",
        "estimated_cost_usd": f"{total_cost:,.0f}",
        "mfu": mfu,
    }

# 示例
result = estimate_training_cost(
    model_params=7e9,
    training_tokens=2e12,
    num_gpus=256,
    gpu_type="A100-80GB",
    cloud_provider="Lambda"
)
for k, v in result.items():
    print(f"{k}: {v}")

def measure_mfu(
    model,
    batch_tokens: int,
    elapsed_seconds: float,
    theoretical_peak_tflops: float
) -> float:
    """
    实测 MFU
    MFU = actual_FLOP_rate / theoretical_peak_FLOP_rate
    """
    # 估算模型的实际 FLOPs
    num_params = sum(p.numel() for p in model.parameters())
    actual_flops = 6 * num_params * batch_tokens  # 前向(2) + 反向(4)

    actual_tflops = actual_flops / elapsed_seconds / 1e12
    mfu = actual_tflops / theoretical_peak_tflops

    return mfu

# MFU 提升策略:
# 1. 使用 Flash Attention (最小化内存 I/O)
# 2. 启用 torch.compile (内核融合)
# 3. 谨慎使用激活检查点 (会降低速度)
# 4. 选择最优批次大小 (最大化 GPU 占用率)
# 5. 让通信与计算重叠 (FSDP/DeepSpeed 设置)

#!/usr/bin/env python3
"""
大规模 LLM 预训练完整示例
GPT 风格模型, DeepSpeed ZeRO-2 + Flash Attention
"""
import os
import time
import math
import torch
import torch.nn as nn
import deepspeed
from torch.utils.data import DataLoader, IterableDataset
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    get_cosine_schedule_with_warmup,
)

# 配置
TRAINING_CONFIG = {
    "model_name": "meta-llama/Llama-2-7b",
    "total_tokens": 1_000_000_000,   # 1B token (演示用)
    "max_seq_len": 2048,
    "global_batch_size": 1024,       # 全局批次 (以 token 计: 1024 * 2048 = 2M)
    "micro_batch_per_gpu": 4,
    "learning_rate": 3e-4,
    "min_lr": 3e-5,
    "warmup_tokens": 20_000_000,     # 2% 预热
    "weight_decay": 0.1,
    "grad_clip": 1.0,
    "save_every_steps": 1000,
    "eval_every_steps": 500,
    "log_every_steps": 10,
}

class StreamingTokenDataset(IterableDataset):
    """流式 token 数据集 (处理大容量数据)"""
    def __init__(self, data_path: str, seq_len: int, rank: int, world_size: int):
        self.data_path = data_path
        self.seq_len = seq_len
        self.rank = rank
        self.world_size = world_size

    def __iter__(self):
        # 各 rank 处理不同的数据区间
        with open(self.data_path, "rb") as f:
            f.seek(self.rank * self.seq_len * 2)  # 以 uint16 为基准
            while True:
                chunk = f.read(self.seq_len * 2 * self.world_size)
                if not chunk:
                    break
                tokens = torch.frombuffer(chunk, dtype=torch.uint16).long()
                if len(tokens) < self.seq_len + 1:
                    break
                input_ids = tokens[:self.seq_len]
                labels = tokens[1:self.seq_len + 1]
                yield {"input_ids": input_ids, "labels": labels}

def train():
    deepspeed.init_distributed()
    rank = torch.distributed.get_rank()
    world_size = torch.distributed.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))

    # 加载模型
    config = AutoConfig.from_pretrained(TRAINING_CONFIG["model_name"])
    with deepspeed.zero.Init():
        model = AutoModelForCausalLM.from_config(config)

    # 初始化 DeepSpeed 引擎
    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model,
        config={
            "train_micro_batch_size_per_gpu": TRAINING_CONFIG["micro_batch_per_gpu"],
            "gradient_accumulation_steps": (
                TRAINING_CONFIG["global_batch_size"]
                // TRAINING_CONFIG["micro_batch_per_gpu"]
                // world_size
            ),
            "bf16": {"enabled": True},
            "zero_optimization": {
                "stage": 2,
                "overlap_comm": True,
                "contiguous_gradients": True,
                "allgather_bucket_size": 2e8,
                "reduce_bucket_size": 2e8,
            },
            "gradient_clipping": TRAINING_CONFIG["grad_clip"],
            "optimizer": {
                "type": "AdamW",
                "params": {
                    "lr": TRAINING_CONFIG["learning_rate"],
                    "betas": [0.9, 0.95],
                    "eps": 1e-8,
                    "weight_decay": TRAINING_CONFIG["weight_decay"],
                }
            },
            "steps_per_print": TRAINING_CONFIG["log_every_steps"],
        }
    )

    # 数据集
    train_dataset = StreamingTokenDataset(
        "/data/train_tokens.bin",
        TRAINING_CONFIG["max_seq_len"],
        rank, world_size
    )
    train_loader = DataLoader(
        train_dataset,
        batch_size=TRAINING_CONFIG["micro_batch_per_gpu"],
        num_workers=4,
        pin_memory=True,
    )

    # 训练循环
    total_tokens = 0
    target_tokens = TRAINING_CONFIG["total_tokens"]

    for batch in train_loader:
        if total_tokens >= target_tokens:
            break

        input_ids = batch["input_ids"].to(model_engine.device)
        labels = batch["labels"].to(model_engine.device)

        # 学习率调度
        warmup_tokens = TRAINING_CONFIG["warmup_tokens"]
        lr = cosine_schedule_with_warmup(
            total_tokens, warmup_tokens, target_tokens,
            TRAINING_CONFIG["learning_rate"], TRAINING_CONFIG["min_lr"]
        )
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr

        # 前向/反向传播
        outputs = model_engine(input_ids=input_ids, labels=labels)
        loss = outputs.loss
        model_engine.backward(loss)
        model_engine.step()

        # 更新 token 计数
        batch_tokens = input_ids.numel() * world_size
        total_tokens += batch_tokens

        # 日志记录 (仅 rank 0)
        if rank == 0 and model_engine.global_steps % TRAINING_CONFIG["log_every_steps"] == 0:
            print(f"Tokens: {total_tokens/1e9:.2f}B, "
                  f"Loss: {loss.item():.4f}, "
                  f"LR: {lr:.2e}, "
                  f"PPL: {math.exp(loss.item()):.1f}")

        # 保存检查点
        if model_engine.global_steps % TRAINING_CONFIG["save_every_steps"] == 0:
            model_engine.save_checkpoint(
                "./checkpoints",
                tag=f"step_{model_engine.global_steps}"
            )

    if rank == 0:
        print("训练完成!")
        model_engine.save_checkpoint("./checkpoints", tag="final")

def cosine_schedule_with_warmup(step, warmup, total, max_lr, min_lr):
    if step < warmup:
        return max_lr * step / max(warmup, 1)
    progress = (step - warmup) / max(total - warmup, 1)
    return min_lr + (max_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * progress))

if __name__ == "__main__":
    train()

deepspeed --num_nodes=4 --num_gpus=8 \
    --master_addr=node0 --master_port=29500 \
    pretrain.py