Added model parameter count

app.py

@@ -10,79 +10,117 @@ bpe_ps = 4.4  # determined by tokenizer
 n_ctx_base = 8192
 n_heads = 20
 n_vocab = 128000
-n_layers = 26
+n_layers = 26 # Used for BPE model and BLT Global model
 
 # Fixed local model parameters
 local_d_model = 1024
 local_g_size = 1
 local_n_ctx = 512  # in bytes
 local_n_heads = 16
-local_n_vocab = 256
+local_n_vocab = 256 # Used for BLT Local model
 local_d_model_k = local_d_model / local_n_heads
 local_d_ff_multiplier = 4
 
-def openai_flops_per_token(n_layers, n_heads, d_model, n_ctx, n_vocab, ff_ratio=4):
+def openai_flops_per_token(n_layers_val, n_heads_val, d_model_val, n_ctx_val, n_vocab_val, ff_ratio=4):
     """Open AI method for forward pass FLOPs counting of decoder-only Transformer"""
-    d_attn = d_model // n_heads
-    d_ff = d_model * ff_ratio
+    d_attn = d_model_val // n_heads_val
+    d_ff = d_model_val * ff_ratio
 
-    embeddings = 4 * d_model
-    attn_qkv = 2 * n_layers * d_model * 3 * (d_attn * n_heads)
-    attn_mask = 2 * n_layers * n_ctx * (d_attn * n_heads)
-    attn_project = 2 * n_layers * (d_attn * n_heads) * d_model
-    ff = 2 * n_layers * 2 * d_model * d_ff
-    logits = 2 * d_model * n_vocab
+    embeddings = 4 * d_model_val # FLOPs for embeddings - not parameter count
+    attn_qkv = 2 * n_layers_val * d_model_val * 3 * (d_attn * n_heads_val)
+    attn_mask = 2 * n_layers_val * n_ctx_val * (d_attn * n_heads_val)
+    attn_project = 2 * n_layers_val * (d_attn * n_heads_val) * d_model_val
+    ff = 2 * n_layers_val * 2 * d_model_val * d_ff
+    logits = 2 * d_model_val * n_vocab_val
 
     return embeddings + attn_qkv + attn_mask + attn_project + ff + logits
 
 
-def cross_attention_flops_per_token(n_layers, n_ctx_cross_attn_kv_len, d_model):
+def cross_attention_flops_per_token(n_layers_ca, n_ctx_cross_attn_kv_len, d_model_ca):
     ca_qo_proj_flops = (
         # Cross Attention QO FLOPs + backward
-        2 * 4 * d_model**2
+        2 * 4 * d_model_ca**2
     )
-    ca_context_flops = 4 * n_ctx_cross_attn_kv_len * d_model
-    return n_layers * (ca_qo_proj_flops + ca_context_flops)
+    ca_context_flops = 4 * n_ctx_cross_attn_kv_len * d_model_ca
+    return n_layers_ca * (ca_qo_proj_flops + ca_context_flops)
 
 
-def calculate_flops(blt_ps, d_model, local_n_layers):
+def calculate_flops(blt_ps, d_model_slider, local_n_layers_slider):
     # BPE calculations
     n_ctx = int(n_ctx_base / bpe_ps)
-    bpe_flops_per_token = openai_flops_per_token(n_layers, n_heads, d_model, n_ctx, n_vocab)
-    bpe_per_byte = bpe_flops_per_token / bpe_ps
+    bpe_flops_per_token_val = openai_flops_per_token(n_layers, n_heads, d_model_slider, n_ctx, n_vocab)
+    bpe_per_byte = bpe_flops_per_token_val / bpe_ps
 
     # BLT Global calculations
     blt_n_ctx = int(n_ctx_base / blt_ps)
