fsdp + bugfixes

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app.py

@@ -2,16 +2,19 @@ import gradio as gr
 import pandas as pd
 from functools import partial
 from defaults import DEFAULTS
+from details import DETAILS
 from state import Model, Parallelism, Training
 from calculator import MemoryCalculation
 from dtypes import DType
+from gradio.themes import ThemeClass as Theme
+from limitations import LIMITATIONS
 
 # Create a Number component for natural numbers (positive integers)
 NaturalNumber = partial(gr.Number, minimum=1, step=1, precision=0, interactive=True)
 
+colors = {
 
-def greet(name, intensity) -> str:
-    return "Hello, " + name + "!" * int(intensity)
+}
 
 
 def create_parallelism_block():
@@ -22,30 +25,103 @@ def create_parallelism_block():
             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
+
+            fsdp_enabled = gr.Checkbox(label="FSDP (Fully Sharded Data Parallel)", value=False)
+            fsdp_parallelism = NaturalNumber(label="FSDP Parallelism", value=1, interactive=False, elem_classes="disabled-field")
+            fsdp_strategy = gr.Radio(
+                choices=["Zero-1", "Zero-2", "Zero-3"],
+                label="FSDP Strategy",
+                value="Zero-1",
+                interactive=False,
+                elem_classes="disabled-field"
+            )
+
+            # Toggle FSDP fields interactivity based on FSDP checkbox
+            fsdp_enabled.change(
+                fn=lambda x: [
+                    gr.update(interactive=x, elem_classes=[] if x else ["disabled-field"]),
+                    gr.update(interactive=x, elem_classes=[] if x else ["disabled-field"])
+                ],
+                inputs=fsdp_enabled,
+                outputs=[fsdp_parallelism, fsdp_strategy]
+            )
+
+            return tp, pp, cp, ep, fsdp_enabled, fsdp_parallelism, fsdp_strategy
 
 
 def create_model_block():
     with gr.Column():
         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)
+        layers = NaturalNumber(label="Number of Layers", value=48)
+        vocab = NaturalNumber(label="Vocab Size", value=262144)
+        hidden = NaturalNumber(label="Hidden Dim", value=3840)
+        intermediate = NaturalNumber(label="Intermediate Dim", value=15360)
         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)
+        active_experts = NaturalNumber(label="Active Experts", value=2, interactive=False, elem_classes="disabled-field")
+        total_experts = NaturalNumber(label="Total Experts", value=8, interactive=False, elem_classes="disabled-field")
+        weight_tied_embeddings = gr.Checkbox(label="Weight Tied Embeddings", value=True)
 
-        # Toggle expert fields visibility based on MoE checkbox
+        # Toggle expert fields interactivity based on MoE checkbox
         is_moe.change(
-            fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
+            fn=lambda x: [
+                gr.update(interactive=x, elem_classes=[] if x else ["disabled-field"]),
+                gr.update(interactive=x, elem_classes=[] if x else ["disabled-field"])
+            ],
             inputs=is_moe,
             outputs=[active_experts, total_experts]
         )
 
