part2/adam_mini: final checkpoint

part2/adam_mini/README.md

@@ -0,0 +1,14 @@
+---
+language: [en]
+tags: [causal-lm, optimizers, anlp-assignment-2]
+---
+# ANLP A2 Part 2 — Adam-mini
+
+Dense decoder-only Transformer (27,269,632 params) trained from scratch on
+`browndw/human-ai-parallel-corpus` (causal LM, 1x dataset) with **Adam-mini**
+(Memory-efficient) implemented from scratch.
+
+* tokens: 41,648,128
+* final validation loss: 3.9736
+* final test BLEU (64-token continuation): 1.43
+* hyper-parameters: `{'lr': 0.0005, 'betas': [0.9, 0.95], 'eps': 1e-08, 'weight_decay': 0.1}`

part2/adam_mini/config.json

@@ -0,0 +1,100 @@
+{
+  "model": {
+    "vocab_size": 16384,
+    "D_model": 512,
+    "N_ctx": 512,
+    "n_layers": 6,
+    "n_heads": 8,
+    "d_ff": 2048,
+    "dropout": 0.0,
+    "rope_theta": 10000.0,
+    "tie_embeddings": true,
+    "norm_eps": 1e-05,
+    "use_moe": false,
+    "n_experts": 4,
+    "n_shared_experts": 0,
+    "top_k": 1,
+    "expert_d_ff": null,
+    "moe_param_match": "total",
+    "normalize_topk": true,
+    "aux_loss_coef": 0.01,
+    "z_loss_coef": 0.001
+  },
+  "optimizer": "adam_mini",
+  "info": {
+    "optimizer": "adam_mini",
+    "optimizer_name": "Adam-mini",
+    "category": "Memory-efficient",
+    "optimizer_hparams": {
+      "lr": 0.0005,
+      "betas": [
+        0.9,
+        0.95
+      ],
+      "eps": 1e-08,
+      "weight_decay": 0.1
+    },
+    "train_hparams": {
+      "dataset_passes": 1.0,
+      "data_fraction": 1.0,
+      "batch_size": 32,
+      "grad_accum_steps": 1,
+      "grad_clip": 1.0,
+      "schedule": "cosine",
+      "warmup_frac": 0.02,
+      "min_lr_ratio": 0.1,
+      "eval_interval": 0.1,
+      "val_max_tokens": 1000000,
+      "log_every": 20,
+      "save_every_steps": 500,
+      "compile": false
+    },
+    "model": {
+      "vocab_size": 16384,
+      "D_model": 512,
+      "N_ctx": 512,
+      "n_layers": 6,
+      "n_heads": 8,
+      "d_ff": 2048,
+      "dropout": 0.0,
+      "rope_theta": 10000.0,
+      "tie_embeddings": true,
+      "norm_eps": 1e-05,
+      "use_moe": false,
+      "n_experts": 4,
+      "n_shared_experts": 0,
+      "top_k": 1,
+      "expert_d_ff": null,
+      "moe_param_match": "total",
+      "normalize_topk": true,
+      "aux_loss_coef": 0.01,
+      "z_loss_coef": 0.001
+    },
+    "params": {
+      "total": 27269632,
+      "embedding": 8388608,
+      "non_embedding": 18881024,
+      "ffn_total": 12582912,
+      "ffn_active": 12582912,
+      "active": 27269632,
+      "active_non_embedding": 18881024
+    },
+    "tokens_seen": 41648128,
+    "dataset_tokens": 41663488,
+    "data_fraction": 1.0,
+    "steps": 2542,
+    "wall_clock_sec": 523.851163001993,
+    "tokens_per_sec": 79503.7425541452,
+    "peak_memory_mb": 6705.42626953125,
+    "optimizer_state_floats": 35679853,
+    "final_train_loss": 3.8838891983032227,
+    "test_loss": 3.9793649074968354,
+    "test_ppl": 53.483056750824076,
+    "diverged": false,
+    "device": "cuda",
+    "precision": "fp16",
+    "wandb_url": "https://wandb.ai/sanyamagrawal2005-iiit-hyderabad/anlp-assignment-2/runs/544bb65e",
+    "hf_url": null,
+    "tag": null
+  }
+}

part2/adam_mini/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e3513b00cf2a98fa339b9af4796b79c1c9a5a9e46f77a3b467c930ec348223d
+size 109083288