import gradio as gr def estimate_transformer_stats(batch_size, seq_len, num_layers, hidden_dim, vocab_size, show_breakdown): B = batch_size S = seq_len L = num_layers D = hidden_dim V = vocab_size # --- Parameters --- num_params = L * 12 * (D ** 2) + D * V # --- FLOPs --- (using 2 * m * n * p per matmul) attn_proj_flops = 2 * 3 * S * D * D attn_score_flops = 2 * S * D * S attn_out_proj_flops = 2 * S * D * D ffn_flops = 2 * 2 * S * D * 4 * D logit_flops = 2 * S * D * V / L total_layer_flops = attn_proj_flops + attn_score_flops + attn_out_proj_flops + ffn_flops + logit_flops total_flops = 6 * B * L * total_layer_flops output_lines = [ f"Parameters: P = 12 * L * D^2 + D * V", f" = 12 * {L} * {D}^2 + {D} * {V} = {num_params:.2e}", f"", f"FLOPs per layer (per sequence):", f" Attention Projections (QKV): 2 * 3 * S * D^2 = 2 * 3 * {S} * {D}^2 = {attn_proj_flops:.2e}", f" Attention Scores (QKᵀ): 2 * S * D * S = 2 * {S} * {D} * {S} = {attn_score_flops:.2e}", f" Attention Output Proj: 2 * S * D^2 = 2 * {S} * {D}^2 = {attn_out_proj_flops:.2e}", f" Feedforward Network: 2 * 2 * S * D * 4D = 2*2*{S}*{D}*{4*D} = {ffn_flops:.2e}", f" Logits: 2 * S * D * V / L = 2*{S}*{D}*{V} / {L} = {logit_flops:.2e}", f"", f"Layer Total FLOPs = {total_layer_flops:.2e}", f"", f"Total Training FLOPs = 6 * B * L * Layer_FLOPs", f" = 6 * {B} * {L} * {total_layer_flops:.2e} = {total_flops:.2e}" ] if show_breakdown: total_all = attn_proj_flops + attn_score_flops + attn_out_proj_flops + ffn_flops + logit_flops output_lines.append("\nComponent-wise totals across training batch:") output_lines.append(f" - QKV Projections: {attn_proj_flops * B * L:.2e} ({100 * attn_proj_flops / total_all:.1f}%)") output_lines.append(f" - Attention Scores: {attn_score_flops * B * L:.2e} ({100 * attn_score_flops / total_all:.1f}%)") output_lines.append(f" - Attention Output: {attn_out_proj_flops * B * L:.2e} ({100 * attn_out_proj_flops / total_all:.1f}%)") output_lines.append(f" - FFN: {ffn_flops * B * L:.2e} ({100 * ffn_flops / total_all:.1f}%)") output_lines.append(f" - Logits: {logit_flops * B * L:.2e} ({100 * logit_flops / total_all:.1f}%)") return "\n".join(output_lines) iface = gr.Interface( fn=estimate_transformer_stats, inputs=[ gr.Number(label="Batch Size", value=1), gr.Number(label="Sequence Length", value=2048), gr.Number(label="Number of Layers", value=24), gr.Number(label="Hidden Size (d_model)", value=2048), gr.Number(label="Vocabulary Size", value=50272), gr.Checkbox(label="Show FLOPs Breakdown", value=True), ], outputs=gr.Textbox(label="Estimates"), title="Transformer Parameter and FLOPs Estimator", description="Estimates parameter count and training FLOPs for decoder-only Transformers (like OPT/GPT). Shows formulas and per-component breakdown." ) if __name__ == "__main__": iface.launch()