-        # 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
+        presets = gr.Dropdown(["Custom"] + list(DEFAULTS.keys()), label="Presets", value="Gemma3 12B", interactive=True)
+
+        # Populate model parameters when preset is selected
+        def populate_from_preset(preset_name):
+            if preset_name and preset_name in DEFAULTS:
+                model = DEFAULTS[preset_name]
+                return [
+                    gr.update(value=model.num_layers),
+                    gr.update(value=model.vocab_size),
+                    gr.update(value=model.hidden_dim),
+                    gr.update(value=model.intermediate_size),
+                    gr.update(value=model.is_moe),
+                    gr.update(value=model.active_experts, interactive=model.is_moe),
+                    gr.update(value=model.total_experts, interactive=model.is_moe),
+                    gr.update(value=model.weight_tied_embeddings)
+                ]
+            return [gr.update() for _ in range(8)]
+
+        # Switch to "Custom" when user manually edits values
+        def switch_to_custom(layers_val, vocab_val, hidden_val, intermediate_val, is_moe_val, active_experts_val, total_experts_val, weight_tied_val, current_preset):
+            # Don't switch to custom if a preset is being applied
+            if current_preset and current_preset in DEFAULTS:
+                model = DEFAULTS[current_preset]
+                # Check if current values match the preset exactly
+                if (layers_val == model.num_layers and
+                    vocab_val == model.vocab_size and
+                    hidden_val == model.hidden_dim and
+                    intermediate_val == model.intermediate_size and
+                    is_moe_val == model.is_moe and
+                    active_experts_val == model.active_experts and
+                    total_experts_val == model.total_experts and
+                    weight_tied_val == model.weight_tied_embeddings):
+                    return gr.update()  # Keep current preset
+
+            return gr.update(value="Custom")
+
+        presets.change(
+            fn=populate_from_preset,
+            inputs=presets,
+            outputs=[layers, vocab, hidden, intermediate, is_moe, active_experts, total_experts, weight_tied_embeddings]
+        )
+
+        # Add change listeners to all model parameter inputs
+        for input_component in [layers, vocab, hidden, intermediate, is_moe, active_experts, total_experts, weight_tied_embeddings]:
+            input_component.change(
+                fn=switch_to_custom,
+                inputs=[layers, vocab, hidden, intermediate, is_moe, active_experts, total_experts, weight_tied_embeddings, presets],
+                outputs=presets
+            )
+
+        return layers, vocab, hidden, intermediate, active_experts, total_experts, is_moe, presets, weight_tied_embeddings
 
 
 def create_training_block():
@@ -58,12 +134,15 @@ def create_training_block():
             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)
+        param_dtype = gr.Dropdown(DType.values(), label="Parameter Dtype", value=DType.FP32.value, interactive=False, elem_classes="disabled-field")
+        reduce_dtype = gr.Dropdown(DType.values(), label="Reduce Dtype", value=DType.FP32.value, interactive=False, elem_classes="disabled-field")
 
-        # Toggle dtype fields visibility based on mixed precision checkbox
+        # Toggle dtype fields interactivity based on mixed precision checkbox
         mixed_precision.change(
-            fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
+            fn=lambda x: [
+                gr.update(interactive=x, elem_classes=[] if x else ["disabled-field"]),
+                gr.update(interactive=x, elem_classes=[] if x else ["disabled-field"])
+            ],
             inputs=mixed_precision,
             outputs=[param_dtype, reduce_dtype]
         )
@@ -71,14 +150,14 @@ def create_training_block():
         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):
