Update app.py

app.py

@@ -65,47 +65,155 @@ with gr.Blocks() as demo:
     with gr.Tabs():
         # Memory Calculation Tab
         with gr.TabItem("Memory Calculation"):
-            hf_model_name_or_path = gr.Textbox(label="HuggingFace Model Name or Path (optional)", 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)
+            hf_model_name_or_path = gr.Textbox(
+                label="HuggingFace Model Name or Path",
+                info="Name of the HuggingFace Hub repository or the local file path for it"
+            )
+            num_gpus = gr.Number(
+                label="Number of GPUs", 
+                value=1, 
+                info="Number of GPUs used for training"
+            )
+            tensor_parallel_size = gr.Number(
+                label="Tensor Parallel Size", 
+                value=1, 
+                info="Tensor parallel degree (1 if not used)"
+            )
+            pipeline_parallel_size = gr.Number(
+                label="Pipeline Parallel Size", 
+                value=1, 
+                info="Pipeline parallel degree (1 if not used)"
+            )
+            batch_size_per_gpu = gr.Number(
+                label="Batch Size per GPU", 
+                value=8, 
+                info="Batch size per GPU"
+            )
+            sequence_length = gr.Number(
+                label="Sequence Length", 
+                value=2048, 
+                info="Sequence length used for training"
+            )
+            vocab_size = gr.Number(
+                label="Vocab Size", 
+                value=51200, 
+                info="How many tokens are in the embedding layer"
+            )
+            hidden_size = gr.Number(
+                label="Hidden Size", 
+                value=6144, 
+                info="Dimension of the model's hidden size"
+            )
+            num_attention_heads = gr.Number(
+                label="Number of Attention Heads", 
+                value=64, 
+                info="Number of attention heads used in the model"
+            )
+            num_layers = gr.Number(
+                label="Number of Layers", 
+                value=44, 
+                info="Number of transformer layers used in the model"
+            )
+            ffn_expansion_factor = gr.Number(
+                label="FFN Expansion Factor", 
+                value=4, 
+                info="How much the MLP hidden size expands"
+            )
+            is_mixed_precision = gr.Checkbox(
+                label="Mixed Precision", 
+                value=True, 
+                info="Whether mixed precision is enabled"
+            )
+            misc_mem_gib = gr.Number(
+                label="Miscellaneous Memory Overhead (GiB)", 
+                value=5, 
+                info="Miscellaneous memory overhead per GPU by DL frameworks, communication libraries, etc."
+            )
 
             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)
+            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
+            )
 
-            hf_model_name_or_path.change(fn=update_from_hf_model, 
+            hf_model_name_or_path.change(
+                fn=update_from_hf_model, 
                 inputs=[hf_model_name_or_path], 
-                outputs=[num_layers, hidden_size, num_attention_heads, vocab_size, sequence_length, memory_result])
+                outputs=[num_layers, hidden_size, num_attention_heads, vocab_size, sequence_length, memory_result]
+            )
 
         # Parameter Calculation Tab
         with gr.TabItem("Parameter Calculation"):
-            hf_model_name_or_path = gr.Textbox(label="HuggingFace Model Name or Path (optional)", value="")
-            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)
+            hf_model_name_or_path = gr.Textbox(
+                label="HuggingFace Model Name or Path",
+                info="Name of the HuggingFace Hub repository or the local file path for it"
+            )
+            vocab_size = gr.Number(
+                label="Vocab Size", 
+                value=51200, 
+                info="How many tokens are in the embedding layer"
+            )
+            tied_embeddings = gr.Checkbox(
+                label="Tied Embeddings", 
+                value=False, 
+                info="Whether embeddings are tied (shared between input and output)"
+            )
+            hidden_size = gr.Number(
+                label="Hidden Size", 
+                value=6144, 
+                info="Dimension of the model's hidden size"
+            )
+            sequence_length = gr.Number(
+                label="Sequence Length", 
+                value=2048, 
+                info="Sequence length used for training"
+            )
+            num_layers = gr.Number(
+                label="Number of Layers", 
+                value=44, 
+                info="Number of transformer layers used in the model"
+            )
+            ffn_expansion_factor = gr.Number(
+                label="FFN Expansion Factor", 
+                value=4, 
+                info="How much the MLP hidden size expands"
+            )
+            num_mlp_linears = gr.Number(
+                label="Number of Linear Layers per MLP Block", 
+                value=2, 
+                info="How many linear layers per MLP block"
+            )
+            kv_size_ratio = gr.Number(
+                label="KV Size Ratio", 
+                value=1.0, 
+                info="Ratio of total query heads to key/value heads. 1.0 for MHA, 1/num_attention_heads for MQA"
+            )
 
             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)
+                moe = gr.Checkbox(
+                    label="MoE", 
+                    value=False, 
+                    info="Whether the model is MoE"
+                )
+                num_experts = gr.Number(
+                    label="Number of Experts", 
+                    value=8, 
+                    info="Number of experts for MoE"
+                )
+                expert_interval = gr.Number(
+                    label="Expert Interval", 
+                    value=1, 
+                    info="Expert interval for MoE"
+                )
+                topk = gr.Number(
+                    label="Top k Routing", 
+                    value=1, 
+                    info="Top k routing for MoE"
+                )
 
             param_result = gr.Textbox(label="Parameter Calculation Result", interactive=False)
             calc_param_button = gr.Button("Calculate Parameters")