#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""estimate_training.py —— 训练算力/成本估算器（7.5/2.4 的立项三行公式，工具化）
用法示例：
  python estimate_training.py --params 7 --tokens 1000 --gpus 8 --gpu "4090" --price 2
  python estimate_training.py --params 0.5 --tokens 5 --gpus 1 --gpu "T4" --price 0
依赖：仅标准库。
"""
import argparse

# 常见卡的有效算力（bf16, TFLOPs，按 40% 利用率折算的粗经验值；可自行改）
GPU_FLOPS_TFLOPS = {
    "4090": 165 * 0.4, "3090": 71 * 0.4, "A100": 312 * 0.4, "H100": 990 * 0.4,
    "V100": 125 * 0.35, "T4": 65 * 0.3, "L40S": 362 * 0.4, "A800": 312 * 0.4,
}

def main():
    ap = argparse.ArgumentParser(description="训练算力与成本估算（FLOPs ≈ 6ND）")
    ap.add_argument("--params", type=float, required=True, help="参数量，单位：B（如 7 表示 7B）")
    ap.add_argument("--tokens", type=float, required=True, help="训练 token 数，单位：B（如 100 表示 1000 亿）")
    ap.add_argument("--gpus", type=int, default=1, help="GPU 数量")
    ap.add_argument("--gpu", type=str, default="4090", help=f"GPU 型号，可选：{list(GPU_FLOPS_TFLOPS)}")
    ap.add_argument("--price", type=float, default=2.0, help="单卡时租价，元/小时")
    ap.add_argument("--util", type=float, default=0.4, help="MFU 假设（模型算力利用率），默认 0.4")
    args = ap.parse_args()

    if args.gpu in GPU_FLOPS_TFLOPS:
        eff_tflops = GPU_FLOPS_TFLOPS[args.gpu] / 0.4 * args.util  # 按用户 MFU 修正
    else:
        print(f"[warn] 未知 GPU '{args.gpu}'，按 100 TFLOPs 有效值估算（自行修改脚本字典）")
        eff_tflops = 100.0

    flops_total = 6 * args.params * 1e9 * args.tokens * 1e9          # FLOPs
    flops_per_sec = eff_tflops * 1e12 * args.gpus                     # 集群有效 FLOPs/s
    seconds = flops_total / flops_per_sec
    hours = seconds / 3600
    days = hours / 24
    cost = hours * args.price * args.gpus

    print("=" * 56)
    print(f"模型           : {args.params:g} B 参数")
    print(f"数据           : {args.tokens:g} B token")
    print(f"训练 FLOPs     : {flops_total:.3e}  （≈ {flops_total/1e21:.2f} ZFLOPs）")
    print(f"集群有效算力   : {args.gpus} × {eff_tflops:.0f} TFLOPs（MFU={args.util:.0%}）")
    print("-" * 56)
    print(f"预计耗时       : {hours:,.1f} 小时（≈ {days:.1f} 天）")
    print(f"预计费用       : {cost:,.0f} 元（{args.price} 元/卡时）")
    print("-" * 56)
    # 缩放参考：同预算下其他规模（Chinchilla 20:1 的对照）
    opt_tokens = args.params * 20
    print(f"[参考] Chinchilla 配比下 {args.params:g}B 的最优 token ≈ {opt_tokens:g}B；"
          f"你本次为 {args.tokens/args.params:.0f}:1（>20 即'过训'路线，2.4 §2）")
    print("=" * 56)
    print("提示：本估算是数量级工具。排期请再加 20-30% 余量（数据/断点/重试）。")

if __name__ == "__main__":
    main()