+def calculate(tp, pp, cp, ep, fsdp_enabled, fsdp_parallelism, fsdp_strategy, layers, vocab, hidden, intermediate, active_experts, total_experts, is_moe, weight_tied_embeddings, 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
+        weight_tied_embeddings=weight_tied_embeddings,
         active_experts=int(active_experts),
         total_experts=int(total_experts),
         is_moe=is_moe
@@ -88,7 +167,10 @@ def calculate(tp, pp, cp, ep, layers, vocab, hidden, intermediate, active_expert
         tensor_parallelism=int(tp),
         pipeline_parallelism=int(pp),
         context_parallelism=int(cp),
-        expert_parallelism=int(ep)
+        expert_parallelism=int(ep),
+        fsdp_enabled=fsdp_enabled,
+        fsdp_parallelism=int(fsdp_parallelism),
+        fsdp_strategy=fsdp_strategy
     )
 
     training_config = Training(
@@ -111,48 +193,118 @@ def calculate(tp, pp, cp, ep, layers, vocab, hidden, intermediate, active_expert
     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
-        ]
-    })
+    # Calculate total memory
+    total_memory = param_memory + activation_memory + gradient_memory + optimizer_memory
+
+    # Round to 1 decimal place for display
+    param_gb = round(param_memory / 1e9, 1)
+    activation_gb = round(activation_memory / 1e9, 1)
+    gradient_gb = round(gradient_memory / 1e9, 1)
+    optimizer_gb = round(optimizer_memory / 1e9, 1)
+    total_gb = round(total_memory / 1e9, 1)
+
+    # Create DataFrame for stacked bar plot
+    # Start with stacked total bar, then add individual bars
+    individual_data = []
+
+    # Stacked total bar first - create separate rows for each component within total
+    for mem_type, gb_val in [
+        ('Activation', activation_gb),
+        ('Optimizer', optimizer_gb),
+        ('Gradient', gradient_gb),
+        ('Parameter', param_gb)
+    ]:
+        individual_data.append({
+            'Component': f'Total Memory\n{total_gb} GB',
+            'Memory (GB)': gb_val,
+            'Type': mem_type
+        })
+
+    # Individual component bars
+    for component, gb_val, mem_type in [
+        (f'Parameter Memory\n{param_gb} GB', param_gb, 'Parameter'),
+        (f'Gradient Memory\n{gradient_gb} GB', gradient_gb, 'Gradient'),
+        (f'Optimizer Memory\n{optimizer_gb} GB', optimizer_gb, 'Optimizer'),
+        (f'Activation Memory\n{activation_gb} GB', activation_gb, 'Activation')
+    ]:
+        individual_data.append({
+            'Component': component,
+            'Memory (GB)': gb_val,
+            'Type': mem_type
+        })
+
+    memory_data = pd.DataFrame(individual_data)
+
+    # Define pastel color map
+    color_map = {
+        'Parameter': '#B6E5D8',   # Light Mint
+        'Gradient': '#FFB6C1',    # Light Pink
+        'Optimizer': '#C7B3FF',   # Light Purple
+        'Activation': '#FFD1A9',  # Light Peach
+    }
 
     return gr.BarPlot(
         value=memory_data,
         x="Component",
         y="Memory (GB)",
+        color="Type",
+        color_map=color_map,
         title="LLM Memory Usage Breakdown",
         container=False,
-        y_lim=[0, None]
+        y_lim=[0, None],
+        sort=[
+            f'Total Memory\n{total_gb} GB',
+            f'Parameter Memory\n{param_gb} GB',
+            f'Gradient Memory\n{gradient_gb} GB',
+            f'Optimizer Memory\n{optimizer_gb} GB',
+            f'Activation Memory\n{activation_gb} GB'
+        ]
     )
 
+css = """
+/* Style for disabled components to make them visually obvious */
+.disabled-field input,
+.disabled-field select,
+.disabled-field textarea {
+    opacity: 0.4 !important;
+    background-color: #f5f5f5 !important;
+    color: #999 !important;
+    cursor: not-allowed !important;
+    text-decoration: line-through;
+}
+
+.disabled-field label {
+    opacity: 0.5 !important;
+    color: #999 !important;
+}
+"""
 
-with gr.Blocks(theme='gstaff/xkcd') as demo:
-    with gr.Sidebar():
-        gr.Textbox("## LLM Memory Visualizer")
+theme = Theme.from_hub("gstaff/xkcd")
+# otherwise invisible in light mode
+theme.checkbox_label_text_color=theme.block_label_text_color
+with gr.Blocks(theme=theme, css=css) as demo:
     with gr.Column():
         with gr.Row(equal_height=True):
-            tp, pp, cp, ep = create_parallelism_block()
-            layers, vocab, hidden, intermediate, active_experts, total_experts, is_moe, presets = create_model_block()
+            tp, pp, cp, ep, fsdp_enabled, fsdp_parallelism, fsdp_strategy = create_parallelism_block()
+            layers, vocab, hidden, intermediate, active_experts, total_experts, is_moe, presets, weight_tied_embeddings = 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.BarPlot(label="Memory Usage Breakdown")
 
         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],
+            inputs=[tp, pp, cp, ep, fsdp_enabled, fsdp_parallelism, fsdp_strategy, layers, vocab, hidden, intermediate, active_experts, total_experts, is_moe, weight_tied_embeddings, seq_len, batch_size, gradient_checkpointing, grad_accumulation, precision, mixed_precision, param_dtype, reduce_dtype],
             outputs=output
         )
 
+        # Limitations and Comments section
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("# Limitations")
+                gr.Markdown(LIMITATIONS)
+            with gr.Column():
+                gr.Markdown("# Comments and Details")
+                gr.Markdown(DETAILS)
+
 
 demo.launch()

calculator.py

@@ -1,21 +1,27 @@
 from state import Model as Model, Parallelism, Training
 from dtypes import DType
+from math import ceil
 
 
 class MemoryCalculation:
-    def __init__(self, modelconfig: Model, parallelismconfig: Parallelism, trainingconfig: 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:
+    def calculate_num_parameters_per_layer(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.
 
-        #self tax
+        # self tax
         b, s = self.training.batch_size, self.training.sequence_length
         h, i, l, v, e = (
             self.model.hidden_dim,
@@ -30,45 +36,77 @@ class MemoryCalculation:
             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
+        layer_norm_attn_in = 2 * h  # not tp sharded
         qkv = 3 * h * h / tp
-        attn_output_proj = h * 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_norm_mlp_in = 2 * h  # not tp sharded
+        mlp_up_proj = (h * i + i) / tp
+        mlp_gate_proj = (h * i + i) / tp
+        mlp_down_proj = (i * h + h) / tp
+        mlp = layer_norm_mlp_in + mlp_up_proj + mlp_gate_proj + mlp_down_proj
+        if self.model.is_moe:
+            router = h * e + e  # assuming replicated for simplicity
+            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
+        return layer
 
-
-        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 small
+    def calculate_unshardeable_parameters(self) -> float:
+        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.total_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
+        final_layer_norm = 2 * h  # not tp sharded
+        # hush linter
         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
+            total_params = input_embedding + unembedding + final_layer_norm
+        elif pp > 1:
+            total_params = max(input_embedding, unembedding) + final_layer_norm
         return total_params
 
+    def calculate_fsdp_sharded_parameters(self) -> float:
+        if not self.parallelism.fsdp_enabled:
+            return self.calculate_num_parameters()
+        else:
+            return (
+                self.calculate_num_parameters_per_layer()
+                * ceil(
+                    (self.model.num_layers - 1) / self.parallelism.pipeline_parallelism
+                )
+                / self.parallelism.fsdp_parallelism
+                + self.calculate_unshardeable_parameters()
+                + self.calculate_num_parameters_per_layer()
+            )
+
+    def calculate_num_parameters(self) -> float:
+        return (
+            self.calculate_num_parameters_per_layer()
+            * ceil(self.model.num_layers / self.parallelism.pipeline_parallelism)
+            + self.calculate_unshardeable_parameters()
+        )
+
     def calculate_activation_parameters(self) -> float:
         # https://blog.eleuther.ai/transformer-math/#activations-and-batch-size
         # https://arxiv.org/abs/2205.05198
@@ -95,29 +133,24 @@ class MemoryCalculation:
             layer = s * b * h / cp / tp  # only keep initial input to layer
             layers = layer * l
             embed = 0
-            final_layer_out = (
-                s * b * h / cp / sp
-            )
-            final_norm = s * b * h / cp / sp  
+            final_layer_out = s * b * h / cp / sp
+            final_norm = s * b * h / cp / sp
             unembed = s * b * v / cp / tp
-            logits = s * b * v / cp / sp # come back to this
-            num_params = (
-                embed + layers + final_layer_out + final_norm + unembed + logits
-            )
+            num_params = embed + layers + final_layer_out + final_norm + unembed
             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 
-            attn_norm = s * b * h / cp / sp 
+            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 / sp 
+            mlp_norm = s * b * h / cp / sp
 
             mlp_up = s * b * i / cp / tp
             mlp_gate = s * b * i / cp / tp
@@ -125,53 +158,77 @@ class MemoryCalculation:
             mlp_down = s * b * h / cp / tp
             if self.model.is_moe:
                 router = (
-                s * b * e / cp / sp)  # makes sense to sp shard if mlp_norm out is sp sharded
+                    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
+                experts = expert * ae / ep
                 mlp = mlp_norm + router + experts
             else:
                 mlp = mlp_norm + mlp_up + mlp_gate + hadamard_swiglu + mlp_down
             layer = attn + mlp
-            layers = layer * l # no decrease from PP because schedules will increase microbatches
+            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 / sp  
+            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
-            )
+            num_params = embed + layers + final_layer_out + final_norm + unembed
             return num_params
-        
+
     def calculate_parameter_memory(self) -> float:
+        if self.parallelism.fsdp_enabled and self.parallelism.fsdp_strategy == "Zero-3":
+            params = self.calculate_fsdp_sharded_parameters()
+        else:
+            params = self.calculate_num_parameters()
         if self.training.mixed_precision:
-            master_copy = self.calculate_num_parameters() * self.training.precision
-            working_copy = self.calculate_num_parameters() * self.training.param_dtype
+            master_copy = params * self.training.precision
+            working_copy = params * self.training.param_dtype
             return master_copy + working_copy
         else:
-            return self.calculate_num_parameters() * self.training.precision
-    
+            return params * self.training.precision
+
     def calculate_gradient_memory(self) -> float:
         # https://blog.eleuther.ai/transformer-math/#gradients
-        return (
-            self.calculate_num_parameters() * 4
-        )  # gradients are same size as parameters
+        if self.parallelism.fsdp_enabled and self.parallelism.fsdp_strategy in ("Zero-3", "Zero-2"):
+            params = self.calculate_fsdp_sharded_parameters()
+        else:
+            params = self.calculate_num_parameters()
+        grad_accumulation = 0
+        if self.training.grad_accumulation:
+            if self.training.mixed_precision:
+                grad_accumulation = (
+                    params * self.training.reduce_dtype
+                )
+            else:
+                grad_accumulation = (
+                    params * self.training.precision
+                )
+        if self.training.mixed_precision:
+            gradients = params * self.training.param_dtype
+        else:
+            gradients = params * self.training.precision
+        return grad_accumulation + gradients
 
     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_num_parameters() * DType.FP32
-        )  # Adam optimizer with 2 states per parameter, assume always fp32
-
+        if self.parallelism.fsdp_enabled:
+            return (
+                2 * self.calculate_num_parameters() * DType.FP32
+            )  / self.parallelism.fsdp_parallelism # don't gather a layer unlike params and grads
+        else:
+            return (
+                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
-            )
-
+            )  # not impacted by fsdp