-    blt_global_flops_per_token = openai_flops_per_token(n_layers, n_heads, d_model, blt_n_ctx, n_vocab=0)
+    blt_global_flops_per_token = openai_flops_per_token(n_layers, n_heads, d_model_slider, blt_n_ctx, n_vocab_val=0) # n_vocab=0 for global
     blt_global_flops_per_byte = blt_global_flops_per_token / blt_ps
 
     # BLT Local calculations
     local_models_transformer_flops_per_byte = openai_flops_per_token(
-        local_n_layers, local_n_heads, local_d_model, local_n_ctx, local_n_vocab
+        local_n_layers_slider, local_n_heads, local_d_model, local_n_ctx, local_n_vocab, ff_ratio=local_d_ff_multiplier
     )
     encoder_model_ca_flops_per_byte = cross_attention_flops_per_token(
-        local_n_layers/2, local_n_ctx, local_d_model
+        local_n_layers_slider / 2, local_n_ctx, local_d_model
     )
     decoder_model_ca_flops_per_byte = cross_attention_flops_per_token(
-        local_n_layers/2, local_n_ctx // blt_ps, local_d_model
+        local_n_layers_slider / 2, local_n_ctx // blt_ps, local_d_model
     )
     local_models_cross_attention_flops_per_byte = encoder_model_ca_flops_per_byte + decoder_model_ca_flops_per_byte
     local_models_flops = local_models_transformer_flops_per_byte + local_models_cross_attention_flops_per_byte
 
     # Calculate advantage
     blt_total = local_models_flops + blt_global_flops_per_byte
-    advantage = 100 * ((blt_total - bpe_per_byte) / bpe_per_byte)
+    advantage = 100 * ((blt_total - bpe_per_byte) / bpe_per_byte) if bpe_per_byte != 0 else 0
+
 
     return {
         'bpe_per_byte': bpe_per_byte,
         'blt_global': blt_global_flops_per_byte,
         'blt_local': local_models_flops,
         'blt_total': blt_total,
-        'advantage': advantage
+        'advantage': advantage,
     }
 
-def create_visualization(blt_ps, d_model, local_n_layers):
-    results = calculate_flops(blt_ps, d_model, local_n_layers)
+def format_params_display(num_params):
+    """Formats number of parameters into a string with M or B units."""
+    if num_params is None:
+        return ""
+    if abs(num_params) >= 1_000_000_000:
+        return f"{num_params / 1_000_000_000:.1f}B Params"
+    elif abs(num_params) >= 1_000_000:
+        return f"{num_params / 1_000_000:.1f}M Params"
+    else: # For numbers less than 1M
+        return f"{num_params / 1_000_000:.2f}M Params"
+
+
+def create_visualization(blt_ps, d_model_slider, local_n_layers_slider):
+    results = calculate_flops(blt_ps, d_model_slider, local_n_layers_slider)
+
+    # Calculate model parameters
+    # BPE Model Parameters: 12 * N * D^2 + 2 * V * D
+    # N = n_layers (global), D = d_model_slider, V = n_vocab (global)
+    bpe_model_params = (12 * n_layers * d_model_slider**2) + (2 * n_vocab * d_model_slider)
+
+    # BLT Model Parameters
+    # Global Component: 12 * N * D^2 (no main vocab projection)
+    # N = n_layers (global), D = d_model_slider
+    blt_global_internal_params = 12 * n_layers * d_model_slider**2
+
+    # Local Component Transformer Part: 12 * N_local * D_local^2 + 2 * V_local * D_local
+    # N_local = local_n_layers_slider, D_local = local_d_model, V_local = local_n_vocab
+    blt_local_transformer_params = (12 * local_n_layers_slider * local_d_model**2) + \
+                                   (2 * local_n_vocab * local_d_model)
+    # Local Component Cross-Attention Part: N_local * 4 * D_local^2 (estimated)
+    # This assumes 4*D^2 params per CA block (Q,K,V,O projections within local_d_model or from global to local)
+    # and local_n_layers_slider effective CA blocks.
+    blt_local_ca_params = local_n_layers_slider * 4 * local_d_model**2
+    blt_local_total_internal_params = blt_local_transformer_params + blt_local_ca_params
+
+    blt_total_model_params = blt_global_internal_params + blt_local_total_internal_params
+
+    bpe_params_str = format_params_display(bpe_model_params)
+    blt_params_str = format_params_display(blt_total_model_params)
 
     # Create the figure with subplots for better control
     fig = go.Figure()
@@ -143,24 +181,55 @@ def create_visualization(blt_ps, d_model, local_n_layers):
         ),
         barmode='stack',
         showlegend=True,
