finish up calcs

app.py

@@ -1,9 +1,11 @@
 import gradio as gr
 from defaults import DEFAULTS
 
-def greet(name, intensity)->str:
+
+def greet(name, intensity) -> str:
     return "Hello, " + name + "!" * int(intensity)
 
+
 def create_parallelism_block():
     with gr.Column():
         gr.Markdown("# Parallelism Parameters")
@@ -13,29 +15,33 @@ def create_parallelism_block():
         ep = gr.Number(label="Expert Parallelism", value=1, interactive=True)
         return tp, pp, cp, ep
 
+
 def create_model_block():
     with gr.Column():
         gr.Markdown("# Model Parameters")
         layers = gr.Number(label="Number of Layers", value=32, interactive=True)
         vocab = gr.Number(label="Vocab Size", value=32000, interactive=True)
         hidden = gr.Number(label="Hidden Dim", value=4096, interactive=True)
-        intermediate = gr.Number(label="Intermediate Dim", value=11008, interactive=True)
+        intermediate = gr.Number(
+            label="Intermediate Dim", value=11008, interactive=True
+        )
         presets = gr.Dropdown(list(DEFAULTS.keys()), label="Presets", interactive=True)
         return layers, vocab, hidden, intermediate, presets
 
+
 def create_training_block():
     with gr.Column():
-        gr.Markdown('# Training Parameters')
+        gr.Markdown("# Training Parameters")
         seq_len = gr.Number(label="Sequence Length", value=8192, interactive=True)
         batch_size = gr.Number(label="Batch Size", value=8, interactive=True)
         return seq_len, batch_size
 
-def calculate(*args)->int:
+
+def calculate(*args) -> int:
     out = 1
     for arg in args:
         out *= arg
     return arg
-    
 
 
 with gr.Blocks() as demo:
@@ -46,8 +52,8 @@ with gr.Blocks() as demo:
             seq_len, batch_size = create_training_block()
         calculate_button = gr.Button("Calculate")
         output = gr.Number(label="Output")
-        
-        calculate_button.click(fn=calculate, inputs=[tp,pp,cp,ep],outputs=output)
-        
+
+        calculate_button.click(fn=calculate, inputs=[tp, pp, cp, ep], outputs=output)
+
 
 demo.launch()

calculator.py

