File size: 6,107 Bytes
8abcf2d 4483569 8abcf2d 4483569 8abcf2d 4483569 8abcf2d 4483569 ee7c71e 4483569 8abcf2d 4483569 8abcf2d 4483569 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 |
import gradio as gr
import math
# Helper function to pretty-print message sizes
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])
# ---- Transformer Parameter Calculation ---- #
def calc_params(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):
if tied_embeddings:
embedding_params = hidden_size * vocab_size
else:
embedding_params = 2 * hidden_size * vocab_size
position_embedding_params = hidden_size * sequence_length
attention_params = int(2 * (1 + kv_size_ratio) * num_layers * hidden_size * hidden_size)
layernorm_params = 13 * num_layers * hidden_size
if moe:
num_expert_layers = num_layers / expert_interval
ffn_expert_params = num_mlp_linears * ffn_expansion_factor * num_expert_layers * num_experts * hidden_size * hidden_size
ffn_dense_params = num_mlp_linears * ffn_expansion_factor * (num_layers - num_expert_layers) * hidden_size * hidden_size
ffn_params = ffn_expert_params + ffn_dense_params
gating_params = num_expert_layers * hidden_size * num_experts
else:
ffn_params = num_mlp_linears * ffn_expansion_factor * num_layers * hidden_size * hidden_size
total_params = embedding_params + attention_params + ffn_params + position_embedding_params + layernorm_params
if moe:
total_params += gating_params
result = f"""
Embedding parameters: {convert_params(embedding_params)}
Attention parameters: {convert_params(attention_params)}
FFN parameters: {convert_params(ffn_params)}
{'Gating parameters: ' + convert_params(gating_params) if moe else ''}
Total Params in the Model: {convert_params(total_params)}
"""
return result
# ---- Memory Calculation Code (from the second script) ---- #
def calc_mem(args):
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.kv_size_ratio) * args.num_layers * args.hidden_size * args.hidden_size)
mlp_params = args.num_mlp_linears * args.num_layers * args.hidden_size * args.ffn_expansion_factor * args.hidden_size
total_params = embed_params + positional_params + ln_params + attention_params + mlp_params
bytes_per_param = args.low_prec_bytes_per_val if args.is_mixed_precision else args.high_prec_bytes_per_val
model_mem = total_params * bytes_per_param
per_gpu_model_mem = model_mem / (args.tensor_parallel_size * args.pipeline_parallel_size)
per_gpu_mem_gib = per_gpu_model_mem / 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=[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()
|