app.py

import gradio as gr
import math
from transformers import AutoConfig  # Required for Hugging Face integration

# ---- Helper Functions ---- #
def convert_params(params):
    if params == 0:
        return "0"
    size_name = ("", "K", "M", "B", "T", "P", "E", "Z", "Y")
    i = int(math.floor(math.log(params, 1000)))
    p = math.pow(1000, i)
    s = round(params / p, 2)
    return "%s %s" % (s, size_name[i])

# Set defaults for missing arguments
def set_defaults(args, defaults):
    for key, value in defaults.items():
        if getattr(args, key) is None:
            setattr(args, key, value)
    return args

# Set value if it's None, else use the config value
def set_if_none(args, key, config, config_key, defaults):
    if getattr(args, key) is None:
        setattr(args, key, config.get(config_key, defaults[key]))
    return args

# Get Hugging Face model arguments
def get_hf_model_args(args, defaults):
    if args.hf_model_name_or_path:
        try:
            config = AutoConfig.from_pretrained(args.hf_model_name_or_path, trust_remote_code=True).to_dict()
        except Exception as e:
            raise gr.Error(f"Error fetching Hugging Face model: {str(e)}")
        
        # Update arguments with Hugging Face model config values
        args.num_layers = config.get("num_hidden_layers", defaults["num_layers"])
        args.hidden_size = config.get("hidden_size", defaults["hidden_size"])
        args.num_attention_heads = config.get("num_attention_heads", defaults["num_attention_heads"])
        args.vocab_size = config.get("vocab_size", defaults["vocab_size"])
        args.sequence_length = config.get("max_position_embeddings", defaults["sequence_length"])

    return set_defaults(args, defaults)

# ---- Memory Calculation ---- #
def calc_mem(hf_model_name_or_path, num_gpus, tensor_parallel_size, pipeline_parallel_size, batch_size_per_gpu, sequence_length, vocab_size, hidden_size, num_attention_heads, num_layers, ffn_expansion_factor, is_mixed_precision, misc_mem_gib):
    
    # Define defaults
    defaults = {
        "num_layers": 44,
        "hidden_size": 6144,
        "num_attention_heads": 64,
        "vocab_size": 51200,
        "sequence_length": 2048,
        "ffn_expansion_factor": 4,
    }
    
    # Create a simple args object to simulate parsed arguments
    class Args:
        def __init__(self, **kwargs):
            for key, value in kwargs.items():
                setattr(self, key, value)

    args = Args(hf_model_name_or_path=hf_model_name_or_path, num_gpus=num_gpus, tensor_parallel_size=tensor_parallel_size,
                pipeline_parallel_size=pipeline_parallel_size, batch_size_per_gpu=batch_size_per_gpu, sequence_length=sequence_length,
                vocab_size=vocab_size, hidden_size=hidden_size, num_attention_heads=num_attention_heads, num_layers=num_layers,
                ffn_expansion_factor=ffn_expansion_factor, is_mixed_precision=is_mixed_precision, misc_mem_gib=misc_mem_gib)

    # Fetch Hugging Face model args if a model is provided
    args = get_hf_model_args(args, defaults)

    dp_degree = args.num_gpus / (args.tensor_parallel_size * args.pipeline_parallel_size)
    embed_params = 2 * args.vocab_size * args.hidden_size
    positional_params = args.hidden_size * args.sequence_length
    ln_params = 8 * args.hidden_size * args.num_layers + (2 * args.hidden_size)
    attention_params = int(2 * (1 + args.ffn_expansion_factor) * args.num_layers * args.hidden_size * args.hidden_size)
    mlp_params = args.ffn_expansion_factor * args.num_layers * args.hidden_size * args.hidden_size
    total_params = embed_params + positional_params + ln_params + attention_params + mlp_params

    bytes_per_param = 2 if args.is_mixed_precision else 4
    model_mem = total_params * bytes_per_param
    per_gpu_mem_gib = (model_mem / (args.tensor_parallel_size * args.pipeline_parallel_size)) / 1024**3 + args.misc_mem_gib

    return f"Per-GPU Memory Required for Training: {per_gpu_mem_gib:.2f} GiB"

# ---- Gradio Interface ---- #
with gr.Blocks() as demo:
    with gr.Tabs():
        with gr.TabItem("Parameter Calculation"):
            vocab_size = gr.Number(label="Vocab Size", value=51200)
            tied_embeddings = gr.Checkbox(label="Tied Embeddings", value=False)
            hidden_size = gr.Number(label="Hidden Size", value=6144)
            sequence_length = gr.Number(label="Sequence Length", value=2048)
            num_layers = gr.Number(label="Number of Layers", value=44)
            ffn_expansion_factor = gr.Number(label="FFN Expansion Factor", value=4)
            num_mlp_linears = gr.Number(label="Number of Linear Layers per MLP Block", value=2)
            kv_size_ratio = gr.Number(label="KV Size Ratio", value=1.0)

            with gr.Accordion("MoE Parameters", open=False):
                moe = gr.Checkbox(label="MoE", value=False)
                num_experts = gr.Number(label="Number of Experts", value=8)
                expert_interval = gr.Number(label="Expert Interval", value=1)
                topk = gr.Number(label="Top k Routing", value=1)

            result = gr.Textbox(label="Output", interactive=False)
            calculate_button = gr.Button("Calculate")
            calculate_button.click(calc_params, inputs=[vocab_size, tied_embeddings, hidden_size, sequence_length, num_layers, moe, num_experts, expert_interval, topk, ffn_expansion_factor, num_mlp_linears, kv_size_ratio], outputs=result)

        with gr.TabItem("Memory Calculation"):
            hf_model_name_or_path = gr.Textbox(label="HuggingFace Model Name or Path", value="")
            num_gpus = gr.Number(label="Number of GPUs", value=1)
            tensor_parallel_size = gr.Number(label="Tensor Parallel Size", value=1)
            pipeline_parallel_size = gr.Number(label="Pipeline Parallel Size", value=1)
            batch_size_per_gpu = gr.Number(label="Batch Size per GPU", value=8)
            sequence_length = gr.Number(label="Sequence Length", value=2048)
            vocab_size = gr.Number(label="Vocab Size", value=51200)
            hidden_size = gr.Number(label="Hidden Size", value=6144)
            num_attention_heads = gr.Number(label="Number of Attention Heads", value=64)
            num_layers = gr.Number(label="Number of Layers", value=44)
            ffn_expansion_factor = gr.Number(label="FFN Expansion Factor", value=4)
            is_mixed_precision = gr.Checkbox(label="Mixed Precision", value=True)
            misc_mem_gib = gr.Number(label="Misc Memory Overhead (GiB)", value=5)

            memory_result = gr.Textbox(label="Memory Calculation Result", interactive=False)
            calc_memory_button = gr.Button("Calculate Memory")
            calc_memory_button.click(calc_mem, inputs=[hf_model_name_or_path, num_gpus, tensor_parallel_size, pipeline_parallel_size, batch_size_per_gpu, sequence_length, vocab_size, hidden_size, num_attention_heads, num_layers, ffn_expansion_factor, is_mixed_precision, misc_mem_gib], outputs=memory_result)

demo.launch()