added support for precision

.gitignore

@@ -0,0 +1,26 @@
+# Python
+__pycache__/
+*.pyc
+*.pyo
+
+# Virtual environments
+.env
+venv/
+.venv/
+
+# Linting/Formatting
+.ruff_cache/
+
+# macOS
+.DS_Store
+
+# IDE
+.vscode/
+.idea/
+
+# Claude Code
+claude/
+
+# Gradio
+gradio_cached_examples/
+flagged/

app.py

@@ -1,5 +1,13 @@
 import gradio as gr
+import pandas as pd
+from functools import partial
 from defaults import DEFAULTS
+from state import Model, Parallelism, Training
+from calculator import MemoryCalculation
+from dtypes import DType
+
+# Create a Number component for natural numbers (positive integers)
+NaturalNumber = partial(gr.Number, minimum=1, step=1, precision=0, interactive=True)
 
 
 def greet(name, intensity) -> str:
@@ -8,52 +16,143 @@ def greet(name, intensity) -> str:
 
 def create_parallelism_block():
     with gr.Column():
-        gr.Markdown("# Parallelism Parameters")
-        tp = gr.Number(label="Tensor Parallelism", value=1, interactive=True)
-        pp = gr.Number(label="Pipeline Parallelism", value=1, interactive=True)
-        cp = gr.Number(label="Context Parallelism", value=1, interactive=True)
-        ep = gr.Number(label="Expert Parallelism", value=1, interactive=True)
-        return tp, pp, cp, ep
+        gr.Markdown("# Parallelism")
+        with gr.Group():
+            tp = NaturalNumber(label="Tensor Parallelism", value=1)
+            pp = NaturalNumber(label="Pipeline Parallelism", value=1)
+            cp = NaturalNumber(label="Context Parallelism", value=1)
+            ep = NaturalNumber(label="Expert Parallelism", value=1)
+            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
+        gr.Markdown("# Model Architecture")
+        layers = NaturalNumber(label="Number of Layers", value=32)
+        vocab = NaturalNumber(label="Vocab Size", value=32000)
+        hidden = NaturalNumber(label="Hidden Dim", value=4096)
+        intermediate = NaturalNumber(label="Intermediate Dim", value=11008)
+        is_moe = gr.Checkbox(label="Mixture of Experts (MoE)", value=False)
+        active_experts = NaturalNumber(label="Active Experts", value=2, visible=False)
+        total_experts = NaturalNumber(label="Total Experts", value=8, visible=False)
+
+        # Toggle expert fields visibility based on MoE checkbox
+        is_moe.change(
+            fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
+            inputs=is_moe,
+            outputs=[active_experts, total_experts]
         )
-        presets = gr.Dropdown(list(DEFAULTS.keys()), label="Presets", interactive=True)
-        return layers, vocab, hidden, intermediate, presets
+
+        # not ready yet
+        # presets = gr.Dropdown(list(DEFAULTS.keys()), label="Presets", interactive=True)
+        return layers, vocab, hidden, intermediate, active_experts, total_experts, is_moe, presets
 
 
 def create_training_block():
     with gr.Column():
-        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
+        gr.Markdown("# Training Config")
+        seq_len = NaturalNumber(label="Sequence Length", value=8192)
+        batch_size = NaturalNumber(label="Batch Size", info="If you are using gradient accumulation, enter microbatch size", value=8)
+        with gr.Row():
+            gradient_checkpointing = gr.Checkbox(label="Gradient Checkpointing", value=False)
+            grad_accumulation = gr.Checkbox(label="Gradient Accumulation", value=False)
+        precision = gr.Dropdown(DType.values(), label="Precision", value=DType.FP32.value, interactive=True)
+        mixed_precision = gr.Checkbox(label="Mixed Precision", value=False)
+        param_dtype = gr.Dropdown(DType.values(), label="Parameter Dtype", value=DType.FP32.value, interactive=True, visible=False)
+        reduce_dtype = gr.Dropdown(DType.values(), label="Reduce Dtype", value=DType.FP32.value, interactive=True, visible=False)
+
+        # Toggle dtype fields visibility based on mixed precision checkbox
+        mixed_precision.change(
+            fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
+            inputs=mixed_precision,
+            outputs=[param_dtype, reduce_dtype]
+        )
+
+        return seq_len, batch_size, gradient_checkpointing, grad_accumulation, precision, mixed_precision, param_dtype, reduce_dtype
+
 