@@ -0,0 +1,162 @@
+from state import Model as Model, Parallelism, Training
+
+
+class MemoryCalculation:
+    def __init__(self, model: Model, parallelism: Parallelism, training: Training):
+        self.model = model
+        self.parallelism = parallelism
+        self.training = training
+
+    def calculate_num_parameters(self) -> float:
+        # https://michaelwornow.net/2024/01/18/counting-params-in-transformer
+        # https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=memory_usage_in_transformers
+
+        # Biases are not added/omitted on a per-model basis for simplicity.
+        # Just include them where they could appear. They're small in comparison to weights anyway and it forms an upper bound.
+
+        b, s = self.training.batch_size, self.training.sequence_length
+        h, i, l, v, e = (
+            self.model.hidden_dim,
+            self.model.intermediate_size,
+            self.model.num_layers,
+            self.model.vocab_size,
+            self.model.experts,
+        )
+        tp, pp, ep = (
+            self.parallelism.tensor_parallelism,
+            self.parallelism.pipeline_parallelism,
+            self.parallelism.expert_parallelism,
+        )
+
+        # Embedding layers
+        input_embedding = v * h / tp
+        unembedding = 0
+        if not self.model.weight_tied_embeddings:
+            unembedding = h * v / tp
+
+        # Attention
+        # weights and biases = *2
+        layer_norm_attn_in = 2 * h # not tp sharded
+        qkv = 3 * h * h / tp
+        attn_output_proj = h * h + h / tp
+        attn = layer_norm_attn_in + qkv + attn_output_proj
+
+        # MLP
+        layer_norm_mlp_in = 2 * h # not tp sharded
+        router = h * e + e # assuming replicated for simplicity
+        mlp_up_proj = h * i + i / tp
+        mlp_gate_proj = h * i + i / tp
+        mlp_down_proj = i * h + h / tp
+        expert = mlp_up_proj + mlp_gate_proj + mlp_down_proj
+        experts = expert * e / ep
+        mlp = layer_norm_mlp_in + router + experts
+
+        layer = attn + mlp
+        layers = layer * l
+
+
+        final_layer_norm = 2 * h # not tp sharded
+
+        # pp and weight tying makes knowing where to embed layer challenging
+        # going to assume "worst" case and it's at the end with final layer norm
+        # even though that's pretty smalle
+        if pp == 1:
+            total_params = input_embedding + layers + unembedding + final_layer_norm
+        if pp > 1:
+            total_params = max(input_embedding, unembedding) + layers/pp + final_layer_norm
+        return total_params
+
+    def calculate_parameter_memory(self) -> float:
+        return (
+            self.calculate_num_parameters() * 4
+        )  # assuming 4 bytes (32 bits) per parameter
+
+    def calculate_activation_parameters(self) -> float:
+        # https://blog.eleuther.ai/transformer-math/#activations-and-batch-size
+        # https://arxiv.org/abs/2205.05198
+        # pp not considered since most pp schemes will run multiple concurrent batches to reduce the bubble
+        b, s = self.training.batch_size, self.training.sequence_length
+        h, i, l, v, e, ae = (
+            self.model.hidden_dim,
+            self.model.intermediate_size,
+            self.model.num_layers,
+            self.model.vocab_size,
+            self.model.active_experts,
+        )
+        tp, cp, pp, ep = (
+            self.parallelism.tensor_parallelism,
+            self.parallelism.context_parallelism,
+            self.parallelism.pipeline_parallelism,
+            self.parallelism.expert_parallelism,
+        )
+        if self.training.gradient_checkpointing:
+            # full recomputation
+            embed = 0
+            layer = s * b * h / cp / tp  # only keep initial input to layer
+            layers = layer * l
+            embed = 0
+            final_layer_out = (
+                s * b * h / cp / tp
+            )  # both sequence and tensor parallelism
+            final_norm = s * b * h / cp / tp  # both sequence and tensor parallelism
+            unembed = s * b * v / cp / tp
+            logits = s * b * v / cp / tp  # both vocab and tensor parallelism
+            num_params = (
+                embed + layers + final_layer_out + final_norm + unembed + logits
+            )
+            return num_params
+        else:
+            # assume flash attention ie do selective recomputation
+            # assume tensor parallel + sequence parallel as described in https://arxiv.org/abs/2205.05198
+            # the variables calculate the activation outputs
+            # Attention Block
+            layer_in = s * b * h / cp / tp  # both sequence and context parallelism
+            attn_norm = s * b * h / cp / tp  # both sequence and context parallelism
+            flash = s * b * h / cp / tp
+            # everything else is recalculated by flash attention
+            projection = s * b * h / cp / tp
+            attn = layer_in + attn_norm + flash + projection
+            # MLP Block
+            mlp_norm = s * b * h / cp / tp  # both sequence and context parallelism
+            router = (
+                s * b * e / cp / tp
+            )  # makes sense to sp shard if mlp_norm out is sp sharded
+            mlp_up = s * b * i / cp / tp
+            mlp_gate = s * b * i / cp / tp
+            hadamard_swiglu = s * b * i / cp / tp
+            mlp_down = s * b * h / cp / tp
+            expert = mlp_up + mlp_gate + hadamard_swiglu + mlp_down
+            experts = expert * ae
+            mlp = mlp_norm + router + experts
+            layer = attn + mlp
+            layers = layer * l
+            # Other
+            embed = 0
+            final_layer_out = (
+                s * b * h / cp / tp
+            )  # both sequence and context parallelism
+            final_norm = s * b * h / cp / tp  # both sequence and context parallelism
+            unembed = s * b * v / cp / tp
+            logits = s * b * v / cp / tp
+            num_params = (
+                embed + layers + final_layer_out + final_norm + unembed + logits
+            )
+            return num_params
+
+    def calculate_activation_memory(self) -> float:
+        return (
+            self.calculate_activation_parameters() * 4
+        )  # assuming 4 bytes (32 bits) per activation
+
+    def calculate_gradient_memory(self) -> float:
+        # https://blog.eleuther.ai/transformer-math/#gradients
+        return (
+            self.calculate_parameter_memory()
+        )  # gradients are same size as parameters
+
+    def calculate_optimizer_memory(self) -> float:
+        # https://blog.eleuther.ai/transformer-math/#optimizer-states
+        # https://www.determined.ai/blog/act-mem-2, https://web.archive.org/web/20250308172134/https://www.determined.ai/blog/act-mem-2
+        return (
+            2 * self.calculate_parameter_memory()
+        )  # Adam optimizer with 3 states per parameter

defaults.py

@@ -1,16 +1,25 @@
 from state import ModelState
 
-GEMMA3_270M = ModelState(vocab_size=256000, num_layers=9, hidden_dim=1152, intermediate_size=4608)
-GEMMA3_1B = ModelState(vocab_size=262208, num_layers=26, hidden_dim=2304, intermediate_size=9216)
-GEMMA3_4B = ModelState(vocab_size=262208, num_layers=28, hidden_dim=3072, intermediate_size=12288)
-GEMMA3_12B = ModelState(vocab_size=262208, num_layers=42, hidden_dim=4608, intermediate_size=18432)
-GEMMA3_27B = ModelState(vocab_size=262208, num_layers=46, hidden_dim=6144, intermediate_size=24576)
+GEMMA3_270M = ModelState(
+    vocab_size=256000, num_layers=9, hidden_dim=1152, intermediate_size=4608
+)
+GEMMA3_1B = ModelState(
+    vocab_size=262208, num_layers=26, hidden_dim=2304, intermediate_size=9216
+)
+GEMMA3_4B = ModelState(
+    vocab_size=262208, num_layers=28, hidden_dim=3072, intermediate_size=12288
+)
+GEMMA3_12B = ModelState(
+    vocab_size=262208, num_layers=42, hidden_dim=4608, intermediate_size=18432
+)
+GEMMA3_27B = ModelState(
+    vocab_size=262208, num_layers=46, hidden_dim=6144, intermediate_size=24576
+)
 
 DEFAULTS = {
     "Gemma3 270M": GEMMA3_270M,
     "Gemma3 1B": GEMMA3_1B,
     "Gemma3 4B": GEMMA3_4B,
     "Gemma3 12B": GEMMA3_12B,
-    "Gemma3 27B": GEMMA3_27B
+    "Gemma3 27B": GEMMA3_27B,
 }
-

state.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 
+
 @dataclass
 class Model:
     vocab_size: int
@@ -16,7 +17,8 @@ class Parallelism:
     context_parallelism: int
     expert_parallelism: int
 
+
 @dataclass
 class Training:
     sequence_length: int
-    batch_size: int
+    batch_size: int