Add completed TinyStories eostok-apr artifact

ts-fivearm/eostok-apr/resolved_config.json

@@ -0,0 +1,225 @@
+{
+  "command": [
+    "/home/ishank/mstok/src/mstok/text/training/train_joint.py",
+    "--config",
+    "configs/text/tinystories/arms/eostok_apr.yaml"
+  ],
+  "config": {
+    "codec": {
+      "checkpoint_path": null,
+      "cls_hash_bits": 0,
+      "cls_hash_weight": 0.1,
+      "compression_factor": 1,
+      "context_length": 64,
+      "dec_layers": 1,
+      "dropout": 0.1,
+      "embed_dim": 384,
+      "enc_attn_window": 0,
+      "enc_conv_kernel": 3,
+      "enc_conv_layers": 0,
+      "enc_layers": 0,
+      "in_vocab_size": 8192,
+      "n_layers": 6,
+      "pre_quant_dropout": 0.25,
+      "pre_quant_groupnorm": 8,
+      "pre_quant_norm_type": "group",
+      "use_qk_norm": false,
+      "vector_quantizer": {
+        "aux": {
+          "fine_drop": {
+            "anneal_steps": 0,
+            "min_keep": 1,
+            "prob": 0.5,
+            "prob_start": null
+          }
+        },
+        "clss": "multiscale_residual_vector_quantizer",
+        "commitment_cost": 0.25,
+        "decay": 0.99,
+        "ema_scale_by_world_size": false,
+        "epsilon": 1e-05,
+        "kmeans_init": false,
+        "kmeans_iters": 10,
+        "learned_all_sampling": false,
+        "learned_l1_cascade": true,
+        "learned_l1_cascade_base": 4,
+        "learned_l1_cascade_norm": true,
+        "learned_l1_cascade_norm_final": true,
+        "learned_l1_sampling": true,
+        "levels": {
+          "end": 64,
+          "manual_levels": [
+            1,
+            4,
+            9,
+            16,
+            25,
+            36,
+            49,
+            64
+          ],
+          "num_levels": 8,
+          "scheduler": "linear",
+          "start": 1,
+          "use_manual_levels": true
+        },
+        "pre_quant_norm_type": "group",
+        "quant_resi": {
+          "arch": "blend",
+          "enabled": true,
+          "hidden_mult": 1,
+          "kernel_size": 3,
+          "learnable_ratio": false,
+          "ratio": 0.5,
+          "share_mode": 0
+        },
+        "restart_threshold": 0.0
+      },
+      "vqvae_vocab_size": [
+        8192,
+        8192,
+        8192,
+        8192,
+        8192,
+        8192,
+        8192,
+        8192
+      ]
+    },
+    "dataset": {
+      "pad_token_id": 1,
+      "train_source": "data/ncm/tinystories/train_bpe_8192.pickle",
+      "validate_source": "data/ncm/tinystories/valid_bpe_8192.pickle"
+    },
+    "experiment": "tinystories",
+    "mixed_precision": true,
+    "ncp": {
+      "attn_pattern": "block_diagonal",
+      "bias": true,
+      "dropout": 0.1,
+      "init_body_from_gpt2": false,
+      "l0_autoregressive": false,
+      "mlp_ratio": 4.0,
+      "n_embd": 768,
+      "n_head": 12,
+      "n_kv_heads": 12,
+      "n_layer": 12,
+      "predict_level0": false,
+      "prefix_len": 0,
+      "prefix_pad_token_id": null,
+      "query_level0": true,
+      "rope_base": 10000,
+      "shared_head_mlp": false,
+      "shared_head_mlp_blocks": 1,
+      "shared_head_mlp_legacy": false,
+      "shared_head_mlp_ln": false,
+      "shared_head_mlp_mult": 1.0,
+      "shared_head_mlp_residual": true,
+      "shared_output_head": true,
+      "tie_heads_bias": false,
+      "tie_heads_learn_scale": true,
+      "tie_heads_to_codebook": false,
+      "tie_heads_trainable": false,
+      "use_flex_attention": false,
+      "use_geglu": false,
+      "use_level_encoding": true,
+      "use_positional_encoding": true,
+      "use_qk_norm": true,
+      "use_relu2": false,
+      "use_rope": true,
+      "use_swiglu": false,
+      "vocab_size": [
+        8192,
+        8192,
+        8192,
+        8192,
+        8192,
+        8192,
+        8192,
+        8192
+      ]
+    },
+    "objectives": {
+      "apr_warmup_steps": 500,
+      "apr_weight": 1.0,
+      "mstok_reconstruction_weight": 1.0,
+      "mstok_residual_weight": 1.0,
+      "mstok_warmup_steps": 500,
+      "mstok_weight": 0.1,
+      "nsp_weight": 1.0,
+      "prediction_temperature": 1.0,
+      "soft_assignment_temperature": 1.2,
+      "teacher_ema_decay": 0.999,
+      "teacher_ema_warmup_steps": 1000
+    },
+    "optimization": {
+      "beta_1": 0.9,
+      "beta_2": 0.95,
+      "codec_beta_2": 0.95,
+      "codec_lr": 0.0005,
+      "codec_min_lr_ratio": 0.2,
+      "codec_warmup_iters": 500,
+      "generator_beta_2": 0.99,
+      "generator_lr": 0.0005,
+      "generator_min_lr_ratio": 0.02,
+      "generator_warmup_iters": 150,
+      "grad_accumulation_steps": 1,
+      "level_loss_alpha": 0.25,
+      "lr": 0.0005,
+      "max_grad_norm": 1.0,
+      "min_lr": 0.0001,
+      "reset_lr": false,
+      "token_corruption_prob": 0.0,
+      "warmup_iters": 500,
+      "weight_decay": 0.01
+    },
+    "output_dir": "outputs/ncm/tinystories-eostok-apr-full-v2-5epoch-ddp",
+    "pre_tokenizer": {
+      "merges_file_path": "data/ncm/tinystories/bpe_merges_8192.pickle",
+      "name": "bpe",
+      "special_tokens": {
+        "eos_token": "<|endoftext|>",
+        "pad_token": "<|pad|>"
+      },
+      "tokenizer": "bpe",
+      "vocab_file_path": "data/ncm/tinystories/bpe_vocab_8192.pickle"
+    },
+    "task": "summarization",
+    "torch_compile": {
+      "backend": "inductor",
+      "dynamic": false,
+      "enable": true,
+      "fullgraph": false,
+      "mode": "reduce-overhead",
+      "scope": "modules"
+    },
+    "training": {
+      "allow_random_codec": false,
+      "batch_size": 128,
+      "checkpoint_interval": 16563,
+      "demo_batch_size": 4,
+      "eval_batch_size": 16,
+      "eval_interval": 250,
+      "expected_global_batch_size": 512,
+      "expected_world_size": 4,
+      "find_unused_parameters": false,
+      "log_interval": 10,
+      "max_steps": 82813,
+      "n_epochs": 5.0,
+      "seed": 42,
+      "steps_per_epoch": null,
+      "stop_after_iters": null,
+      "total_iters": null,
+      "tqdm_interval": 10,
+      "val_iters": 20
+    },
+    "training_regime": "eostok_apr",
+    "wandb": {
+      "enable": true,
+      "group": "text-e2e-v2-5epoch-ddp",
+      "name": null,
+      "project": "mstok-tinystories",
+      "resume": null
+    }
+  }
+}