+def calculate(tp, pp, cp, ep, layers, vocab, hidden, intermediate, active_experts, total_experts, is_moe, seq_len, batch_size, gradient_checkpointing, grad_accumulation, precision, mixed_precision, param_dtype, reduce_dtype):
+    # Create state objects
+    model_config = Model(
+        vocab_size=int(vocab),
+        num_layers=int(layers),
+        hidden_dim=int(hidden),
+        intermediate_size=int(intermediate),
+        weight_tied_embeddings=True,  # Default assumption
+        active_experts=int(active_experts),
+        total_experts=int(total_experts),
+        is_moe=is_moe
+    )
 
-def calculate(*args) -> int:
-    out = 1
-    for arg in args:
-        out *= arg
-    return arg
+    parallelism_config = Parallelism(
+        tensor_parallelism=int(tp),
+        pipeline_parallelism=int(pp),
+        context_parallelism=int(cp),
+        expert_parallelism=int(ep)
+    )
 
+    training_config = Training(
+        sequence_length=int(seq_len),
+        batch_size=int(batch_size),
+        gradient_checkpointing=gradient_checkpointing,
+        grad_accumulation=grad_accumulation,
+        precision=DType(precision),
+        mixed_precision=mixed_precision,
+        param_dtype=DType(param_dtype),
+        reduce_dtype=DType(reduce_dtype)
+    )
 
-with gr.Blocks() as demo:
+    # Calculate different memory components
+    calc = MemoryCalculation(model_config, parallelism_config, training_config)
+
+    # Get all memory calculations
+    param_memory = calc.calculate_parameter_memory()
+    activation_memory = calc.calculate_activation_memory()
+    gradient_memory = calc.calculate_gradient_memory()
+    optimizer_memory = calc.calculate_optimizer_memory()
+
+    # Create DataFrame for bar plot
+    memory_data = pd.DataFrame({
+        'Component': [
+            'Parameter Memory',
+            'Activation Memory',
+            'Gradient Memory',
+            'Optimizer Memory'
+        ],
+        'Memory (GB)': [
+            param_memory / 1e9,
+            activation_memory / 1e9,
+            gradient_memory / 1e9,
+            optimizer_memory / 1e9
+        ]
+    })
+
+    return gr.BarPlot(
+        value=memory_data,
+        x="Component",
+        y="Memory (GB)",
+        title="LLM Memory Usage Breakdown",
+        container=False,
+        y_lim=[0, None]
+    )
+
+
+with gr.Blocks(theme='gstaff/xkcd') as demo:
+    with gr.Sidebar():
+        gr.Textbox("## LLM Memory Visualizer")
     with gr.Column():
-        with gr.Row():
+        with gr.Row(equal_height=True):
             tp, pp, cp, ep = create_parallelism_block()
-            layers, vocab, hidden, intermediate, presets = create_model_block()
-            seq_len, batch_size = create_training_block()
+            layers, vocab, hidden, intermediate, active_experts, total_experts, is_moe, presets = create_model_block()
+            seq_len, batch_size, gradient_checkpointing, grad_accumulation, precision, mixed_precision, param_dtype, reduce_dtype = create_training_block()
         calculate_button = gr.Button("Calculate")
-        output = gr.Number(label="Output")
+        output = gr.BarPlot(label="Memory Usage Breakdown")
 
