[WIP] differentially private training

configs/llama/munin-7b-core-pt-dp.toml

@@ -0,0 +1,49 @@
+model_name = "llama3"
+flavor = "Comma7B"
+tokenizer_name = "common-pile/comma-v0.1-2t"
+
+# job 
+job_name = "munin-7b-core-pt-dp"
+wandb_project = "munin-7b-core-pt-dp"
+enable_wandb = false
+
+# parallelism
+num_nodes = 1
+data_parallel_shard_degree = 8
+data_parallel_replicate_degree = 1
+
+# training settings
+train_batch_size = 4
+seq_len = 4096
+gradient_accumulation_steps = 2
+train_num_steps = 60097
+scheduler = "linear_warmup_constant_sqrt_decay"
+warmup_steps = 1000
+cooldown_steps = 1000
+checkpoint_interval = 500
+forced_load_path = "/work/training/maester/comma-v0.1-2t-dcp/"
+compile = true
+enable_cut_cross_entropy = false
+ac_mode = "none"
+selective_ac_option = "op"
+
+dp_enabled = true
+dp_clip_norm = 1.0
+dp_noise_multiplier = 1.0
+dp_num_privacy_units = 10000
+dp_delta = 1e-6
+
+[dataset]
+bos_token = 2
+eos_token = 1
+data_dirs = [
+    "/work/data/dfm-common-pile-16_3/",
+]
+dataset_weights = "1.0"
+
+[opt_cfg] # must specify *all* fields here, will not merge with defaults
+lr = 1e-5
+betas = [0.9, 0.95]
+weight_decay = 0.1
+eps = 1e-9
+fused = true

maester/config.py

@@ -157,6 +157,14 @@ class Config(BaseSettings):
     enable_async_tensor_parallel: bool = False
     enable_compiled_autograd: bool = True
 
+    # differential privacy
+    dp_enabled: bool = False
+    dp_clip_norm: float = 1.0
+    dp_noise_multiplier: float = 1.0
+    dp_num_privacy_units: int = 1
+    dp_delta: float = 1e-6
+    dp_assert: bool = False
+
     # profiling
     enable_profiling: bool = True
     enable_memory_snapshot: bool = False

maester/models/gemma/__init__.py

@@ -3,6 +3,25 @@
 __all__ = ["GemmaTextModel", "ModelArgs"]
 
 gemma3_configs = {
+    "270M": ModelArgs(
+        vocab_size=262_144,
+        dim=640,
+        n_layers=18,
+        n_heads=4,
+        num_key_value_heads=1,
+        head_dim=256,
+        intermediate_size=2048,
+        attn_types=["local_sliding", "local_sliding", "local_sliding", "local_sliding", "local_sliding", "global"],
+        use_post_ffw_norm=True,
+        use_pre_ffw_norm=True,
+        sliding_window_size=512,
+        rope_wave_length={
+            "local_sliding": 10_000,
+            "global": 1_000_000,
+        },
+        use_qk_norm=True,
+        vision_config=None,
+    ),
     "1B": ModelArgs(
         vocab_size=262_144,  # Actual size from google/gemma-3-1b-pt tokenizer
         dim=1152,