-        height=600,
+        height=650, # Increased height slightly for param text
         template="plotly_white",
         font=dict(size=14),
         bargap=0.3,
-        plot_bgcolor='white'
+        plot_bgcolor='white',
+        margin=dict(b=100) # Add bottom margin for parameter text
     )
 
     fig.add_annotation(
         x='BLT',
-        y=results['blt_total'] * 1.1,  # Position above stacked bar
-        text=f"Total: {results['blt_total']:.2e}",
+        y=results['blt_total'] * 1.05,  # Position above stacked bar, adjust if needed
+        text=f"Total FLOPs/Byte: {results['blt_total']:.2e}",
         showarrow=False,
-        font=dict(size=12, color="black", weight="bold"),
-        bgcolor="white",
+        font=dict(size=12, color="black"), # Removed bold to differentiate from param text
+        bgcolor="rgba(255,255,255,0.5)", # Slight background for readability
         bordercolor="black",
-        borderwidth=1
+        borderwidth=1,
+        xanchor='center',
+        yanchor='bottom'
+    )
+
+    # Add parameter count annotations at the bottom of bars
+    fig.add_annotation(
+        x='BPE',
+        y=0,
+        text=bpe_params_str,
+        showarrow=False,
+        xref="x",
+        yref="paper", # Use paper coordinates for y to position below x-axis
+        yanchor='top',
+        xanchor='center',
+        yshift=-35, # Adjust this value to position correctly below the bar
+        font=dict(size=11, color="black", weight="bold"),
     )
 
+    fig.add_annotation(
+        x='BLT',
+        y=0,
+        text=blt_params_str,
+        showarrow=False,
+        xref="x",
+        yref="paper",
+        yanchor='top',
+        xanchor='center',
+        yshift=-35, # Adjust this value
+        font=dict(size=11, color="black", weight="bold"),
+    )
+
+
     # Update traces to ensure proper stacking
     fig.update_traces(textfont_size=10)
 
@@ -171,14 +240,15 @@ with gr.Blocks(title="BLT vs BPE FLOPs Comparison") as demo:
     gr.Markdown("""
     # BLT vs BPE FLOPs Comparison
 
-    This interactive visualization compares the computational efficiency (FLOPs per byte) between:
+    This interactive visualization compares the computational efficiency (FLOPs per byte) and total model parameters between:
     - **BPE (Byte Pair Encoding)**: Traditional transformer architecture
     - **BLT (Byte Latent Transformer)**: Novel architecture with Global and Local components with a dynamic patch size to segment bytes.
 
     A few things you'll notice:
     1. Patch size reduces global model FLOPs but not local model
-    2. Increasing patch size and global model dimension doesn't change total FLOPs
+    2. Increasing patch size and global model dimension doesn't change total FLOPs (Note: FLOPs yes, parameters will change with d_model)
     3. In smaller BLTs, local models constitute a larger portion of the total FLOPs
+    Parameter counts are displayed below each bar.
     """)
 
     with gr.Row():
@@ -198,75 +268,88 @@ with gr.Blocks(title="BLT vs BPE FLOPs Comparison") as demo:
                 maximum=8192,
                 value=2560,
                 step=128,
-                label="Model Dimension (d_model)",
-                info="Hidden dimension size of the model"
+                label="Global Model Dimension (d_model)",
+                info="Hidden dimension size of the BPE model and BLT's Global model"
             )
 
             local_n_layers_slider = gr.Slider(
                 minimum=2,
-                maximum=24,
+                maximum=24, # Max value for local_n_layers
                 value=10,
-                step=2,
+                step=2,     # Ensure even numbers for CA split
                 label="Local Model Layers (local_n_layers)",
-                info="Number of layers in the local model"
+                info="Number of layers in the BLT's local model"
             )
 
             gr.Markdown("### Fixed Parameters")