-        calculate_button.click(fn=calculate, inputs=[tp, pp, cp, ep], outputs=output)
+        calculate_button.click(
+            fn=calculate,
+            inputs=[tp, pp, cp, ep, layers, vocab, hidden, intermediate, active_experts, total_experts, is_moe, seq_len, batch_size, gradient_checkpointing, grad_accumulation, precision, mixed_precision, param_dtype, reduce_dtype],
+            outputs=output
+        )
 
 
 demo.launch()

calculator.py

@@ -1,11 +1,12 @@
 from state import Model as Model, Parallelism, Training
+from dtypes import DType
 
 
 class MemoryCalculation:
-    def __init__(self, model: Model, parallelism: Parallelism, training: Training):
-        self.model = model
-        self.parallelism = parallelism
-        self.training = training
+    def __init__(self, modelconfig: Model, parallelismconfig: Parallelism, trainingconfig: Training):
+        self.model = modelconfig
+        self.parallelism = parallelismconfig
+        self.training = trainingconfig
 
     def calculate_num_parameters(self) -> float:
         # https://michaelwornow.net/2024/01/18/counting-params-in-transformer
@@ -14,13 +15,14 @@ class MemoryCalculation:
         # 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.
 
+        #self tax
         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,
+            self.model.total_experts,
         )
         tp, pp, ep = (
             self.parallelism.tensor_parallelism,
@@ -59,18 +61,14 @@ class MemoryCalculation:
 
         # 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
+        # even though that's pretty small
+        total_params = 0
         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
@@ -81,6 +79,7 @@ class MemoryCalculation:
             self.model.intermediate_size,
             self.model.num_layers,
             self.model.vocab_size,
+            self.model.total_experts,
             self.model.active_experts,
         )
         tp, cp, pp, ep = (
@@ -89,6 +88,7 @@ class MemoryCalculation:
             self.parallelism.pipeline_parallelism,
             self.parallelism.expert_parallelism,
         )
+        sp = tp
         if self.training.gradient_checkpointing:
             # full recomputation
             embed = 0
@@ -96,11 +96,11 @@ class MemoryCalculation:
             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
+                s * b * h / cp / sp
+            )
+            final_norm = s * b * h / cp / sp  
             unembed = s * b * v / cp / tp
-            logits = s * b * v / cp / tp  # both vocab and tensor parallelism
+            logits = s * b * v / cp / sp # come back to this
             num_params = (
                 embed + layers + final_layer_out + final_norm + unembed + logits
             )
@@ -110,53 +110,68 @@ class MemoryCalculation:
             # 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
+            layer_in = s * b * h / cp / tp 
+            attn_norm = s * b * h / cp / sp 
             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_norm = s * b * h / cp / sp 
+
             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
+            if self.model.is_moe:
+                router = (
+                s * b * e / cp / sp)  # makes sense to sp shard if mlp_norm out is sp sharded
+                expert = mlp_up + mlp_gate + hadamard_swiglu + mlp_down
+                experts = expert * ae
+                mlp = mlp_norm + router + experts
+            else:
+                mlp = mlp_norm + mlp_up + mlp_gate + hadamard_swiglu + mlp_down
             layer = attn + mlp
-            layers = layer * l
+            layers = layer * l # no decrease from PP because schedules will increase microbatches
             # 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
+            final_norm = s * b * h / cp / sp  
             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_parameter_memory(self) -> float:
+        if self.training.mixed_precision:
+            master_copy = self.calculate_num_parameters() * self.training.precision
+            working_copy = self.calculate_num_parameters() * self.training.param_dtype
+            return master_copy + working_copy
+        else:
+            return self.calculate_num_parameters() * self.training.precision
+    
     def calculate_gradient_memory(self) -> float:
         # https://blog.eleuther.ai/transformer-math/#gradients
         return (
-            self.calculate_parameter_memory()
+            self.calculate_num_parameters() * 4
         )  # 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
