#!/usr/bin/env python3
"""Instantiate a ~8.3B DeepSeek-V3 MoE model with random weights.

Run this on a GPU-equipped machine (e.g. 1× NVL H100) so the dense
initialization completes quickly:

    python scripts/benchmarks/build_deepseek_v3_8b.py --output deepseek-v3-8b-moe
"""

from __future__ import annotations

import argparse
from pathlib import Path

import torch
from transformers import DeepseekV3Config, DeepseekV3ForCausalLM

DTYPE_MAP = {
    "float32": torch.float32,
    "bfloat16": torch.bfloat16,
    "float16": torch.float16,
}


def build_config() -> DeepseekV3Config:
    """Return a DeepSeek V3 configuration totaling ~8.3B parameters."""

    return DeepseekV3Config(
        vocab_size=32_000,
        hidden_size=3_072,
        intermediate_size=8_192,
        moe_intermediate_size=2_560,
        num_hidden_layers=20,
        num_attention_heads=24,
        num_key_value_heads=24,
        n_routed_experts=18,
        num_experts_per_tok=4,
        n_group=6,
        topk_group=4,
        kv_lora_rank=192,
        q_lora_rank=384,
        max_position_embeddings=2_048,
        rope_theta=10_000.0,
        rope_interleave=True,
        hidden_act="silu",
        initializer_range=0.02,
        attention_dropout=0.0,
        attention_bias=False,
        n_shared_experts=1,
        routed_scaling_factor=2.5,
        norm_topk_prob=True,
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Directory to save the generated model",
    )
    parser.add_argument(
        "--dtype",
        default="bfloat16",
        choices=DTYPE_MAP.keys(),
        help="Storage dtype for the checkpoint",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=0,
        help="Torch RNG seed for reproducibility",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    torch.manual_seed(args.seed)

    output_dir = args.output
    output_dir.mkdir(parents=True, exist_ok=True)

    config = build_config()
    model = DeepseekV3ForCausalLM(config)

    dtype = DTYPE_MAP[args.dtype]
    model.to(dtype=dtype)

    param_count = sum(p.numel() for p in model.parameters())
    print(f"Initialized DeepSeek-V3 MoE with {param_count / 1e9:.3f}B parameters")

    model.save_pretrained(output_dir, safe_serialization=True)
    config.save_pretrained(output_dir)
    print(f"Saved model and config to {output_dir.resolve()}")


if __name__ == "__main__":
    main()