-            gr.Markdown("""
-            - **BPE's bytes per token**: 4.4
-            - **BPE/BLT Number of Layers**: 26
-            - **BPE/BLT Number of Heads**: 20
-            - **BPE's Vocabulary Size**: 128,000
-            - **BPE/BLT Context Length**: 8,192 bytes
-            - **Local Model Dimension**: 1,024
-            - **Local Model Heads**: 16
+            gr.Markdown(f"""
+            - **BPE's bytes per token (bpe_ps)**: {bpe_ps}
+            - **BPE/BLT Global - Num Layers (n_layers)**: {n_layers}
+            - **BPE/BLT Global - Num Heads (n_heads)**: {n_heads}
+            - **BPE - Vocabulary Size (n_vocab)**: {n_vocab:,}
+            - **BPE/BLT - Context Length (n_ctx_base)**: {n_ctx_base:,} bytes
+            - **BLT Local - Model Dimension (local_d_model)**: {local_d_model}
+            - **BLT Local - Num Heads (local_n_heads)**: {local_n_heads}
+            - **BLT Local - Vocabulary Size (local_n_vocab)**: {local_n_vocab}
+            - **BLT Local - FF Multiplier (local_d_ff_multiplier)**: {local_d_ff_multiplier}
             """)
 
-            gr.Markdown("### Current Values")
+            gr.Markdown("### Current Values & Totals")
             info_text = gr.Markdown("")
 
         with gr.Column(scale=2):
-            plot = gr.Plot(label="FLOPs Comparison")
+            plot = gr.Plot(label="FLOPs Comparison & Model Parameters")
 
     # Set up interactivity
-    def update_plot(blt_ps, d_model, local_n_layers):
-        fig = create_visualization(blt_ps, d_model, local_n_layers)
+    def update_plot_and_info(blt_ps_val, d_model_val, local_n_layers_val):
+        fig = create_visualization(blt_ps_val, d_model_val, local_n_layers_val)
+        results = calculate_flops(blt_ps_val, d_model_val, local_n_layers_val)
+
+        # Recalculate parameters for info text (could also be returned by create_visualization or calculate_flops)
+        bpe_model_p = (12 * n_layers * d_model_val**2) + (2 * n_vocab * d_model_val)
+        blt_global_p = 12 * n_layers * d_model_val**2
+        blt_local_transformer_p = (12 * local_n_layers_val * local_d_model**2) + \
+                                  (2 * local_n_vocab * local_d_model)
+        blt_local_ca_p = local_n_layers_val * 4 * local_d_model**2
+        blt_local_total_internal_p = blt_local_transformer_p + blt_local_ca_p
+        blt_total_model_p = blt_global_p + blt_local_total_internal_p
 
-        # Calculate values for info display
-        results = calculate_flops(blt_ps, d_model, local_n_layers)
         info_str = f"""
         **BPE FLOPs/byte**: {results['bpe_per_byte']:.2e}
+        **BPE Total Params**: {format_params_display(bpe_model_p)}
 
         **BLT Global FLOPs/byte**: {results['blt_global']:.2e}
-
         **BLT Local FLOPs/byte**: {results['blt_local']:.2e}
-
         **BLT Total FLOPs/byte**: {results['blt_total']:.2e}
-        """
+        **BLT Total Params**: {format_params_display(blt_total_model_p)}
+        (Global: {format_params_display(blt_global_p)}, Local: {format_params_display(blt_local_total_internal_p)})
 
+        **BLT Advantage (FLOPs/byte vs BPE)**: {results['advantage']:.1f}%
+        """
         return fig, info_str
 
     # Update plot when any slider changes
+    inputs_list = [blt_ps_slider, d_model_slider, local_n_layers_slider]
     blt_ps_slider.change(
-        update_plot,
-        inputs=[blt_ps_slider, d_model_slider, local_n_layers_slider],
+        update_plot_and_info,
+        inputs=inputs_list,
         outputs=[plot, info_text]
     )
     d_model_slider.change(
-        update_plot,
-        inputs=[blt_ps_slider, d_model_slider, local_n_layers_slider],
+        update_plot_and_info,
+        inputs=inputs_list,
         outputs=[plot, info_text]
     )
     local_n_layers_slider.change(
-        update_plot,
-        inputs=[blt_ps_slider, d_model_slider, local_n_layers_slider],
+        update_plot_and_info,
+        inputs=inputs_list,
         outputs=[plot, info_text]
     )
 
     # Initial plot
     demo.load(
-        update_plot,
-        inputs=[blt_ps_slider, d_model_slider, local_n_layers_slider],
+        update_plot_and_info,
+        inputs=inputs_list,
         outputs=[plot, info_text]
     )