File size: 7,005 Bytes
8abcf2d
 
ef8c30b
8abcf2d
ef8c30b
8abcf2d
 
 
 
 
 
 
 
 
ef8c30b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4483569
 
 
 
ef8c30b
 
4483569
 
ef8c30b
4483569
ef8c30b
8abcf2d
4483569
 
ef8c30b
4483569
 
 
 
 
 
 
 
 
 
 
8abcf2d
4483569
 
 
 
 
ee7c71e
4483569
 
 
8abcf2d
4483569
 
 
 
 
 
 
 
 
 
 
 
 
 
8abcf2d
4483569
 
ef8c30b
8abcf2d
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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()