defaults.py

@@ -1,24 +1,116 @@
 from state import Model
 
+# https://huggingface.co/google/gemma-3-270m/blob/main/config.json
 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
+    vocab_size=262144,
+    num_layers=18,
+    hidden_dim=640,
+    intermediate_size=2048,
+    weight_tied_embeddings=True,
+    active_experts=1,
+    total_experts=1,
+    is_moe=False,
 )
 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
+    vocab_size=262144,
+    num_layers=26,
+    hidden_dim=1152,
+    intermediate_size=6912,
+    weight_tied_embeddings=True,
+    active_experts=1,
+    total_experts=1,
+    is_moe=False,
 )
 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
+    vocab_size=262144,
+    num_layers=34,
+    hidden_dim=2560,
+    intermediate_size=10240,
+    weight_tied_embeddings=True,
+    active_experts=1,
+    total_experts=1,
+    is_moe=False,
 )
 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
+    vocab_size=262144,
+    num_layers=48,
+    hidden_dim=3840,
+    intermediate_size=15360,
+    weight_tied_embeddings=True,
+    active_experts=1,
+    total_experts=1,
+    is_moe=False,
 )
 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
+    vocab_size=262144,
+    num_layers=62,
+    hidden_dim=5376,
+    intermediate_size=21504,
+    weight_tied_embeddings=True,
+    active_experts=1,
+    total_experts=1,
+    is_moe=False,
+)
+# No maverick, don't support non-homogenous layers yet
+
+# https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct/blob/main/config.json
+LLAMA4_SCOUT = Model(
+    vocab_size=202048,
+    num_layers=48,
+    hidden_dim=5120,
+    intermediate_size=8192,
+    weight_tied_embeddings=True,
+    active_experts=2,
+    total_experts=17,
+    is_moe=True,
+)
+
+# https://huggingface.co/unsloth/Llama-3.2-1B-Instruct/blob/main/config.json
+LLAMA3_1B = Model(
+    vocab_size=128256,
+    num_layers=16,
+    hidden_dim=2048,
+    intermediate_size=8192,
+    weight_tied_embeddings=True,
+    active_experts=1,
+    total_experts=1,
+    is_moe=False,
+)
+
+# https://huggingface.co/unsloth/Llama-3.2-3B-Instruct/blob/main/config.json
+LLAMA3_3B = Model(
+    vocab_size=128256,
+    num_layers=28,
+    hidden_dim=3072,
+    intermediate_size=8192,
+    weight_tied_embeddings=True,
+    active_experts=1,
+    total_experts=1,
+    is_moe=False,
+)
+
+# https://huggingface.co/unsloth/llama-3-8b-Instruct/blob/main/config.json
+LLAMA3_8B = Model(
+    vocab_size=128256,
+    num_layers=32,
+    hidden_dim=4096,
+    intermediate_size=14336,
+    weight_tied_embeddings=True,
+    active_experts=1,
+    total_experts=1,
+    is_moe=False,
+)
+
+# https://huggingface.co/unsloth/Llama-3.3-70B-Instruct/blob/main/config.json
+LLAMA3_70B = Model(
+    vocab_size=128256,
+    num_layers=80,
+    hidden_dim=8192,
+    intermediate_size=28672,
+    weight_tied_embeddings=True,
+    active_experts=1,
+    total_experts=1,
+    is_moe=False,
 )
 
 DEFAULTS = {
@@ -27,4 +119,9 @@ DEFAULTS = {
     "Gemma3 4B": GEMMA3_4B,
     "Gemma3 12B": GEMMA3_12B,
     "Gemma3 27B": GEMMA3_27B,
+    "Llama3 1B": LLAMA3_1B,
+    "Llama3 3B": LLAMA3_3B,
+    "Llama3 8B": LLAMA3_8B,
+    "Llama3 70B": LLAMA3_70B,
+    "Llama4 Scout": LLAMA4_SCOUT,
 }

details.py

@@ -0,0 +1,21 @@
+DETAILS = """
+### Resources I found helpful while building this tool:
+- [The Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook)
+- [Reducing Activation Recomputation in Large Transformer Models](https://arxiv.org/abs/2205.05198)
+- [Transformer Math - Michael Wornow](https://michaelwornow.net/2024/01/18/counting-params-in-transformer)
+- [Transformer Math 101](https://blog.eleuther.ai/transformer-math/)
+
+
+### Why this tool?
+While there are some good tools out there already:
+- [Hugging Face Model Memory Estimator](https://huggingface.co/spaces/hf-accelerate/model-memory-usage)
+- [DeepSpeed Model Memory Calculator](https://huggingface.co/spaces/andstor/deepspeed-model-memory-usage)
+- [DeepSpeed Native Utility](https://deepspeed.readthedocs.io/en/latest/memory.html)
+
+None of them had all the features I wanted in one place. I wanted a tool that could:
+- Accept arbitrary model configurations
+- Support FSDP
+- Support 5d parallelism
+- Be interactive and break down memory usage by category, to better inform configurations.
+
+"""

dtypes.py

@@ -31,4 +31,4 @@ class DType(Enum):
 
     def __rmul__(self, other):
         """Multiply number by dtype to get total bytes"""
-        return other * self.bytes_per_element()
+        return other * self.bytes_per_element()

limitations.py

@@ -0,0 +1,27 @@
+LIMITATIONS = """
+This calculator has many limitations and assumptions
+### Assumptions:
+- Your implementation of tensor parallel also incorporates sequence parallel
+- You are doing selective recomputation with flash attention if not doing gradient checkpointing
+- You keep a master copy of the model weights for mixed precision
+    - May not be true for some implementations which cast on the fly
+- You're using Adam optimizer
+- If using PP you're using a schedule that will keep the number of activations roughly the same
+- EP is the number of PPxTP units that share each expert
+- Swiglu activation function
+- Rotary embeddings
+
+### Limitations:
+- Does not support non-homogenous layers
+    - e.g. Llama4 Maverick with alternating dense and sparse layers, iRoPE
+- Does not include memory for kernel or framework overhead
+- Does not include memory for intermediates 
+- Does not include vision layers for multi-modal models
+- Models shared experts as another routed expert per token
+- Does not support different dtypes for different parts of the model
+    - e.g. MXFP4 for GPT-OSS 20 and 120B
+- Have not validated EP/FSDP interaction
+- Doesn't model biases on a per-model basis
+
+Note this is not an exhaustive list, just some of the main ones
+"""

state.py

@@ -20,6 +20,9 @@ class Parallelism:
     pipeline_parallelism: int
     context_parallelism: int
     expert_parallelism: int
+    fsdp_enabled: bool
+    fsdp_parallelism: int
+    fsdp_strategy: str
 
 
 @dataclass