initial calculator

defaults.py

@@ -1,10 +1,10 @@
 from state import ModelState
 
-GEMMA3_270M = ModelState(vocab_size=256000, num_layers=9, hidden_size=1152, intermediate_size=4608)
-GEMMA3_1B = ModelState(vocab_size=262208, num_layers=26, hidden_size=2304, intermediate_size=9216)
-GEMMA3_4B = ModelState(vocab_size=262208, num_layers=28, hidden_size=3072, intermediate_size=12288)
-GEMMA3_12B = ModelState(vocab_size=262208, num_layers=42, hidden_size=4608, intermediate_size=18432)
-GEMMA3_27B = ModelState(vocab_size=262208, num_layers=46, hidden_size=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,

state.py

@@ -1,8 +1,22 @@
 from dataclasses import dataclass
 
 @dataclass
-class ModelState:
+class Model:
     vocab_size: int
     num_layers: int
-    hidden_size: int
-    intermediate_size: int
+    hidden_dim: int
+    intermediate_size: int
+    weight_tied_embeddings: bool
+
+
+@dataclass
+class Parallelism:
+    tensor_parallelism: int
+    pipeline_parallelism: int
+    context_parallelism: int
+    expert_parallelism: int
+
+@dataclass
+class Training:
+    sequence_length: int
+    batch_size: int