+            2 * self.calculate_num_parameters() * DType.FP32
+        )  # Adam optimizer with 2 states per parameter, assume always fp32
+
+    def calculate_activation_memory(self) -> float:
+        if self.training.mixed_precision:
+            return self.calculate_activation_parameters() * self.training.param_dtype
+        else:
+            return (
+                self.calculate_activation_parameters() * self.training.precision
+            )
+

defaults.py

@@ -1,19 +1,24 @@
-from state import ModelState
+from state import Model
 
-GEMMA3_270M = ModelState(
-    vocab_size=256000, num_layers=9, hidden_dim=1152, intermediate_size=4608
+GEMMA3_270M = Model(
+    vocab_size=256000, num_layers=9, hidden_dim=1152, intermediate_size=4608,
+    weight_tied_embeddings=True, active_experts=2, total_experts=8, is_moe=False
 )
-GEMMA3_1B = ModelState(
-    vocab_size=262208, num_layers=26, hidden_dim=2304, intermediate_size=9216
+GEMMA3_1B = Model(
+    vocab_size=262208, num_layers=26, hidden_dim=2304, intermediate_size=9216,
+    weight_tied_embeddings=True, active_experts=2, total_experts=8, is_moe=False
 )
-GEMMA3_4B = ModelState(
-    vocab_size=262208, num_layers=28, hidden_dim=3072, intermediate_size=12288
+GEMMA3_4B = Model(
+    vocab_size=262208, num_layers=28, hidden_dim=3072, intermediate_size=12288,
+    weight_tied_embeddings=True, active_experts=2, total_experts=8, is_moe=False
 )
-GEMMA3_12B = ModelState(
-    vocab_size=262208, num_layers=42, hidden_dim=4608, intermediate_size=18432
+GEMMA3_12B = Model(
+    vocab_size=262208, num_layers=42, hidden_dim=4608, intermediate_size=18432,
+    weight_tied_embeddings=True, active_experts=2, total_experts=8, is_moe=False
 )
-GEMMA3_27B = ModelState(
-    vocab_size=262208, num_layers=46, hidden_dim=6144, intermediate_size=24576
+GEMMA3_27B = Model(
+    vocab_size=262208, num_layers=46, hidden_dim=6144, intermediate_size=24576,
+    weight_tied_embeddings=True, active_experts=2, total_experts=8, is_moe=False
 )
 
 DEFAULTS = {

dtypes.py

@@ -0,0 +1,34 @@
+from enum import Enum
+
+
+class DType(Enum):
+    FP32 = "fp32"
+    FP16 = "fp16"
+    BF16 = "bf16"
+    FP8 = "fp8"
+
+    @classmethod
+    def values(cls):
+        """Return a list of all enum values"""
+        return [dtype.value for dtype in cls]
+
+    def bytes_per_element(self):
+        """Return the number of bytes per element for this dtype"""
+        if self == DType.FP32:
+            return 4
+        elif self == DType.FP16:
+            return 2
+        elif self == DType.BF16:
+            return 2
+        elif self == DType.FP8:
+            return 1
+        else:
+            raise ValueError(f"Unknown dtype: {self}")
+
+    def __mul__(self, other):
+        """Multiply dtype by a number to get total bytes"""
+        return self.bytes_per_element() * other
+
+    def __rmul__(self, other):
+        """Multiply number by dtype to get total bytes"""
+        return other * self.bytes_per_element()

state.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from dtypes import DType
 
 
 @dataclass
@@ -8,6 +9,9 @@ class Model:
     hidden_dim: int
     intermediate_size: int
     weight_tied_embeddings: bool
+    active_experts: int
+    total_experts: int
+    is_moe: bool
 
 
 @dataclass
@@ -22,3 +26,9 @@ class Parallelism:
 class Training:
     sequence_length: int
     batch_size: int
+    gradient_checkpointing: bool
+    grad_accumulation: bool
+    precision: DType
+    mixed_precision: bool
+    param_dtype: DType
+    reduce_dtype: DType