Add application file

app.py

@@ -0,0 +1,544 @@
+import gradio as gr
+import argparse
+
+
+def num_floating_point_operations(args):
+    def calculate_layer_counts():
+        """Calculate the number of attention, Mamba, and MLP layers."""
+        if args.hybrid_override_pattern:
+            counts = {"M": 0, "*": 0, "-": 0}
+            for layer_type in args.hybrid_override_pattern:
+                if layer_type in counts:
+                    counts[layer_type] += 1
+            return counts["*"], counts["M"], counts["-"]
+        else:
+            num_attn_layers = round(args.num_layers * args.hybrid_attention_ratio)
+            num_mlp_layers = round(args.num_layers * args.hybrid_mlp_ratio)
+            num_mamba_layers = args.num_layers - num_attn_layers - num_mlp_layers
+            return num_attn_layers, num_mamba_layers, num_mlp_layers
+
+    def mlp_layer_flops(batch_size, seq_len, hidden_size, expansion=4.0, swiglu=False):
+        """Calculate FLOPs for an MLP layer."""
+        scale_factor = 3.0 / 2.0 if swiglu else 1.0
+        return 4 * expansion * scale_factor * batch_size * seq_len * hidden_size**2
+
+    def attn_layer_flops(
+        batch_size,
+        seq_len,
+        hidden_size,
+        num_heads,
+        gqa=True,
+        gqa_groups=8,
+        kv_channels=None,
+    ):
+        """Calculate FLOPs for an attention layer."""
+        p = (kv_channels * num_heads / hidden_size) if kv_channels else 1
+        g = gqa_groups if gqa else num_heads
+        return (
+            4
+            * batch_size
+            * seq_len
+            * hidden_size
+            * p
+            * (hidden_size + (hidden_size * (g / num_heads)) + (seq_len / 2))
+        )
+
+    def mamba_layer_flops(
+        batch_size, seq_len, hidden_size, state_dim=16, head_dim=64, num_groups=1
+    ):
+        """Calculate FLOPs for a Mamba layer."""
+        # Note (rwaleffe): flops estimate for scan should be updated based on new SSD kernels,
+        # but small percent of overall layer flops
+        d_in = 2 * hidden_size
+        nheads = d_in // head_dim
+        return (
+            (
+                2
+                * batch_size
+                * seq_len
+                * hidden_size
+                * (2 * d_in + 2 * num_groups * state_dim + nheads)
+            )  # in_proj
+            + (7 * batch_size * seq_len * d_in * state_dim)  # scan
+            + (2 * batch_size * seq_len * d_in * hidden_size)  # out_proj
+        )
+
+    def hybrid_flops(
+        batch_size,
+        seq_len,
+        hidden_size,
+        num_attn_layers,
+        num_mamba_layers,
+        num_mlp_layers,
+        mamba_state_dim=128,
+        mamba_head_dim=64,
+        mamba_num_groups=8,
+        num_attn_heads=32,
+        gqa=True,
+        gqa_groups=8,
+        kv_channels=None,
+        mlp_expansion=4.0,
+        swiglu=False,
+        vocab_size=256000,
+    ):
+        """Calculate total FLOPs for the hybrid model."""
+        flops_fwd = (
+            num_attn_layers
+            * attn_layer_flops(
+                batch_size,
+                seq_len,
+                hidden_size,
+                num_attn_heads,
+                gqa,
+                gqa_groups,
+                kv_channels,
+            )
+            + num_mlp_layers
+            * mlp_layer_flops(batch_size, seq_len, hidden_size, mlp_expansion, swiglu)
+            + num_mamba_layers
+            * mamba_layer_flops(
+                batch_size,
+                seq_len,
+                hidden_size,
+                mamba_state_dim,
+                mamba_head_dim,
+                mamba_num_groups,
+            )
+            + (
+                2 * batch_size * seq_len * hidden_size * vocab_size
+            )  # logits computation
+        )
+        return flops_fwd * 3
+
+    def transformer_flops():
+        """Calculate FLOPs for a standard Transformer model."""
+        # TODO(helenn/dnarayanan): Refactor this to reuse the helper methods.
+        # Attention projection size.
+        query_projection_size = args.kv_channels * args.num_attention_heads
+        query_projection_to_hidden_size_ratio = query_projection_size / args.hidden_size
+        # Group Query Attention.
+        if not args.group_query_attention:
+            args.num_query_groups = args.num_attention_heads
+        # MoE.
+        if args.num_experts is None:
+            # Every Transformer MLP is dense.
+            num_dense_layers = args.num_layers
+            num_moe_layers = 0
+            num_experts_routed_to = 0
+            last_layer_is_moe = 0
+        else:
+            # Calculate number of dense and MoE Transformer MLPs.
+            if isinstance(args.moe_layer_freq, int):
+                moe_layer_pattern = [
+                    1 if (i % args.moe_layer_freq == 0) else 0
+                    for i in range(args.num_layers)
+                ]
+            elif isinstance(args.moe_layer_freq, list):
+                moe_layer_pattern = args.moe_layer_freq
+            else:
+                raise RuntimeError("Illegal --moe-layer-freq argument provided!")
+            assert len(moe_layer_pattern) == args.num_layers, (
+                f"Invalid length of moe_layer_pattern: {len(moe_layer_pattern)}, "
+                f"expected {args.num_layers}, "
+                f"current moe layer pattern: {args.moe_layer_freq}"
+            )
+            num_moe_layers = sum(
+                moe_layer_pattern
+            )  # Number of 1s in `moe_layer_pattern`.
+            num_dense_layers = args.num_layers - num_moe_layers
+            num_experts_routed_to = args.moe_router_topk
+            last_layer_is_moe = moe_layer_pattern[-1]
+
+        if args.mtp_num_layers is not None:
+            mtp_num_layers = args.mtp_num_layers
+            num_moe_layers += last_layer_is_moe * mtp_num_layers
+            num_dense_layers += (1 - last_layer_is_moe) * mtp_num_layers
+            num_layers = args.num_layers + mtp_num_layers
+        else:
+            mtp_num_layers = 0
+            num_layers = args.num_layers
+
+        moe_ffn_hidden_size = (
+            args.moe_ffn_hidden_size
+            if args.moe_ffn_hidden_size is not None
+            else args.ffn_hidden_size
+        )
+        shared_expert_ffn_hidden_size = (
+            0
+            if args.moe_shared_expert_intermediate_size is None
+            else args.moe_shared_expert_intermediate_size
+        )
+        # SwiGLU.
+        gated_linear_multiplier = 3 / 2 if args.swiglu else 1
+
+        # The 12x term below comes from the following factors; for more details, see
+        # "APPENDIX: FLOATING-POINT OPERATIONS" in https://arxiv.org/abs/2104.04473.
+        # - 3x: Each GEMM in the model needs to be performed 3 times (forward pass,
+        #       backward wgrad [weight gradient], backward dgrad [data gradient]).
+        # - 2x: GEMMs of a particular size are stacked twice in the standard Transformer model
+        #       architectures implemented in this codebase (e.g., h->ffn_h GEMM and ffn_h->h GEMM
+        #       in MLP layer).
+        # - 2x: A GEMM of a m*n tensor with a n*k tensor requires 2mnk floating-point operations.
+        expansion_factor = 3 * 2 * 2
+
+        if args.multi_latent_attention:
+            assert not args.group_query_attention
+            """
+            Basic arithmetic
+            let B is batch size, s is seq_len, h is embedding dim,
+            for one self_attnetion block (prenorm is not included)
+            qkv projection:  6Bsh^2
+            attn:            2Bs^2h
+            attn over value: 2Bs^2h
+            oproj:           2Bsh^2
+
+            references
+            https://arxiv.org/abs/2305.10403
+            https://arxiv.org/abs/2205.05198
+            """
+            ## MLA
+            if args.q_lora_rank is None:
+                q_term = (
+                    args.hidden_size
+                    * args.num_attention_heads
+                    * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+                )
+            else:
+                q_term = args.q_lora_rank * (
+                    args.hidden_size
+                    + args.num_attention_heads
+                    * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+                    + 1
+                )
+            self_attn_term = (
+                3
+                * 2  # fwd(1) + bwd(2) *FMA
+                * num_layers
+                * (
+                    ## q lora + rope + q norm
+                    q_term
+                    ## kv lora + rope + kv norm
+                    + args.kv_lora_rank
+                    * (
+                        args.hidden_size
+                        + args.num_attention_heads
+                        * (args.qk_head_dim + args.v_head_dim)
+                        + 1
+                    )
+                    + args.hidden_size * args.qk_pos_emb_head_dim
+                    ## o proj
+                    + (args.num_attention_heads * args.v_head_dim) * args.hidden_size
+                    ## core attn
+                    + args.seq_length
+                    * (
+                        args.num_attention_heads
+                        * (args.qk_head_dim + args.qk_pos_emb_head_dim)
+                    )
+                    / 2
+                    + args.seq_length * args.num_attention_heads * args.v_head_dim / 2
+                )
+            )
+
+        else:
+            ## MHA or GQA
+            self_attn_term = (
+                expansion_factor
+                * num_layers
+                * args.hidden_size
+                * args.hidden_size
+                * (
+                    (
+                        1
+                        + (args.num_query_groups / args.num_attention_heads)
+                        # # Only half of the attention matrix is non-zero and needs to be multiplied with V.
+                        + (args.seq_length / args.hidden_size / 2)
+                    )
+                    * query_projection_to_hidden_size_ratio
+                )
+            )
+
+        total_floating_point_operations = (
+            args.batch_size
+            * args.seq_length
+            * (
+                # MLP
+                expansion_factor
+                * num_layers
+                * args.hidden_size
+                * (
+                    # dense layer (deepseek v2, v3 style)
+                    (args.ffn_hidden_size * gated_linear_multiplier)
+                    * (num_dense_layers / num_layers)
+                    # routed experts
+                    + (
+                        moe_ffn_hidden_size
+                        * num_experts_routed_to
+                        * gated_linear_multiplier
+                    )
+                    * (num_moe_layers / num_layers)
+                    # Shared Experts.
+                    + (shared_expert_ffn_hidden_size * gated_linear_multiplier)
+                    * (num_moe_layers / num_layers)
+                )
+                # Self Attention
+                + self_attn_term
+                # MTP norms and proj
+                + 3
+                * 2
+                * mtp_num_layers
+                * (
+                    # MTP eh norm + final nrom
+                    3 * args.hidden_size
+                    # MTH eh proj
+                    + 2 * args.hidden_size * args.hidden_size
+                )
+                # Logit.
+                + 3
+                * 2
+                * args.hidden_size
+                * args.padded_vocab_size
+                * (mtp_num_layers + 1)
+            )
+        )
+        return total_floating_point_operations
+
+    # Main entrypoint for FLOPs calculation.
+    if args.is_hybrid_model:
+        # Calculate the number of each type of layer.
+        num_attn_layers, num_mamba_layers, num_mlp_layers = calculate_layer_counts()
+
+        # Compute hybrid model FLOPs.
+        return hybrid_flops(
+            batch_size=args.batch_size,
+            seq_len=args.seq_length,
+            hidden_size=args.hidden_size,
+            num_attn_layers=num_attn_layers,
+            num_mamba_layers=num_mamba_layers,
+            num_mlp_layers=num_mlp_layers,
+            mamba_state_dim=args.mamba_state_dim,
+            mamba_head_dim=args.mamba_head_dim,
+            mamba_num_groups=args.mamba_num_groups,
+            num_attn_heads=args.num_attention_heads,
+            gqa=args.group_query_attention,
+            gqa_groups=args.num_query_groups,
+            kv_channels=args.kv_channels,
+            mlp_expansion=args.ffn_hidden_size / args.hidden_size,
+            swiglu=args.swiglu,
+            vocab_size=args.padded_vocab_size,
+        )
+    else:
+        # Compute standard Transformer model FLOPs.
+        return transformer_flops()
+
+
+def calculate_flops(args):
+    model_flops = num_floating_point_operations(args)
+    flops_per_token = model_flops / (args.batch_size * args.seq_length)
+    print(f"FLOPs Per Iteration: {model_flops}\nFLOPs Per Token: {flops_per_token}")
+    return model_flops
+
+
+def calculate_mfu(model_flops, *, iter_elapsed_time, num_p800_cards):
+    assert (
+        model_flops and iter_elapsed_time and num_p800_cards
+    ), "Iter elapsed time and P800 cards must be provided"
+    mfu = model_flops / (iter_elapsed_time * num_p800_cards * 3.5e14)
+    print(f"MFU P800 bf16: {mfu:.2%}")
+
+
+def calculate_mfu_web(  is_hybrid_model, group_query_attention, swiglu, num_layers, hidden_size,
+                        ffn_hidden_size, padded_vocab_size, num_attention_heads, kv_channels,
+                        num_experts, moe_layer_freq, moe_router_topk, moe_ffn_hidden_size, moe_shared_expert_intermediate_size,
+                        multi_latent_attention, q_lora_rank, kv_lora_rank, qk_head_dim, v_head_dim, qk_pos_emb_head_dim,
+                        mtp_num_layers, seq_length, batch_size, iter_elapsed_time, num_p800_cards
+                    ):
+    is_hybrid_model = True if is_hybrid_model == "True" else False
+    group_query_attention = True if group_query_attention == "True" else False
+    swiglu = True if swiglu == "True" else False
+    multi_latent_attention = True if multi_latent_attention == "True" else False
+
+    '''
+    为了直接调用calculate_flops(args)接口，这里将参数直接打包
+    '''
+    class parameter:
+        def __init__(self, 
+                        is_hybrid_model, group_query_attention, swiglu, num_layers, hidden_size,
+                        ffn_hidden_size, padded_vocab_size, num_attention_heads, kv_channels,
+                        num_experts, moe_layer_freq, moe_router_topk, moe_ffn_hidden_size, moe_shared_expert_intermediate_size,
+                        multi_latent_attention, q_lora_rank, kv_lora_rank, qk_head_dim, v_head_dim, qk_pos_emb_head_dim,
+                        mtp_num_layers, seq_length, batch_size, iter_elapsed_time, num_p800_cards, 
+                        hybrid_override_pattern=None):
+            self.is_hybrid_model = is_hybrid_model
+            self.group_query_attention = group_query_attention
+            self.swiglu = swiglu
+            self.num_layers = num_layers
+            self.hidden_size = hidden_size
+            self.ffn_hidden_size = ffn_hidden_size
+            self.padded_vocab_size = padded_vocab_size
+            self.num_attention_heads = num_attention_heads
+            self.kv_channels = kv_channels
+            self.num_experts = num_experts
+            self.moe_layer_freq = moe_layer_freq
+            self.moe_router_topk = moe_router_topk
+            self.moe_ffn_hidden_size = moe_ffn_hidden_size
+            self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size
+            self.multi_latent_attention = multi_latent_attention
+            self.q_lora_rank = q_lora_rank
+            self.kv_lora_rank = kv_lora_rank
+            self.qk_head_dim = qk_head_dim
+            self.v_head_dim = v_head_dim
+            self.qk_pos_emb_head_dim = qk_pos_emb_head_dim
+            self.mtp_num_layers = mtp_num_layers
+            self.seq_length = seq_length
+            self.batch_size = batch_size
+            self.iter_elapsed_time = iter_elapsed_time
+            self.num_p800_cards = num_p800_cards
+            self.hybrid_override_pattern = hybrid_override_pattern
+
+    mfu_parameter = parameter(is_hybrid_model, group_query_attention, swiglu, num_layers, hidden_size,
+                        ffn_hidden_size, padded_vocab_size, num_attention_heads, kv_channels,
+                        num_experts, moe_layer_freq, moe_router_topk, moe_ffn_hidden_size, moe_shared_expert_intermediate_size,
+                        multi_latent_attention, q_lora_rank, kv_lora_rank, qk_head_dim, v_head_dim, qk_pos_emb_head_dim,
+                        mtp_num_layers, seq_length, batch_size, iter_elapsed_time, num_p800_cards, 
+                        hybrid_override_pattern=None)
+
+    model_flops = num_floating_point_operations(mfu_parameter)
+    flops_per_token = model_flops / (batch_size * seq_length)
+    print(f"FLOPs Per Iteration: {model_flops}\nFLOPs Per Token: {flops_per_token}")
+
+    assert (
+            model_flops and iter_elapsed_time and num_p800_cards
+        ), "Iter elapsed time and P800 cards must be provided"
+
+    mfu = model_flops / (iter_elapsed_time * num_p800_cards * 3.5e14)
+    print(f"MFU P800 bf16: {mfu:.2%}")
+    return model_flops, flops_per_token, "{:.2f}".format(mfu * 100)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    args = parser.parse_args()
+
+    # Standard Transformer config
+    args.is_hybrid_model = False
+    args.group_query_attention = False
+    args.swiglu = True
+    args.num_layers = 61
+    args.hidden_size = 7168
+    args.ffn_hidden_size = 18432
+    args.padded_vocab_size = 100002
+    args.num_attention_heads = 128
+    args.kv_channels = 128
+
+    # MoE config
+    args.num_experts = 256
+    args.moe_layer_freq = 1
+    args.moe_router_topk = 8
+    args.moe_ffn_hidden_size = 2048
+    args.moe_shared_expert_intermediate_size = 2048
+
+    # MLA config
+    args.multi_latent_attention = True
+    args.q_lora_rank = 1536
+    args.kv_lora_rank = 512
+    args.qk_head_dim = 128
+    args.v_head_dim = 128
+    args.qk_pos_emb_head_dim = 64
+
+    # MTP config
+    args.mtp_num_layers = 1
+
+    # Data config
+    args.seq_length = 4096
+    args.batch_size = 1024
+
+    # mfu config
+    args.iter_elapsed_time = 100
+    args.num_p800_cards = 512
+
+    #calculate_mfu(calculate_flops(args), iter_elapsed_time=args.iter_elapsed_time, num_p800_cards=args.num_p800_cards)
+    with gr.Blocks(title="Compute MFU") as demo:
+        gr.Markdown("## Compute MFU")
+        
+        with gr.Group() as custom_group:
+            gr.Markdown("Standard Transformer config:")
+            with gr.Row():
+                is_hybrid_model = gr.Dropdown(["True", "False"], 
+                                                label="hybrid model", 
+                                                value="True" if args.is_hybrid_model else "False")
+
+                group_query_attention = gr.Dropdown(["True", "False"], 
+                                                label="group query attention", 
+                                                value="True" if args.group_query_attention else "False")
+
+                swiglu = gr.Dropdown(["True", "False"], 
+                                        label="swiglu", 
+                                        value="True" if args.swiglu else "False")
+
+                num_layers = gr.Number(label="num layers", value=args.num_layers, precision=0)
+                hidden_size = gr.Number(label="hidden size", value=args.hidden_size, precision=0)
+                ffn_hidden_size = gr.Number(label="ffn hidden size", value=args.ffn_hidden_size, precision=0)
+                padded_vocab_size = gr.Number(label="padded vocab size", value=args.padded_vocab_size, precision=0)
+                num_attention_heads = gr.Number(label="num attention heads", value=args.num_attention_heads, precision=0)
+                kv_channels = gr.Number(label="kv channels", value=args.kv_channels, precision=0)
+
+        with gr.Group() as custom_group:
+            gr.Markdown("MoE config:")
+            with gr.Row():
+                num_experts = gr.Number(label="num experts", value=args.num_experts, precision=0)
+                moe_layer_freq = gr.Number(label="moe layer freq", value=args.moe_layer_freq, precision=0)
+                moe_router_topk = gr.Number(label="moe router topk", value=args.moe_router_topk, precision=0)
+                moe_ffn_hidden_size = gr.Number(label="moe ffn hidden size", value=args.moe_ffn_hidden_size, precision=0)
+                moe_shared_expert_intermediate_size = gr.Number(label="moe shared expert intermediate size", value=args.moe_shared_expert_intermediate_size, precision=0)
+
+        with gr.Group() as custom_group:
+            gr.Markdown("MLA config:")
+            with gr.Row():
+                multi_latent_attention = gr.Dropdown(["True", "False"], 
+                                                label="multi_latent_attention", 
+                                                value="True" if args.multi_latent_attention else "False")
+                q_lora_rank = gr.Number(label="q lora rank", value=args.q_lora_rank, precision=0)
+                kv_lora_rank = gr.Number(label="kv lora rank", value=args.kv_lora_rank, precision=0)
+                qk_head_dim = gr.Number(label="qk head dim", value=args.qk_head_dim, precision=0)
+                v_head_dim = gr.Number(label="v head dim", value=args.v_head_dim, precision=0)
+                qk_pos_emb_head_dim = gr.Number(label="qk pos emb head dim", value=args.qk_pos_emb_head_dim, precision=0)
+
+        with gr.Group() as custom_group:
+            with gr.Row():
+                with gr.Group():
+                    gr.Markdown("MTP config:")
+                    mtp_num_layers = gr.Number(label="mtp num layers", value=args.mtp_num_layers, precision=0)
+
+                with gr.Group():
+                    gr.Markdown("Data config:")
+                    with gr.Row():
+                        seq_length = gr.Number(label="seq length", value=args.seq_length, precision=0)
+                        batch_size = gr.Number(label="batch size", value=args.batch_size, precision=0)
+
+                with gr.Group():
+                    gr.Markdown("MFU config:")
+                    with gr.Row():
+                        iter_elapsed_time = gr.Number(label="iter elapsed time", value=args.iter_elapsed_time, precision=0)
+                        num_p800_cards = gr.Number(label="num p800 cards", value=args.num_p800_cards, precision=0)
+
+        # 计算结果显示控件
+        with gr.Group() as custom_group:
+            gr.Markdown("Compute results:")
+            with gr.Row():
+                model_flops = gr.Number(label="model flops", precision=0)
+                flops_per_token = gr.Number(label="flops per token", precision=0)
+                # mfu = gr.Number(label="mfu", precision=0)
+                mfu = gr.Textbox(label="MFU P800 bf16")
+
+        # 计算按钮
+        btn = gr.Button("Calculate")
+        btn.click(  fn=calculate_mfu_web, 
+                    inputs=[is_hybrid_model, group_query_attention, swiglu, num_layers, hidden_size,
+                        ffn_hidden_size, padded_vocab_size, num_attention_heads, kv_channels,
+                        num_experts, moe_layer_freq, moe_router_topk, moe_ffn_hidden_size, moe_shared_expert_intermediate_size,
+                        multi_latent_attention, q_lora_rank, kv_lora_rank, qk_head_dim, v_head_dim, qk_pos_emb_head_dim,
+                        mtp_num_layers, seq_length, batch_size, iter_elapsed_time, num_p800_cards], 
+                    outputs=[model_flops, flops_per_token, mfu]
+                )
+
+    # 启动 Gradio 应用
+    demo.launch()