Add files using upload-large-folder tool

configs/.ipynb_checkpoints/pretrain_5090_stage1-checkpoint.json

@@ -0,0 +1,27 @@
+{
+  "seed": 42,
+  "train_dir": "data/pretokenized/train",
+  "val_dir": "data/pretokenized/val",
+  "output_dir": "outputs/pretrain_stage1",
+  "checkpoint_dir": "checkpoints/pretrain_stage1",
+  "init_from": null,
+  "resume_from": null,
+  "seq_len": 2048,
+  "micro_batch_size": 8,
+  "grad_accum_steps": 32,
+  "max_steps": 20000,
+  "warmup_steps": 2000,
+  "learning_rate": 0.003,
+  "min_lr": 0.0003,
+  "weight_decay": 0.1,
+  "beta1": 0.9,
+  "beta2": 0.95,
+  "grad_clip": 1.0,
+  "precision": "bf16",
+  "num_workers": 0,
+  "log_interval": 10,
+  "eval_interval": 250,
+  "eval_batches": 50,
+  "save_interval": 100,
+  "compile_model": false
+}

configs/.ipynb_checkpoints/pretrain_5090_stage2_anneal-checkpoint.json

@@ -0,0 +1,27 @@
+{
+  "seed": 42,
+  "train_dir": "data/pretokenized/train",
+  "val_dir": "data/pretokenized/val",
+  "output_dir": "outputs/pretrain_stage2",
+  "checkpoint_dir": "checkpoints/pretrain_stage2",
+  "init_from": "checkpoints/pretrain_stage1/last.pt",
+  "resume_from": null,
+  "seq_len": 8192,
+  "micro_batch_size": 2,
+  "grad_accum_steps": 16,
+  "max_steps": 1000,
+  "warmup_steps": 100,
+  "learning_rate": 0.001,
+  "min_lr": 0.0001,
+  "weight_decay": 0.1,
+  "beta1": 0.9,
+  "beta2": 0.95,
+  "grad_clip": 1.0,
+  "precision": "bf16",
+  "num_workers": 0,
+  "log_interval": 5,
+  "eval_interval": 100,
+  "eval_batches": 20,
+  "save_interval": 50,
+  "compile_model": false
+}

configs/data_mix_10b.json

@@ -0,0 +1,56 @@
+{
+  "tokenizer_sample_documents": 2000000,
+  "tokenizer_min_frequency": 2,
+  "tokenizer_special_tokens": [
+    "<pad>",
+    "<bos>",
+    "<eos>",
+    "<unk>"
+  ],
+  "train_tokens": 10000000000,
+  "val_tokens": 20000000,
+  "shard_size_tokens": 100000000,
+  "sources": [
+    {
+      "name": "fineweb_edu",
+      "path": "HuggingFaceFW/fineweb-edu",
+      "config_name": "sample-10BT",
+      "split": "train",
+      "text_field": "text",
+      "weight": 0.6,
+      "streaming": true,
+      "shuffle_buffer": 10000
+    },
+    {
+      "name": "cosmopedia_v2",
+      "path": "HuggingFaceTB/smollm-corpus",
+      "config_name": "cosmopedia-v2",
+      "split": "train",
+      "text_field": "text",
+      "weight": 0.2,
+      "streaming": true,
+      "shuffle_buffer": 10000
+    },
+    {
+      "name": "the_stack_python",
+      "path": "bigcode/the-stack-dedup",
+      "config_name": null,
+      "data_dir": "data/python",
+      "split": "train",
+      "text_field": "content",
+      "weight": 0.1,
+      "streaming": true,
+      "shuffle_buffer": 2000
+    },
+    {
+      "name": "finemath",
+      "path": "HuggingFaceTB/finemath",
+      "config_name": "finemath-4plus",
+      "split": "train",
+      "text_field": "text",
+      "weight": 0.1,
+      "streaming": true,
+      "shuffle_buffer": 5000
+    }
+  ]
+}

configs/model_70m.json

@@ -0,0 +1,17 @@
+{
+  "vocab_size": 49152,
+  "max_seq_len": 8192,
+  "d_model": 384,
+  "n_layers": 32,
+  "n_heads": 6,
+  "ffn_hidden_dim": 1024,
+  "rope_theta": 10000.0,
+  "rms_norm_eps": 1e-05,
+  "initializer_range": 0.02,
+  "dropout": 0.0,
+  "tie_word_embeddings": true,
+  "bias": false,
+  "pad_token_id": 0,
+  "bos_token_id": 1,
+  "eos_token_id": 2
+}

configs/pretrain_5090_stage1.json

@@ -0,0 +1,27 @@
+{
+  "seed": 42,
+  "train_dir": "data/pretokenized/train",
+  "val_dir": "data/pretokenized/val",
+  "output_dir": "outputs/pretrain_stage1",
+  "checkpoint_dir": "checkpoints/pretrain_stage1",
+  "init_from": null,
+  "resume_from": null,
+  "seq_len": 2048,
+  "micro_batch_size": 8,
+  "grad_accum_steps": 32,
+  "max_steps": 20000,
+  "warmup_steps": 2000,
+  "learning_rate": 0.003,
+  "min_lr": 0.0003,
+  "weight_decay": 0.1,
+  "beta1": 0.9,
+  "beta2": 0.95,
+  "grad_clip": 1.0,
+  "precision": "bf16",
+  "num_workers": 0,
+  "log_interval": 10,
+  "eval_interval": 250,
+  "eval_batches": 50,
+  "save_interval": 100,
+  "compile_model": false
+}

configs/pretrain_5090_stage2_anneal.json

@@ -0,0 +1,27 @@
+{
+  "seed": 42,
+  "train_dir": "data/pretokenized/train",
+  "val_dir": "data/pretokenized/val",
+  "output_dir": "outputs/pretrain_stage2",
+  "checkpoint_dir": "checkpoints/pretrain_stage2",
+  "init_from": "checkpoints/pretrain_stage1/last.pt",
+  "resume_from": null,
+  "seq_len": 8192,
+  "micro_batch_size": 2,
+  "grad_accum_steps": 16,
+  "max_steps": 1000,
+  "warmup_steps": 100,
+  "learning_rate": 0.001,
+  "min_lr": 0.0001,
+  "weight_decay": 0.1,
+  "beta1": 0.9,
+  "beta2": 0.95,
+  "grad_clip": 1.0,
+  "precision": "bf16",
+  "num_workers": 0,
+  "log_interval": 5,
+  "eval_interval": 100,
+  "eval_batches": 20,
+  "save_interval": 50,
+  "compile_model": false
+}

configs/pretrain_mps_dryrun.json

@@ -0,0 +1,27 @@
+{
+  "seed": 42,
+  "train_dir": "data/pretokenized/train",
+  "val_dir": "data/pretokenized/val",
+  "output_dir": "outputs/pretrain_mps_dryrun",
+  "checkpoint_dir": "checkpoints/pretrain_mps_dryrun",
+  "init_from": null,
+  "resume_from": null,
+  "seq_len": 512,
+  "micro_batch_size": 1,
+  "grad_accum_steps": 4,
+  "max_steps": 500,
+  "warmup_steps": 50,
+  "learning_rate": 0.001,
+  "min_lr": 0.0001,
+  "weight_decay": 0.1,
+  "beta1": 0.9,
+  "beta2": 0.95,
+  "grad_clip": 1.0,
+  "precision": "fp32",
+  "num_workers": 0,
+  "log_interval": 1,
+  "eval_interval": 10,
+  "eval_batches": 2,
+  "save_interval": 10,
+  "compile_model": false
+}

configs/sft_5090.json

@@ -0,0 +1,26 @@
+{
+  "seed": 42,
+  "dataset_path": "data/sft/processed",
+  "output_dir": "outputs/sft",
+  "checkpoint_dir": "checkpoints/sft",
+  "init_from": "checkpoints/pretrain_stage2/last.pt",
+  "resume_from": null,
+  "seq_len": 2048,
+  "micro_batch_size": 8,
+  "grad_accum_steps": 16,
+  "max_steps": 5000,
+  "warmup_steps": 200,
+  "learning_rate": 0.0005,
+  "min_lr": 5e-05,
+  "weight_decay": 0.01,
+  "beta1": 0.9,
+  "beta2": 0.95,
+  "grad_clip": 1.0,
+  "precision": "bf16",
+  "num_workers": 0,
+  "log_interval": 10,
+  "eval_interval": 100,
+  "eval_batches": 50,
+  "save_interval": 200,
+  "compile_model": false
+}

configs/sft_data_smoltalk.json

@@ -0,0 +1,12 @@
+{
+  "path": "HuggingFaceTB/smoltalk",
+  "config_name": null,
+  "split": "train",
+  "revision": null,
+  "streaming": false,
+  "shuffle": true,
+  "format": "messages",
+  "messages_field": "messages",
+  "val_examples": 2000,
+  "max_train_examples": 200000
+}

data/README.md

@@ -0,0 +1,3 @@
+---
+license: mit
+---

data/pretokenized/dataset_summary.json

@@ -0,0 +1,198 @@
+{
+  "tokenizer": {
+    "vocab_size": 49152,
+    "special_tokens": {
+      "pad_token": "<pad>",
+      "bos_token": "<bos>",
+      "eos_token": "<eos>",
+      "unk_token": "<unk>",
+      "pad_token_id": 0,
+      "bos_token_id": 1,
+      "eos_token_id": 2,
+      "unk_token_id": 3
+    },
+    "data_config": {
+      "sources": [
+        {
+          "name": "fineweb_edu",
+          "path": "HuggingFaceFW/fineweb-edu",
+          "split": "train",
+          "weight": 0.6,
+          "text_field": "text",
+          "config_name": "sample-10BT",
+          "data_dir": null,
+          "revision": null,
+          "streaming": true,
+          "shuffle_buffer": 10000,
+          "sample_documents": null
+        },
+        {
+          "name": "cosmopedia_v2",
+          "path": "HuggingFaceTB/smollm-corpus",
+          "split": "train",
+          "weight": 0.2,
+          "text_field": "text",
+          "config_name": "cosmopedia-v2",
+          "data_dir": null,
+          "revision": null,
+          "streaming": true,
+          "shuffle_buffer": 10000,
+          "sample_documents": null
+        },
+        {
+          "name": "the_stack_python",
+          "path": "bigcode/the-stack-dedup",
+          "split": "train",
+          "weight": 0.1,
+          "text_field": "content",
+          "config_name": null,
+          "data_dir": "data/python",
+          "revision": null,
+          "streaming": true,
+          "shuffle_buffer": 2000,
+          "sample_documents": null
+        },
+        {
+          "name": "finemath",
+          "path": "HuggingFaceTB/finemath",
+          "split": "train",
+          "weight": 0.1,
+          "text_field": "text",
+          "config_name": "finemath-4plus",
+          "data_dir": null,
+          "revision": null,
+          "streaming": true,
+          "shuffle_buffer": 5000,
+          "sample_documents": null
+        }
+      ],
+      "tokenizer_sample_documents": 2000000,
+      "tokenizer_min_frequency": 2,
+      "tokenizer_special_tokens": [
+        "<pad>",
+        "<bos>",
+        "<eos>",
+        "<unk>"
+      ],
+      "train_tokens": 10000000000,
+      "val_tokens": 20000000,
+      "shard_size_tokens": 100000000
+    }
+  },
+  "data_config": {
+    "sources": [
+      {
+        "name": "fineweb_edu",
+        "path": "HuggingFaceFW/fineweb-edu",
+        "split": "train",
+        "weight": 0.6,
+        "text_field": "text",
+        "config_name": "sample-10BT",
+        "data_dir": null,
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 10000,
+        "sample_documents": null
+      },
+      {
+        "name": "cosmopedia_v2",
+        "path": "HuggingFaceTB/smollm-corpus",
+        "split": "train",
+        "weight": 0.2,
+        "text_field": "text",
+        "config_name": "cosmopedia-v2",
+        "data_dir": null,
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 10000,
+        "sample_documents": null
+      },
+      {
+        "name": "the_stack_python",
+        "path": "bigcode/the-stack-dedup",
+        "split": "train",
+        "weight": 0.1,
+        "text_field": "content",
+        "config_name": null,
+        "data_dir": "data/python",
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 2000,
+        "sample_documents": null
+      },
+      {
+        "name": "finemath",
+        "path": "HuggingFaceTB/finemath",
+        "split": "train",
+        "weight": 0.1,
+        "text_field": "text",
+        "config_name": "finemath-4plus",
+        "data_dir": null,
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 5000,
+        "sample_documents": null
+      }
+    ],
+    "tokenizer_sample_documents": 2000000,
+    "tokenizer_min_frequency": 2,
+    "tokenizer_special_tokens": [
+      "<pad>",
+      "<bos>",
+      "<eos>",
+      "<unk>"
+    ],
+    "train_tokens": 10000000000,
+    "val_tokens": 20000000,
+    "shard_size_tokens": 100000000
+  },
+  "mixing_strategy": "global_interleaving_weighted_progress_balancing",
+  "train_target_tokens": 10000000000,
+  "val_target_tokens": 20000000,
+  "train_tokens_written": 10000000000,
+  "val_tokens_written": 20000000,
+  "train_shards": 100,
+  "val_shards": 1,
+  "sources": {
+    "fineweb_edu": {
+      "path": "HuggingFaceFW/fineweb-edu",
+      "data_dir": null,
+      "split": "train",
+      "train_target_tokens": 6000000000,
+      "val_target_tokens": 12000000,
+      "train_tokens_written": 6000000000,
+      "val_tokens_written": 12000000,
+      "documents_used": 5922817
+    },
+    "cosmopedia_v2": {
+      "path": "HuggingFaceTB/smollm-corpus",
+      "data_dir": null,
+      "split": "train",
+      "train_target_tokens": 2000000000,
+      "val_target_tokens": 4000000,
+      "train_tokens_written": 2000000000,
+      "val_tokens_written": 4000000,
+      "documents_used": 2792704
+    },
+    "the_stack_python": {
+      "path": "bigcode/the-stack-dedup",
+      "data_dir": "data/python",
+      "split": "train",
+      "train_target_tokens": 1000000000,
+      "val_target_tokens": 2000000,
+      "train_tokens_written": 1000000000,
+      "val_tokens_written": 2000000,
+      "documents_used": 684540
+    },
+    "finemath": {
+      "path": "HuggingFaceTB/finemath",
+      "data_dir": null,
+      "split": "train",
+      "train_target_tokens": 1000000000,
+      "val_target_tokens": 2000000,
+      "train_tokens_written": 1000000000,
+      "val_tokens_written": 2000000,
+      "documents_used": 692367
+    }
+  }
+}

data/pretokenized/train/train_manifest.json

@@ -0,0 +1,502 @@
+[
+  {
+    "path": "train_00000.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00001.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00002.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00003.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00004.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00005.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00006.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00007.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00008.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00009.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00010.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00011.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00012.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00013.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00014.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00015.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00016.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00017.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00018.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00019.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00020.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00021.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00022.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00023.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00024.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00025.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00026.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00027.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00028.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00029.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00030.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00031.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00032.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00033.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00034.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00035.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00036.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00037.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00038.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00039.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00040.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00041.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00042.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00043.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00044.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00045.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00046.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00047.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00048.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00049.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00050.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00051.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00052.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00053.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00054.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00055.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00056.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00057.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00058.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00059.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00060.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00061.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00062.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00063.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00064.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00065.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00066.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00067.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00068.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00069.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00070.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00071.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00072.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00073.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00074.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00075.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00076.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00077.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00078.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00079.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00080.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00081.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00082.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00083.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00084.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00085.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00086.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00087.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00088.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00089.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00090.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00091.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00092.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00093.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00094.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00095.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00096.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00097.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00098.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  },
+  {
+    "path": "train_00099.bin",
+    "num_tokens": 100000000,
+    "dtype": "uint16"
+  }
+]

data/pretokenized/val/val_manifest.json

@@ -0,0 +1,7 @@
+[
+  {
+    "path": "val_00000.bin",
+    "num_tokens": 20000000,
+    "dtype": "uint16"
+  }
+]

data/tokenizer/logs/train_tokenizer_20260312_114030.log

@@ -0,0 +1,11 @@
+2026-03-12 11:40:30,043 | INFO | Tokenizer training started
+2026-03-12 11:40:30,044 | INFO | Log file: data/tokenizer/logs/train_tokenizer_20260312_114030.log
+2026-03-12 11:40:30,044 | INFO | Arguments | data_config=configs/data_mix_10b.json output_dir=data/tokenizer vocab_size=49152 seed=42
+2026-03-12 11:40:30,044 | INFO | Tokenizer config | sample_documents=2,000,000 min_frequency=2 special_tokens=['<pad>', '<bos>', '<eos>', '<unk>'] num_sources=4
+2026-03-12 11:40:30,044 | INFO | Tokenizer source start | name=fineweb_edu path=HuggingFaceFW/fineweb-edu data_dir=None split=train text_field=text limit_docs=1,200,000 streaming=True
+2026-03-12 11:51:35,669 | INFO | Tokenizer source start | name=cosmopedia_v2 path=HuggingFaceTB/smollm-corpus data_dir=None split=train text_field=text limit_docs=400,000 streaming=True
+2026-03-12 11:55:58,013 | INFO | Tokenizer source start | name=the_stack_python path=bigcode/the-stack-dedup data_dir=data/python split=train text_field=content limit_docs=200,000 streaming=True
+2026-03-12 12:00:03,620 | INFO | Tokenizer source start | name=finemath path=HuggingFaceTB/finemath data_dir=None split=train text_field=text limit_docs=200,000 streaming=True
+2026-03-12 12:08:46,619 | INFO | Tokenizer saved | path=data/tokenizer/tokenizer.json
+2026-03-12 12:08:46,630 | INFO | Tokenizer summary | vocab_size=49152 pad_id=0 bos_id=1 eos_id=2 unk_id=3
+2026-03-12 12:08:46,630 | INFO | Tokenizer metadata saved | path=data/tokenizer/tokenizer_meta.json

data/tokenizer/tokenizer_meta.json

@@ -0,0 +1,80 @@
+{
+  "vocab_size": 49152,
+  "special_tokens": {
+    "pad_token": "<pad>",
+    "bos_token": "<bos>",
+    "eos_token": "<eos>",
+    "unk_token": "<unk>",
+    "pad_token_id": 0,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "unk_token_id": 3
+  },
+  "data_config": {
+    "sources": [
+      {
+        "name": "fineweb_edu",
+        "path": "HuggingFaceFW/fineweb-edu",
+        "split": "train",
+        "weight": 0.6,
+        "text_field": "text",
+        "config_name": "sample-10BT",
+        "data_dir": null,
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 10000,
+        "sample_documents": null
+      },
+      {
+        "name": "cosmopedia_v2",
+        "path": "HuggingFaceTB/smollm-corpus",
+        "split": "train",
+        "weight": 0.2,
+        "text_field": "text",
+        "config_name": "cosmopedia-v2",
+        "data_dir": null,
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 10000,
+        "sample_documents": null
+      },
+      {
+        "name": "the_stack_python",
+        "path": "bigcode/the-stack-dedup",
+        "split": "train",
+        "weight": 0.1,
+        "text_field": "content",
+        "config_name": null,
+        "data_dir": "data/python",
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 2000,
+        "sample_documents": null
+      },
+      {
+        "name": "finemath",
+        "path": "HuggingFaceTB/finemath",
+        "split": "train",
+        "weight": 0.1,
+        "text_field": "text",
+        "config_name": "finemath-4plus",
+        "data_dir": null,
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 5000,
+        "sample_documents": null
+      }
+    ],
+    "tokenizer_sample_documents": 2000000,
+    "tokenizer_min_frequency": 2,
+    "tokenizer_special_tokens": [
+      "<pad>",
+      "<bos>",
+      "<eos>",
+      "<unk>"
+    ],
+    "train_tokens": 10000000000,
+    "val_tokens": 20000000,
+    "shard_size_tokens": 100000000
+  }
+}

data/tokenizer/tokenizer_summary.json

@@ -0,0 +1,80 @@
+{
+  "vocab_size": 49152,
+  "special_tokens": {
+    "pad_token": "<pad>",
+    "bos_token": "<bos>",
+    "eos_token": "<eos>",
+    "unk_token": "<unk>",
+    "pad_token_id": 0,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "unk_token_id": 3
+  },
+  "data_config": {
+    "sources": [
+      {
+        "name": "fineweb_edu",
+        "path": "HuggingFaceFW/fineweb-edu",
+        "split": "train",
+        "weight": 0.6,
+        "text_field": "text",
+        "config_name": "sample-10BT",
+        "data_dir": null,
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 10000,
+        "sample_documents": null
+      },
+      {
+        "name": "cosmopedia_v2",
+        "path": "HuggingFaceTB/smollm-corpus",
+        "split": "train",
+        "weight": 0.2,
+        "text_field": "text",
+        "config_name": "cosmopedia-v2",
+        "data_dir": null,
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 10000,
+        "sample_documents": null
+      },
+      {
+        "name": "the_stack_python",
+        "path": "bigcode/the-stack-dedup",
+        "split": "train",
+        "weight": 0.1,
+        "text_field": "content",
+        "config_name": null,
+        "data_dir": "data/python",
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 2000,
+        "sample_documents": null
+      },
+      {
+        "name": "finemath",
+        "path": "HuggingFaceTB/finemath",
+        "split": "train",
+        "weight": 0.1,
+        "text_field": "text",
+        "config_name": "finemath-4plus",
+        "data_dir": null,
+        "revision": null,
+        "streaming": true,
+        "shuffle_buffer": 5000,
+        "sample_documents": null
+      }
+    ],
+    "tokenizer_sample_documents": 2000000,
+    "tokenizer_min_frequency": 2,
+    "tokenizer_special_tokens": [
+      "<pad>",
+      "<bos>",
+      "<eos>",
+      "<unk>"
+    ],
+    "train_tokens": 10000000000,
+    "val_tokens": 20000000,
+    "shard_size_tokens": 100000000
+  }
+}

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140730.jsonl

@@ -0,0 +1,2 @@
+{"event": "run_started", "timestamp": "2026-03-13T14:07:30", "log_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140730.log", "metrics_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140730.jsonl", "model_config": {"vocab_size": 49152, "max_seq_len": 8192, "d_model": 384, "n_layers": 32, "n_heads": 6, "ffn_hidden_dim": 1024, "rope_theta": 10000.0, "rms_norm_eps": 1e-05, "initializer_range": 0.02, "dropout": 0.0, "tie_word_embeddings": true, "bias": false, "pad_token_id": 0, "bos_token_id": 1, "eos_token_id": 2}, "train_config": {"seed": 42, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val", "output_dir": "outputs/pretrain_mps_dryrun", "checkpoint_dir": "checkpoints/pretrain_mps_dryrun", "init_from": null, "resume_from": null, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "max_steps": 20, "warmup_steps": 5, "learning_rate": 0.001, "min_lr": 0.0001, "weight_decay": 0.1, "beta1": 0.9, "beta2": 0.95, "grad_clip": 1.0, "precision": "fp32", "num_workers": 0, "log_interval": 1, "eval_interval": 10, "eval_batches": 2, "save_interval": 10, "compile_model": false}, "args": {"model_config": "configs/model_70m.json", "train_config": "configs/pretrain_mps_dryrun.json", "max_steps_override": null}}
+{"event": "runtime_summary", "timestamp": "2026-03-13T14:07:34", "device": "mps", "precision": "fp32", "compile_model": false, "parameters": 75571584, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "tokens_per_step": 2048, "num_train_shards": 100, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val"}

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140730.log

@@ -0,0 +1,10 @@
+2026-03-13 14:07:30,831 | INFO | Pretraining started
+2026-03-13 14:07:30,832 | INFO | Log file: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140730.log
+2026-03-13 14:07:30,832 | INFO | Metrics JSONL: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140730.jsonl
+2026-03-13 14:07:30,832 | INFO | Arguments | model_config=configs/model_70m.json train_config=configs/pretrain_mps_dryrun.json max_steps_override=None
+2026-03-13 14:07:30,832 | INFO | Model config | {'vocab_size': 49152, 'max_seq_len': 8192, 'd_model': 384, 'n_layers': 32, 'n_heads': 6, 'ffn_hidden_dim': 1024, 'rope_theta': 10000.0, 'rms_norm_eps': 1e-05, 'initializer_range': 0.02, 'dropout': 0.0, 'tie_word_embeddings': True, 'bias': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2}
+2026-03-13 14:07:30,832 | INFO | Train config | {'seed': 42, 'train_dir': 'data/pretokenized/train', 'val_dir': 'data/pretokenized/val', 'output_dir': 'outputs/pretrain_mps_dryrun', 'checkpoint_dir': 'checkpoints/pretrain_mps_dryrun', 'init_from': None, 'resume_from': None, 'seq_len': 512, 'micro_batch_size': 1, 'grad_accum_steps': 4, 'max_steps': 20, 'warmup_steps': 5, 'learning_rate': 0.001, 'min_lr': 0.0001, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'precision': 'fp32', 'num_workers': 0, 'log_interval': 1, 'eval_interval': 10, 'eval_batches': 2, 'save_interval': 10, 'compile_model': False}
+2026-03-13 14:07:34,596 | INFO | Device summary | device=mps precision=fp32 compile_model=False
+2026-03-13 14:07:34,597 | INFO | Model summary | parameters=75.57M
+2026-03-13 14:07:34,597 | INFO | Batch summary | seq_len=512 micro_batch_size=1 grad_accum_steps=4 tokens_per_step=2,048
+2026-03-13 14:07:34,597 | INFO | Dataset summary | train_dir=data/pretokenized/train val_dir=data/pretokenized/val num_train_shards=100

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140907.jsonl

@@ -0,0 +1,27 @@
+{"event": "run_started", "timestamp": "2026-03-13T14:09:07", "log_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140907.log", "metrics_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140907.jsonl", "model_config": {"vocab_size": 49152, "max_seq_len": 8192, "d_model": 384, "n_layers": 32, "n_heads": 6, "ffn_hidden_dim": 1024, "rope_theta": 10000.0, "rms_norm_eps": 1e-05, "initializer_range": 0.02, "dropout": 0.0, "tie_word_embeddings": true, "bias": false, "pad_token_id": 0, "bos_token_id": 1, "eos_token_id": 2}, "train_config": {"seed": 42, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val", "output_dir": "outputs/pretrain_mps_dryrun", "checkpoint_dir": "checkpoints/pretrain_mps_dryrun", "init_from": null, "resume_from": null, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "max_steps": 20, "warmup_steps": 5, "learning_rate": 0.001, "min_lr": 0.0001, "weight_decay": 0.1, "beta1": 0.9, "beta2": 0.95, "grad_clip": 1.0, "precision": "fp32", "num_workers": 0, "log_interval": 1, "eval_interval": 10, "eval_batches": 2, "save_interval": 10, "compile_model": false}, "args": {"model_config": "configs/model_70m.json", "train_config": "configs/pretrain_mps_dryrun.json", "max_steps_override": null}}
+{"event": "runtime_summary", "timestamp": "2026-03-13T14:09:10", "device": "mps", "precision": "fp32", "compile_model": false, "parameters": 75571584, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "tokens_per_step": 2048, "num_train_shards": 100, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val"}
+{"event": "train", "timestamp": "2026-03-13T14:09:12", "step": 1, "loss": 10.848917245864868, "lr": 0.0002, "tok_per_sec": 961.8014053409653, "grad_norm": 5.573695659637451, "tokens_seen": 2048, "elapsed_sec": 2.1293377080000937, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:13", "step": 2, "loss": 10.763139724731445, "lr": 0.0004, "tok_per_sec": 1605.8989070685525, "grad_norm": 8.322466850280762, "tokens_seen": 4096, "elapsed_sec": 1.275298208987806, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:14", "step": 3, "loss": 10.356749057769775, "lr": 0.0006000000000000001, "tok_per_sec": 2736.6722939747565, "grad_norm": 2.6283912658691406, "tokens_seen": 6144, "elapsed_sec": 0.7483541250112467, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:14", "step": 4, "loss": 10.376826286315918, "lr": 0.0008, "tok_per_sec": 2756.3866613090086, "grad_norm": 2.217130184173584, "tokens_seen": 8192, "elapsed_sec": 0.7430017089936882, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:15", "step": 5, "loss": 10.231549263000488, "lr": 0.001, "tok_per_sec": 2715.4393876891622, "grad_norm": 12.64534854888916, "tokens_seen": 10240, "elapsed_sec": 0.754205750010442, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:16", "step": 6, "loss": 9.938905477523804, "lr": 0.001, "tok_per_sec": 2725.141290121042, "grad_norm": 1.7282862663269043, "tokens_seen": 12288, "elapsed_sec": 0.7515206669922918, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:17", "step": 7, "loss": 9.66855764389038, "lr": 0.0009901664203302125, "tok_per_sec": 2738.185267024283, "grad_norm": 1.9499105215072632, "tokens_seen": 14336, "elapsed_sec": 0.7479406250058673, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:17", "step": 8, "loss": 9.335453271865845, "lr": 0.0009610954559391703, "tok_per_sec": 2751.8910854624123, "grad_norm": 1.7210659980773926, "tokens_seen": 16384, "elapsed_sec": 0.7442154999589548, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:18", "step": 9, "loss": 9.267512798309326, "lr": 0.0009140576474687263, "tok_per_sec": 2708.5891084687337, "grad_norm": 1.762829065322876, "tokens_seen": 18432, "elapsed_sec": 0.7561132080154493, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:19", "step": 10, "loss": 8.833673238754272, "lr": 0.0008511087728614862, "tok_per_sec": 2765.4777296002535, "grad_norm": 1.700391173362732, "tokens_seen": 20480, "elapsed_sec": 0.7405592090217397, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:09:19", "step": 10, "val_loss": 9.096094608306885, "perplexity": 8920.386982370957, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:09:33", "step": 10, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000010.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 20480}
+{"event": "train", "timestamp": "2026-03-13T14:09:34", "step": 11, "loss": 8.795855522155762, "lr": 0.0007750000000000001, "tok_per_sec": 135.3521374189279, "grad_norm": 1.4899625778198242, "tokens_seen": 22528, "elapsed_sec": 15.130902540986426, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:35", "step": 12, "loss": 8.558577060699463, "lr": 0.0006890576474687264, "tok_per_sec": 2659.660815260197, "grad_norm": 1.5879555940628052, "tokens_seen": 24576, "elapsed_sec": 0.7700230000191368, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:36", "step": 13, "loss": 8.595118284225464, "lr": 0.0005970378084704442, "tok_per_sec": 2709.453151326185, "grad_norm": 1.3136154413223267, "tokens_seen": 26624, "elapsed_sec": 0.7558720840024762, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:36", "step": 14, "loss": 8.341074705123901, "lr": 0.000502962191529556, "tok_per_sec": 2569.4064364370934, "grad_norm": 1.2977045774459839, "tokens_seen": 28672, "elapsed_sec": 0.7970712499809451, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:37", "step": 15, "loss": 8.268006086349487, "lr": 0.0004109423525312737, "tok_per_sec": 2272.894269158833, "grad_norm": 1.197304368019104, "tokens_seen": 30720, "elapsed_sec": 0.9010537919821218, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:38", "step": 16, "loss": 8.30242395401001, "lr": 0.0003250000000000001, "tok_per_sec": 2636.6892030166064, "grad_norm": 1.1259286403656006, "tokens_seen": 32768, "elapsed_sec": 0.7767316669924185, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:39", "step": 17, "loss": 8.6144118309021, "lr": 0.00024889122713851394, "tok_per_sec": 2628.5409836443887, "grad_norm": 0.9170812368392944, "tokens_seen": 34816, "elapsed_sec": 0.7791394590167329, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:40", "step": 18, "loss": 8.439870119094849, "lr": 0.00018594235253127368, "tok_per_sec": 2699.6840414438493, "grad_norm": 1.6393400430679321, "tokens_seen": 36864, "elapsed_sec": 0.758607292023953, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:40", "step": 19, "loss": 7.915311574935913, "lr": 0.00013890454406082956, "tok_per_sec": 2709.5957985032933, "grad_norm": 1.111694097518921, "tokens_seen": 38912, "elapsed_sec": 0.755832290975377, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:09:41", "step": 20, "loss": 7.964773654937744, "lr": 0.00010983357966978745, "tok_per_sec": 2689.7526879403435, "grad_norm": 1.00663423538208, "tokens_seen": 40960, "elapsed_sec": 0.7614082919899374, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:09:41", "step": 20, "val_loss": 8.758275032043457, "perplexity": 6363.125917448135, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:09:54", "step": 20, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000020.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 40960}
+{"event": "run_finished", "timestamp": "2026-03-13T14:09:54", "final_step": 20, "tokens_seen": 40960}

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140907.log

@@ -0,0 +1,14 @@
+2026-03-13 14:09:07,112 | INFO | Pretraining started
+2026-03-13 14:09:07,112 | INFO | Log file: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140907.log
+2026-03-13 14:09:07,112 | INFO | Metrics JSONL: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_140907.jsonl
+2026-03-13 14:09:07,112 | INFO | Arguments | model_config=configs/model_70m.json train_config=configs/pretrain_mps_dryrun.json max_steps_override=None
+2026-03-13 14:09:07,112 | INFO | Model config | {'vocab_size': 49152, 'max_seq_len': 8192, 'd_model': 384, 'n_layers': 32, 'n_heads': 6, 'ffn_hidden_dim': 1024, 'rope_theta': 10000.0, 'rms_norm_eps': 1e-05, 'initializer_range': 0.02, 'dropout': 0.0, 'tie_word_embeddings': True, 'bias': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2}
+2026-03-13 14:09:07,112 | INFO | Train config | {'seed': 42, 'train_dir': 'data/pretokenized/train', 'val_dir': 'data/pretokenized/val', 'output_dir': 'outputs/pretrain_mps_dryrun', 'checkpoint_dir': 'checkpoints/pretrain_mps_dryrun', 'init_from': None, 'resume_from': None, 'seq_len': 512, 'micro_batch_size': 1, 'grad_accum_steps': 4, 'max_steps': 20, 'warmup_steps': 5, 'learning_rate': 0.001, 'min_lr': 0.0001, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'precision': 'fp32', 'num_workers': 0, 'log_interval': 1, 'eval_interval': 10, 'eval_batches': 2, 'save_interval': 10, 'compile_model': False}
+2026-03-13 14:09:10,064 | INFO | Device summary | device=mps precision=fp32 compile_model=False
+2026-03-13 14:09:10,065 | INFO | Model summary | parameters=75.57M
+2026-03-13 14:09:10,065 | INFO | Batch summary | seq_len=512 micro_batch_size=1 grad_accum_steps=4 tokens_per_step=2,048
+2026-03-13 14:09:10,065 | INFO | Dataset summary | train_dir=data/pretokenized/train val_dir=data/pretokenized/val num_train_shards=100
+2026-03-13 14:09:19,703 | INFO | Eval step | step=10 val_loss=9.0961 perplexity=8920.39
+2026-03-13 14:09:33,612 | INFO | Checkpoint saved | step=10 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000010.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt
+2026-03-13 14:09:41,833 | INFO | Eval step | step=20 val_loss=8.7583 perplexity=6363.13
+2026-03-13 14:09:54,172 | INFO | Checkpoint saved | step=20 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000020.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_141224.jsonl

@@ -0,0 +1,27 @@
+{"event": "run_started", "timestamp": "2026-03-13T14:12:24", "log_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_141224.log", "metrics_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_141224.jsonl", "model_config": {"vocab_size": 49152, "max_seq_len": 8192, "d_model": 384, "n_layers": 32, "n_heads": 6, "ffn_hidden_dim": 1024, "rope_theta": 10000.0, "rms_norm_eps": 1e-05, "initializer_range": 0.02, "dropout": 0.0, "tie_word_embeddings": true, "bias": false, "pad_token_id": 0, "bos_token_id": 1, "eos_token_id": 2}, "train_config": {"seed": 42, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val", "output_dir": "outputs/pretrain_mps_dryrun", "checkpoint_dir": "checkpoints/pretrain_mps_dryrun", "init_from": null, "resume_from": null, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "max_steps": 20, "warmup_steps": 5, "learning_rate": 0.001, "min_lr": 0.0001, "weight_decay": 0.1, "beta1": 0.9, "beta2": 0.95, "grad_clip": 1.0, "precision": "fp32", "num_workers": 0, "log_interval": 1, "eval_interval": 10, "eval_batches": 2, "save_interval": 10, "compile_model": false}, "args": {"model_config": "configs/model_70m.json", "train_config": "configs/pretrain_mps_dryrun.json", "max_steps_override": null}}
+{"event": "runtime_summary", "timestamp": "2026-03-13T14:12:27", "device": "mps", "precision": "fp32", "compile_model": false, "parameters": 75571584, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "tokens_per_step": 2048, "num_train_shards": 100, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val"}
+{"event": "train", "timestamp": "2026-03-13T14:12:28", "step": 1, "loss": 10.848917245864868, "lr": 0.0002, "tok_per_sec": 1572.3567196374945, "grad_norm": 5.573695659637451, "tokens_seen": 2048, "elapsed_sec": 1.3025034169550054, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:29", "step": 2, "loss": 10.763139724731445, "lr": 0.0004, "tok_per_sec": 1953.4684072997784, "grad_norm": 8.322466850280762, "tokens_seen": 4096, "elapsed_sec": 1.0483916670200415, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:30", "step": 3, "loss": 10.356749057769775, "lr": 0.0006000000000000001, "tok_per_sec": 2746.91589368826, "grad_norm": 2.6283912658691406, "tokens_seen": 6144, "elapsed_sec": 0.7455634170328267, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:31", "step": 4, "loss": 10.376826524734497, "lr": 0.0008, "tok_per_sec": 2765.8319145818245, "grad_norm": 2.217130184173584, "tokens_seen": 8192, "elapsed_sec": 0.7404643750051036, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:32", "step": 5, "loss": 10.231549263000488, "lr": 0.001, "tok_per_sec": 2769.4683023376706, "grad_norm": 12.645360946655273, "tokens_seen": 10240, "elapsed_sec": 0.7394921249942854, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:32", "step": 6, "loss": 9.938905715942383, "lr": 0.001, "tok_per_sec": 2767.895009558497, "grad_norm": 1.7282859086990356, "tokens_seen": 12288, "elapsed_sec": 0.7399124579969794, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:33", "step": 7, "loss": 9.66855764389038, "lr": 0.0009901664203302125, "tok_per_sec": 2751.478083715354, "grad_norm": 1.9499105215072632, "tokens_seen": 14336, "elapsed_sec": 0.7443272080272436, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:34", "step": 8, "loss": 9.335453271865845, "lr": 0.0009610954559391703, "tok_per_sec": 2761.168994935579, "grad_norm": 1.7210659980773926, "tokens_seen": 16384, "elapsed_sec": 0.7417148330132477, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:34", "step": 9, "loss": 9.267512798309326, "lr": 0.0009140576474687263, "tok_per_sec": 2775.2087673169717, "grad_norm": 1.762829065322876, "tokens_seen": 18432, "elapsed_sec": 0.737962500017602, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:35", "step": 10, "loss": 8.833673000335693, "lr": 0.0008511087728614862, "tok_per_sec": 2743.4623249730666, "grad_norm": 1.700391173362732, "tokens_seen": 20480, "elapsed_sec": 0.74650195898721, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:12:35", "step": 10, "val_loss": 9.096094131469727, "perplexity": 8920.382728799992, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:12:53", "step": 10, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000010.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 20480}
+{"event": "train", "timestamp": "2026-03-13T14:12:55", "step": 11, "loss": 8.79585576057434, "lr": 0.0007750000000000001, "tok_per_sec": 106.21275007147676, "grad_norm": 1.4899623394012451, "tokens_seen": 22528, "elapsed_sec": 19.282054166018497, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:55", "step": 12, "loss": 8.558577299118042, "lr": 0.0006890576474687264, "tok_per_sec": 2617.334622430896, "grad_norm": 1.5879555940628052, "tokens_seen": 24576, "elapsed_sec": 0.7824754169560038, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:56", "step": 13, "loss": 8.595118045806885, "lr": 0.0005970378084704442, "tok_per_sec": 2577.6688468954194, "grad_norm": 1.3136155605316162, "tokens_seen": 26624, "elapsed_sec": 0.7945163330296054, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:57", "step": 14, "loss": 8.341074705123901, "lr": 0.000502962191529556, "tok_per_sec": 2652.6070721156225, "grad_norm": 1.2977045774459839, "tokens_seen": 28672, "elapsed_sec": 0.7720706249820068, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:58", "step": 15, "loss": 8.268006086349487, "lr": 0.0004109423525312737, "tok_per_sec": 2697.354489927494, "grad_norm": 1.1973044872283936, "tokens_seen": 30720, "elapsed_sec": 0.7592624579556286, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:58", "step": 16, "loss": 8.302424192428589, "lr": 0.0003250000000000001, "tok_per_sec": 2578.0945986776087, "grad_norm": 1.1259286403656006, "tokens_seen": 32768, "elapsed_sec": 0.794385124987457, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:12:59", "step": 17, "loss": 8.61441159248352, "lr": 0.00024889122713851394, "tok_per_sec": 2190.790230266318, "grad_norm": 0.9170812368392944, "tokens_seen": 34816, "elapsed_sec": 0.9348224999848753, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:13:00", "step": 18, "loss": 8.439870119094849, "lr": 0.00018594235253127368, "tok_per_sec": 2557.2677193427544, "grad_norm": 1.6393400430679321, "tokens_seen": 36864, "elapsed_sec": 0.8008547499775887, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:13:01", "step": 19, "loss": 7.915311813354492, "lr": 0.00013890454406082956, "tok_per_sec": 2489.488983600405, "grad_norm": 1.111694097518921, "tokens_seen": 38912, "elapsed_sec": 0.8226587920216843, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:13:02", "step": 20, "loss": 7.964773654937744, "lr": 0.00010983357966978745, "tok_per_sec": 2517.705047649911, "grad_norm": 1.00663423538208, "tokens_seen": 40960, "elapsed_sec": 0.8134392080246471, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:13:02", "step": 20, "val_loss": 8.758275032043457, "perplexity": 6363.125917448135, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:13:17", "step": 20, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000020.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 40960}
+{"event": "run_finished", "timestamp": "2026-03-13T14:13:17", "final_step": 20, "tokens_seen": 40960}

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_141224.log

@@ -0,0 +1,34 @@
+2026-03-13 14:12:24,605 | INFO | Pretraining started
+2026-03-13 14:12:24,605 | INFO | Log file: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_141224.log
+2026-03-13 14:12:24,605 | INFO | Metrics JSONL: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_141224.jsonl
+2026-03-13 14:12:24,605 | INFO | Arguments | model_config=configs/model_70m.json train_config=configs/pretrain_mps_dryrun.json max_steps_override=None
+2026-03-13 14:12:24,605 | INFO | Model config | {'vocab_size': 49152, 'max_seq_len': 8192, 'd_model': 384, 'n_layers': 32, 'n_heads': 6, 'ffn_hidden_dim': 1024, 'rope_theta': 10000.0, 'rms_norm_eps': 1e-05, 'initializer_range': 0.02, 'dropout': 0.0, 'tie_word_embeddings': True, 'bias': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2}
+2026-03-13 14:12:24,605 | INFO | Train config | {'seed': 42, 'train_dir': 'data/pretokenized/train', 'val_dir': 'data/pretokenized/val', 'output_dir': 'outputs/pretrain_mps_dryrun', 'checkpoint_dir': 'checkpoints/pretrain_mps_dryrun', 'init_from': None, 'resume_from': None, 'seq_len': 512, 'micro_batch_size': 1, 'grad_accum_steps': 4, 'max_steps': 20, 'warmup_steps': 5, 'learning_rate': 0.001, 'min_lr': 0.0001, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'precision': 'fp32', 'num_workers': 0, 'log_interval': 1, 'eval_interval': 10, 'eval_batches': 2, 'save_interval': 10, 'compile_model': False}
+2026-03-13 14:12:27,439 | INFO | Device summary | device=mps precision=fp32 compile_model=False
+2026-03-13 14:12:27,440 | INFO | Model summary | parameters=75.57M
+2026-03-13 14:12:27,440 | INFO | Batch summary | seq_len=512 micro_batch_size=1 grad_accum_steps=4 tokens_per_step=2,048
+2026-03-13 14:12:27,440 | INFO | Dataset summary | train_dir=data/pretokenized/train val_dir=data/pretokenized/val num_train_shards=100
+2026-03-13 14:12:28,743 | INFO | Train step | step=1 loss=10.8489 lr=0.000200 tok_per_sec=1,572 grad_norm=5.5737 tokens_seen=2.05K
+2026-03-13 14:12:29,792 | INFO | Train step | step=2 loss=10.7631 lr=0.000400 tok_per_sec=1,953 grad_norm=8.3225 tokens_seen=4.10K
+2026-03-13 14:12:30,538 | INFO | Train step | step=3 loss=10.3567 lr=0.000600 tok_per_sec=2,747 grad_norm=2.6284 tokens_seen=6.14K
+2026-03-13 14:12:31,280 | INFO | Train step | step=4 loss=10.3768 lr=0.000800 tok_per_sec=2,766 grad_norm=2.2171 tokens_seen=8.19K
+2026-03-13 14:12:32,020 | INFO | Train step | step=5 loss=10.2315 lr=0.001000 tok_per_sec=2,769 grad_norm=12.6454 tokens_seen=10.24K
+2026-03-13 14:12:32,760 | INFO | Train step | step=6 loss=9.9389 lr=0.001000 tok_per_sec=2,768 grad_norm=1.7283 tokens_seen=12.29K
+2026-03-13 14:12:33,505 | INFO | Train step | step=7 loss=9.6686 lr=0.000990 tok_per_sec=2,751 grad_norm=1.9499 tokens_seen=14.34K
+2026-03-13 14:12:34,247 | INFO | Train step | step=8 loss=9.3355 lr=0.000961 tok_per_sec=2,761 grad_norm=1.7211 tokens_seen=16.38K
+2026-03-13 14:12:34,986 | INFO | Train step | step=9 loss=9.2675 lr=0.000914 tok_per_sec=2,775 grad_norm=1.7628 tokens_seen=18.43K
+2026-03-13 14:12:35,733 | INFO | Train step | step=10 loss=8.8337 lr=0.000851 tok_per_sec=2,743 grad_norm=1.7004 tokens_seen=20.48K
+2026-03-13 14:12:35,903 | INFO | Eval step | step=10 val_loss=9.0961 perplexity=8920.38
+2026-03-13 14:12:53,990 | INFO | Checkpoint saved | step=10 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000010.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt
+2026-03-13 14:12:55,016 | INFO | Train step | step=11 loss=8.7959 lr=0.000775 tok_per_sec=106 grad_norm=1.4900 tokens_seen=22.53K
+2026-03-13 14:12:55,798 | INFO | Train step | step=12 loss=8.5586 lr=0.000689 tok_per_sec=2,617 grad_norm=1.5880 tokens_seen=24.58K
+2026-03-13 14:12:56,593 | INFO | Train step | step=13 loss=8.5951 lr=0.000597 tok_per_sec=2,578 grad_norm=1.3136 tokens_seen=26.62K
+2026-03-13 14:12:57,366 | INFO | Train step | step=14 loss=8.3411 lr=0.000503 tok_per_sec=2,653 grad_norm=1.2977 tokens_seen=28.67K
+2026-03-13 14:12:58,126 | INFO | Train step | step=15 loss=8.2680 lr=0.000411 tok_per_sec=2,697 grad_norm=1.1973 tokens_seen=30.72K
+2026-03-13 14:12:58,921 | INFO | Train step | step=16 loss=8.3024 lr=0.000325 tok_per_sec=2,578 grad_norm=1.1259 tokens_seen=32.77K
+2026-03-13 14:12:59,857 | INFO | Train step | step=17 loss=8.6144 lr=0.000249 tok_per_sec=2,191 grad_norm=0.9171 tokens_seen=34.82K
+2026-03-13 14:13:00,660 | INFO | Train step | step=18 loss=8.4399 lr=0.000186 tok_per_sec=2,557 grad_norm=1.6393 tokens_seen=36.86K
+2026-03-13 14:13:01,483 | INFO | Train step | step=19 loss=7.9153 lr=0.000139 tok_per_sec=2,489 grad_norm=1.1117 tokens_seen=38.91K
+2026-03-13 14:13:02,297 | INFO | Train step | step=20 loss=7.9648 lr=0.000110 tok_per_sec=2,518 grad_norm=1.0066 tokens_seen=40.96K
+2026-03-13 14:13:02,479 | INFO | Eval step | step=20 val_loss=8.7583 perplexity=6363.13
+2026-03-13 14:13:17,338 | INFO | Checkpoint saved | step=20 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000020.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142331.jsonl

@@ -0,0 +1,27 @@
+{"event": "run_started", "timestamp": "2026-03-13T14:23:31", "log_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142331.log", "metrics_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142331.jsonl", "model_config": {"vocab_size": 49152, "max_seq_len": 8192, "d_model": 384, "n_layers": 32, "n_heads": 6, "ffn_hidden_dim": 1024, "rope_theta": 10000.0, "rms_norm_eps": 1e-05, "initializer_range": 0.02, "dropout": 0.0, "tie_word_embeddings": true, "bias": false, "pad_token_id": 0, "bos_token_id": 1, "eos_token_id": 2}, "train_config": {"seed": 42, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val", "output_dir": "outputs/pretrain_mps_dryrun", "checkpoint_dir": "checkpoints/pretrain_mps_dryrun", "init_from": null, "resume_from": null, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "max_steps": 20, "warmup_steps": 5, "learning_rate": 0.001, "min_lr": 0.0001, "weight_decay": 0.1, "beta1": 0.9, "beta2": 0.95, "grad_clip": 1.0, "precision": "fp32", "num_workers": 0, "log_interval": 1, "eval_interval": 10, "eval_batches": 2, "save_interval": 10, "compile_model": false}, "args": {"model_config": "configs/model_70m.json", "train_config": "configs/pretrain_mps_dryrun.json", "max_steps_override": null}}
+{"event": "runtime_summary", "timestamp": "2026-03-13T14:23:34", "device": "mps", "precision": "fp32", "compile_model": false, "parameters": 75571584, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "tokens_per_step": 2048, "num_train_shards": 100, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val"}
+{"event": "train", "timestamp": "2026-03-13T14:23:36", "step": 1, "loss": 10.848917245864868, "lr": 0.0002, "tok_per_sec": 1528.4547833159693, "grad_norm": 5.573695659637451, "tokens_seen": 2048, "elapsed_sec": 1.3399153330246918, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:23:37", "step": 2, "loss": 10.763139486312866, "lr": 0.0004, "tok_per_sec": 1823.3959146519999, "grad_norm": 8.322466850280762, "tokens_seen": 4096, "elapsed_sec": 1.1231789999874309, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:23:37", "step": 3, "loss": 10.356749296188354, "lr": 0.0006000000000000001, "tok_per_sec": 2760.206551850419, "grad_norm": 2.6283912658691406, "tokens_seen": 6144, "elapsed_sec": 0.7419734579743817, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:23:38", "step": 4, "loss": 10.376826524734497, "lr": 0.0008, "tok_per_sec": 2758.5098402671138, "grad_norm": 2.217130184173584, "tokens_seen": 8192, "elapsed_sec": 0.7424298329860903, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:23:39", "step": 5, "loss": 10.231549263000488, "lr": 0.001, "tok_per_sec": 2743.529400521807, "grad_norm": 12.645355224609375, "tokens_seen": 10240, "elapsed_sec": 0.7464837080333382, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:23:40", "step": 6, "loss": 9.938905715942383, "lr": 0.001, "tok_per_sec": 2787.348798075642, "grad_norm": 1.7282867431640625, "tokens_seen": 12288, "elapsed_sec": 0.7347483750199899, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:23:40", "step": 7, "loss": 9.66855764389038, "lr": 0.0009901664203302125, "tok_per_sec": 2768.0064610428144, "grad_norm": 1.9499105215072632, "tokens_seen": 14336, "elapsed_sec": 0.7398826660355553, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:23:41", "step": 8, "loss": 9.335453271865845, "lr": 0.0009610954559391703, "tok_per_sec": 2765.3453249466884, "grad_norm": 1.7210659980773926, "tokens_seen": 16384, "elapsed_sec": 0.7405946669750847, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:23:42", "step": 9, "loss": 9.267512559890747, "lr": 0.0009140576474687263, "tok_per_sec": 2776.3259680860633, "grad_norm": 1.762829065322876, "tokens_seen": 18432, "elapsed_sec": 0.7376655419939198, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:23:43", "step": 10, "loss": 8.833673000335693, "lr": 0.0008511087728614862, "tok_per_sec": 2775.777209987709, "grad_norm": 1.700391173362732, "tokens_seen": 20480, "elapsed_sec": 0.7378113750019111, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:23:43", "step": 10, "val_loss": 9.096094131469727, "perplexity": 8920.382728799992, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:24:04", "step": 10, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000010.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 20480}
+{"event": "train", "timestamp": "2026-03-13T14:24:04", "step": 11, "loss": 8.795855522155762, "lr": 0.0007750000000000001, "tok_per_sec": 93.63913751097624, "grad_norm": 1.4899623394012451, "tokens_seen": 22528, "elapsed_sec": 21.871196749969386, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:24:05", "step": 12, "loss": 8.558577060699463, "lr": 0.0006890576474687264, "tok_per_sec": 2713.0807773384895, "grad_norm": 1.5879555940628052, "tokens_seen": 24576, "elapsed_sec": 0.7548614169936627, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:24:06", "step": 13, "loss": 8.595118045806885, "lr": 0.0005970378084704442, "tok_per_sec": 2623.969926754092, "grad_norm": 1.3136155605316162, "tokens_seen": 26624, "elapsed_sec": 0.7804967500269413, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:24:07", "step": 14, "loss": 8.341074705123901, "lr": 0.000502962191529556, "tok_per_sec": 2640.2210539544203, "grad_norm": 1.2977045774459839, "tokens_seen": 28672, "elapsed_sec": 0.7756926250294782, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:24:08", "step": 15, "loss": 8.268006086349487, "lr": 0.0004109423525312737, "tok_per_sec": 2716.846221941928, "grad_norm": 1.1973044872283936, "tokens_seen": 30720, "elapsed_sec": 0.7538152080378495, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:24:08", "step": 16, "loss": 8.302424192428589, "lr": 0.0003250000000000001, "tok_per_sec": 2732.9290763944346, "grad_norm": 1.1259286403656006, "tokens_seen": 32768, "elapsed_sec": 0.7493791250162758, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:24:09", "step": 17, "loss": 8.6144118309021, "lr": 0.00024889122713851394, "tok_per_sec": 2724.595059218572, "grad_norm": 0.9170812368392944, "tokens_seen": 34816, "elapsed_sec": 0.7516713329823688, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:24:10", "step": 18, "loss": 8.439870119094849, "lr": 0.00018594235253127368, "tok_per_sec": 2735.215026577713, "grad_norm": 1.6393400430679321, "tokens_seen": 36864, "elapsed_sec": 0.7487528329947963, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:24:11", "step": 19, "loss": 7.915311813354492, "lr": 0.00013890454406082956, "tok_per_sec": 2730.9775428733506, "grad_norm": 1.111694097518921, "tokens_seen": 38912, "elapsed_sec": 0.7499146250193007, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:24:11", "step": 20, "loss": 7.964773654937744, "lr": 0.00010983357966978745, "tok_per_sec": 2738.180386915235, "grad_norm": 1.00663423538208, "tokens_seen": 40960, "elapsed_sec": 0.7479419580195099, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:24:11", "step": 20, "val_loss": 8.758275032043457, "perplexity": 6363.125917448135, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:24:24", "step": 20, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000020.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 40960}
+{"event": "run_finished", "timestamp": "2026-03-13T14:24:24", "final_step": 20, "tokens_seen": 40960}

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142331.log

@@ -0,0 +1,34 @@
+2026-03-13 14:23:31,892 | INFO | Pretraining started
+2026-03-13 14:23:31,892 | INFO | Log file: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142331.log
+2026-03-13 14:23:31,892 | INFO | Metrics JSONL: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142331.jsonl
+2026-03-13 14:23:31,892 | INFO | Arguments | model_config=configs/model_70m.json train_config=configs/pretrain_mps_dryrun.json max_steps_override=None
+2026-03-13 14:23:31,892 | INFO | Model config | {'vocab_size': 49152, 'max_seq_len': 8192, 'd_model': 384, 'n_layers': 32, 'n_heads': 6, 'ffn_hidden_dim': 1024, 'rope_theta': 10000.0, 'rms_norm_eps': 1e-05, 'initializer_range': 0.02, 'dropout': 0.0, 'tie_word_embeddings': True, 'bias': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2}
+2026-03-13 14:23:31,892 | INFO | Train config | {'seed': 42, 'train_dir': 'data/pretokenized/train', 'val_dir': 'data/pretokenized/val', 'output_dir': 'outputs/pretrain_mps_dryrun', 'checkpoint_dir': 'checkpoints/pretrain_mps_dryrun', 'init_from': None, 'resume_from': None, 'seq_len': 512, 'micro_batch_size': 1, 'grad_accum_steps': 4, 'max_steps': 20, 'warmup_steps': 5, 'learning_rate': 0.001, 'min_lr': 0.0001, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'precision': 'fp32', 'num_workers': 0, 'log_interval': 1, 'eval_interval': 10, 'eval_batches': 2, 'save_interval': 10, 'compile_model': False}
+2026-03-13 14:23:34,726 | INFO | Device summary | device=mps precision=fp32 compile_model=False
+2026-03-13 14:23:34,727 | INFO | Model summary | parameters=75.57M
+2026-03-13 14:23:34,727 | INFO | Batch summary | seq_len=512 micro_batch_size=1 grad_accum_steps=4 tokens_per_step=2,048
+2026-03-13 14:23:34,727 | INFO | Dataset summary | train_dir=data/pretokenized/train val_dir=data/pretokenized/val num_train_shards=100
+2026-03-13 14:23:36,068 | INFO | Train step | step=1 loss=10.8489 lr=0.000200 tok_per_sec=1,528 grad_norm=5.5737 tokens_seen=2.05K
+2026-03-13 14:23:37,192 | INFO | Train step | step=2 loss=10.7631 lr=0.000400 tok_per_sec=1,823 grad_norm=8.3225 tokens_seen=4.10K
+2026-03-13 14:23:37,934 | INFO | Train step | step=3 loss=10.3567 lr=0.000600 tok_per_sec=2,760 grad_norm=2.6284 tokens_seen=6.14K
+2026-03-13 14:23:38,678 | INFO | Train step | step=4 loss=10.3768 lr=0.000800 tok_per_sec=2,759 grad_norm=2.2171 tokens_seen=8.19K
+2026-03-13 14:23:39,425 | INFO | Train step | step=5 loss=10.2315 lr=0.001000 tok_per_sec=2,744 grad_norm=12.6454 tokens_seen=10.24K
+2026-03-13 14:23:40,160 | INFO | Train step | step=6 loss=9.9389 lr=0.001000 tok_per_sec=2,787 grad_norm=1.7283 tokens_seen=12.29K
+2026-03-13 14:23:40,900 | INFO | Train step | step=7 loss=9.6686 lr=0.000990 tok_per_sec=2,768 grad_norm=1.9499 tokens_seen=14.34K
+2026-03-13 14:23:41,641 | INFO | Train step | step=8 loss=9.3355 lr=0.000961 tok_per_sec=2,765 grad_norm=1.7211 tokens_seen=16.38K
+2026-03-13 14:23:42,380 | INFO | Train step | step=9 loss=9.2675 lr=0.000914 tok_per_sec=2,776 grad_norm=1.7628 tokens_seen=18.43K
+2026-03-13 14:23:43,118 | INFO | Train step | step=10 loss=8.8337 lr=0.000851 tok_per_sec=2,776 grad_norm=1.7004 tokens_seen=20.48K
+2026-03-13 14:23:43,296 | INFO | Eval step | step=10 val_loss=9.0961 perplexity=8920.38
+2026-03-13 14:24:04,120 | INFO | Checkpoint saved | step=10 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000010.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt
+2026-03-13 14:24:04,990 | INFO | Train step | step=11 loss=8.7959 lr=0.000775 tok_per_sec=94 grad_norm=1.4900 tokens_seen=22.53K
+2026-03-13 14:24:05,745 | INFO | Train step | step=12 loss=8.5586 lr=0.000689 tok_per_sec=2,713 grad_norm=1.5880 tokens_seen=24.58K
+2026-03-13 14:24:06,526 | INFO | Train step | step=13 loss=8.5951 lr=0.000597 tok_per_sec=2,624 grad_norm=1.3136 tokens_seen=26.62K
+2026-03-13 14:24:07,302 | INFO | Train step | step=14 loss=8.3411 lr=0.000503 tok_per_sec=2,640 grad_norm=1.2977 tokens_seen=28.67K
+2026-03-13 14:24:08,057 | INFO | Train step | step=15 loss=8.2680 lr=0.000411 tok_per_sec=2,717 grad_norm=1.1973 tokens_seen=30.72K
+2026-03-13 14:24:08,806 | INFO | Train step | step=16 loss=8.3024 lr=0.000325 tok_per_sec=2,733 grad_norm=1.1259 tokens_seen=32.77K
+2026-03-13 14:24:09,559 | INFO | Train step | step=17 loss=8.6144 lr=0.000249 tok_per_sec=2,725 grad_norm=0.9171 tokens_seen=34.82K
+2026-03-13 14:24:10,308 | INFO | Train step | step=18 loss=8.4399 lr=0.000186 tok_per_sec=2,735 grad_norm=1.6393 tokens_seen=36.86K
+2026-03-13 14:24:11,058 | INFO | Train step | step=19 loss=7.9153 lr=0.000139 tok_per_sec=2,731 grad_norm=1.1117 tokens_seen=38.91K
+2026-03-13 14:24:11,807 | INFO | Train step | step=20 loss=7.9648 lr=0.000110 tok_per_sec=2,738 grad_norm=1.0066 tokens_seen=40.96K
+2026-03-13 14:24:11,966 | INFO | Eval step | step=20 val_loss=8.7583 perplexity=6363.13
+2026-03-13 14:24:24,399 | INFO | Checkpoint saved | step=20 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000020.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142530.jsonl

@@ -0,0 +1,13 @@
+{"event": "run_started", "timestamp": "2026-03-13T14:25:30", "log_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142530.log", "metrics_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142530.jsonl", "model_config": {"vocab_size": 49152, "max_seq_len": 8192, "d_model": 384, "n_layers": 32, "n_heads": 6, "ffn_hidden_dim": 1024, "rope_theta": 10000.0, "rms_norm_eps": 1e-05, "initializer_range": 0.02, "dropout": 0.0, "tie_word_embeddings": true, "bias": false, "pad_token_id": 0, "bos_token_id": 1, "eos_token_id": 2}, "train_config": {"seed": 42, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val", "output_dir": "outputs/pretrain_mps_dryrun", "checkpoint_dir": "checkpoints/pretrain_mps_dryrun", "init_from": null, "resume_from": null, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "max_steps": 20, "warmup_steps": 5, "learning_rate": 0.001, "min_lr": 0.0001, "weight_decay": 0.1, "beta1": 0.9, "beta2": 0.95, "grad_clip": 1.0, "precision": "fp32", "num_workers": 0, "log_interval": 1, "eval_interval": 10, "eval_batches": 2, "save_interval": 10, "compile_model": false}, "args": {"model_config": "configs/model_70m.json", "train_config": "configs/pretrain_mps_dryrun.json", "max_steps_override": null}}
+{"event": "runtime_summary", "timestamp": "2026-03-13T14:25:33", "device": "mps", "precision": "fp32", "compile_model": false, "parameters": 75571584, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "tokens_per_step": 2048, "num_train_shards": 100, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val"}
+{"event": "train", "timestamp": "2026-03-13T14:25:35", "step": 1, "loss": 10.848917245864868, "lr": 0.0002, "tok_per_sec": 1738.338776128348, "grad_norm": 5.573695659637451, "tokens_seen": 2048, "elapsed_sec": 1.1781362920301035, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:25:35", "step": 2, "loss": 10.763139486312866, "lr": 0.0004, "tok_per_sec": 2202.432846935467, "grad_norm": 8.322466850280762, "tokens_seen": 4096, "elapsed_sec": 0.9298807919840328, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:25:36", "step": 3, "loss": 10.356749296188354, "lr": 0.0006000000000000001, "tok_per_sec": 2756.6912099758943, "grad_norm": 2.6283912658691406, "tokens_seen": 6144, "elapsed_sec": 0.7429196250159293, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:25:37", "step": 4, "loss": 10.376826524734497, "lr": 0.0008, "tok_per_sec": 2738.8579811402797, "grad_norm": 2.217130184173584, "tokens_seen": 8192, "elapsed_sec": 0.7477569169714116, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:25:38", "step": 5, "loss": 10.231549263000488, "lr": 0.001, "tok_per_sec": 2710.2101486388783, "grad_norm": 12.645347595214844, "tokens_seen": 10240, "elapsed_sec": 0.755660958995577, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:25:39", "step": 6, "loss": 9.938905477523804, "lr": 0.001, "tok_per_sec": 2643.8368942648644, "grad_norm": 1.7282862663269043, "tokens_seen": 12288, "elapsed_sec": 0.7746317499550059, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:25:39", "step": 7, "loss": 9.66855764389038, "lr": 0.0009901664203302125, "tok_per_sec": 2696.598737188916, "grad_norm": 1.9499105215072632, "tokens_seen": 14336, "elapsed_sec": 0.7594752499717288, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:25:40", "step": 8, "loss": 9.335453271865845, "lr": 0.0009610954559391703, "tok_per_sec": 2697.2926156977223, "grad_norm": 1.7210659980773926, "tokens_seen": 16384, "elapsed_sec": 0.7592798749683425, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:25:41", "step": 9, "loss": 9.267513036727905, "lr": 0.0009140576474687263, "tok_per_sec": 2700.6039931851096, "grad_norm": 1.762829065322876, "tokens_seen": 18432, "elapsed_sec": 0.7583488749805838, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:25:42", "step": 10, "loss": 8.833672761917114, "lr": 0.0008511087728614862, "tok_per_sec": 2564.5051143205906, "grad_norm": 1.700391173362732, "tokens_seen": 20480, "elapsed_sec": 0.7985946249682456, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:25:42", "step": 10, "val_loss": 9.096094131469727, "perplexity": 8920.382728799992, "eval_batches": 2}

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142530.log

@@ -0,0 +1,21 @@
+2026-03-13 14:25:30,936 | INFO | Pretraining started
+2026-03-13 14:25:30,936 | INFO | Log file: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142530.log
+2026-03-13 14:25:30,936 | INFO | Metrics JSONL: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142530.jsonl
+2026-03-13 14:25:30,936 | INFO | Arguments | model_config=configs/model_70m.json train_config=configs/pretrain_mps_dryrun.json max_steps_override=None
+2026-03-13 14:25:30,936 | INFO | Model config | {'vocab_size': 49152, 'max_seq_len': 8192, 'd_model': 384, 'n_layers': 32, 'n_heads': 6, 'ffn_hidden_dim': 1024, 'rope_theta': 10000.0, 'rms_norm_eps': 1e-05, 'initializer_range': 0.02, 'dropout': 0.0, 'tie_word_embeddings': True, 'bias': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2}
+2026-03-13 14:25:30,937 | INFO | Train config | {'seed': 42, 'train_dir': 'data/pretokenized/train', 'val_dir': 'data/pretokenized/val', 'output_dir': 'outputs/pretrain_mps_dryrun', 'checkpoint_dir': 'checkpoints/pretrain_mps_dryrun', 'init_from': None, 'resume_from': None, 'seq_len': 512, 'micro_batch_size': 1, 'grad_accum_steps': 4, 'max_steps': 20, 'warmup_steps': 5, 'learning_rate': 0.001, 'min_lr': 0.0001, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'precision': 'fp32', 'num_workers': 0, 'log_interval': 1, 'eval_interval': 10, 'eval_batches': 2, 'save_interval': 10, 'compile_model': False}
+2026-03-13 14:25:33,870 | INFO | Device summary | device=mps precision=fp32 compile_model=False
+2026-03-13 14:25:33,870 | INFO | Model summary | parameters=75.57M
+2026-03-13 14:25:33,871 | INFO | Batch summary | seq_len=512 micro_batch_size=1 grad_accum_steps=4 tokens_per_step=2,048
+2026-03-13 14:25:33,871 | INFO | Dataset summary | train_dir=data/pretokenized/train val_dir=data/pretokenized/val num_train_shards=100
+2026-03-13 14:25:35,050 | INFO | Train step | step=1 loss=10.8489 lr=0.000200 tok_per_sec=1,738 grad_norm=5.5737 tokens_seen=2.05K
+2026-03-13 14:25:35,980 | INFO | Train step | step=2 loss=10.7631 lr=0.000400 tok_per_sec=2,202 grad_norm=8.3225 tokens_seen=4.10K
+2026-03-13 14:25:36,724 | INFO | Train step | step=3 loss=10.3567 lr=0.000600 tok_per_sec=2,757 grad_norm=2.6284 tokens_seen=6.14K
+2026-03-13 14:25:37,472 | INFO | Train step | step=4 loss=10.3768 lr=0.000800 tok_per_sec=2,739 grad_norm=2.2171 tokens_seen=8.19K
+2026-03-13 14:25:38,228 | INFO | Train step | step=5 loss=10.2315 lr=0.001000 tok_per_sec=2,710 grad_norm=12.6453 tokens_seen=10.24K
+2026-03-13 14:25:39,004 | INFO | Train step | step=6 loss=9.9389 lr=0.001000 tok_per_sec=2,644 grad_norm=1.7283 tokens_seen=12.29K
+2026-03-13 14:25:39,764 | INFO | Train step | step=7 loss=9.6686 lr=0.000990 tok_per_sec=2,697 grad_norm=1.9499 tokens_seen=14.34K
+2026-03-13 14:25:40,524 | INFO | Train step | step=8 loss=9.3355 lr=0.000961 tok_per_sec=2,697 grad_norm=1.7211 tokens_seen=16.38K
+2026-03-13 14:25:41,283 | INFO | Train step | step=9 loss=9.2675 lr=0.000914 tok_per_sec=2,701 grad_norm=1.7628 tokens_seen=18.43K
+2026-03-13 14:25:42,082 | INFO | Train step | step=10 loss=8.8337 lr=0.000851 tok_per_sec=2,565 grad_norm=1.7004 tokens_seen=20.48K
+2026-03-13 14:25:42,254 | INFO | Eval step | step=10 val_loss=9.0961 perplexity=8920.38

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142559.jsonl

@@ -0,0 +1,61 @@
+{"event": "run_started", "timestamp": "2026-03-13T14:25:59", "log_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142559.log", "metrics_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142559.jsonl", "model_config": {"vocab_size": 49152, "max_seq_len": 8192, "d_model": 384, "n_layers": 32, "n_heads": 6, "ffn_hidden_dim": 1024, "rope_theta": 10000.0, "rms_norm_eps": 1e-05, "initializer_range": 0.02, "dropout": 0.0, "tie_word_embeddings": true, "bias": false, "pad_token_id": 0, "bos_token_id": 1, "eos_token_id": 2}, "train_config": {"seed": 42, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val", "output_dir": "outputs/pretrain_mps_dryrun", "checkpoint_dir": "checkpoints/pretrain_mps_dryrun", "init_from": null, "resume_from": null, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "max_steps": 500, "warmup_steps": 50, "learning_rate": 0.001, "min_lr": 0.0001, "weight_decay": 0.1, "beta1": 0.9, "beta2": 0.95, "grad_clip": 1.0, "precision": "fp32", "num_workers": 0, "log_interval": 1, "eval_interval": 10, "eval_batches": 2, "save_interval": 10, "compile_model": false}, "args": {"model_config": "configs/model_70m.json", "train_config": "configs/pretrain_mps_dryrun.json", "max_steps_override": null}}
+{"event": "runtime_summary", "timestamp": "2026-03-13T14:26:03", "device": "mps", "precision": "fp32", "compile_model": false, "parameters": 75571584, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "tokens_per_step": 2048, "num_train_shards": 100, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val"}
+{"event": "train", "timestamp": "2026-03-13T14:26:04", "step": 1, "loss": 10.848917245864868, "lr": 2e-05, "tok_per_sec": 1704.8695679494026, "grad_norm": 5.573695659637451, "tokens_seen": 2048, "elapsed_sec": 1.2012649169773795, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:05", "step": 2, "loss": 10.897652626037598, "lr": 4e-05, "tok_per_sec": 2132.5760089557198, "grad_norm": 5.0279011726379395, "tokens_seen": 4096, "elapsed_sec": 0.9603409169940278, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:05", "step": 3, "loss": 10.785077571868896, "lr": 6e-05, "tok_per_sec": 2721.481065658625, "grad_norm": 5.114167213439941, "tokens_seen": 6144, "elapsed_sec": 0.7525314160156995, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:06", "step": 4, "loss": 10.634832620620728, "lr": 8e-05, "tok_per_sec": 2755.0287727887576, "grad_norm": 6.422860622406006, "tokens_seen": 8192, "elapsed_sec": 0.7433679169625975, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:07", "step": 5, "loss": 10.747493743896484, "lr": 0.0001, "tok_per_sec": 2752.2115918295694, "grad_norm": 6.580272197723389, "tokens_seen": 10240, "elapsed_sec": 0.7441288330010138, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:08", "step": 6, "loss": 10.633646488189697, "lr": 0.00012, "tok_per_sec": 2738.0719353969716, "grad_norm": 6.525123119354248, "tokens_seen": 12288, "elapsed_sec": 0.7479715830413625, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:08", "step": 7, "loss": 10.477944374084473, "lr": 0.00014000000000000001, "tok_per_sec": 2749.743315053932, "grad_norm": 5.189582824707031, "tokens_seen": 14336, "elapsed_sec": 0.7447967920452356, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:09", "step": 8, "loss": 10.35365605354309, "lr": 0.00016, "tok_per_sec": 2753.8837425710026, "grad_norm": 2.357203960418701, "tokens_seen": 16384, "elapsed_sec": 0.743676999991294, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:10", "step": 9, "loss": 10.339627742767334, "lr": 0.00018, "tok_per_sec": 2752.9484182072647, "grad_norm": 4.98753547668457, "tokens_seen": 18432, "elapsed_sec": 0.7439296669908799, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:11", "step": 10, "loss": 10.207262754440308, "lr": 0.0002, "tok_per_sec": 2755.2380338330304, "grad_norm": 5.554019927978516, "tokens_seen": 20480, "elapsed_sec": 0.7433114579762332, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:26:11", "step": 10, "val_loss": 10.418237686157227, "perplexity": 33464.40736092908, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:26:27", "step": 10, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000010.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 20480}
+{"event": "train", "timestamp": "2026-03-13T14:26:28", "step": 11, "loss": 10.161789417266846, "lr": 0.00021999999999999998, "tok_per_sec": 119.14029052794488, "grad_norm": 2.177887201309204, "tokens_seen": 22528, "elapsed_sec": 17.189818750019185, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:29", "step": 12, "loss": 10.047882318496704, "lr": 0.00024, "tok_per_sec": 2680.171783579867, "grad_norm": 1.9737660884857178, "tokens_seen": 24576, "elapsed_sec": 0.7641301249968819, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:29", "step": 13, "loss": 10.105542421340942, "lr": 0.00026000000000000003, "tok_per_sec": 2702.3713362104313, "grad_norm": 4.096495151519775, "tokens_seen": 26624, "elapsed_sec": 0.7578529170132242, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:30", "step": 14, "loss": 9.940982818603516, "lr": 0.00028000000000000003, "tok_per_sec": 2722.8033630518785, "grad_norm": 1.8798285722732544, "tokens_seen": 28672, "elapsed_sec": 0.7521659579942934, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:31", "step": 15, "loss": 9.852107524871826, "lr": 0.0003, "tok_per_sec": 2738.483668417333, "grad_norm": 3.7492053508758545, "tokens_seen": 30720, "elapsed_sec": 0.7478591249673627, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:32", "step": 16, "loss": 9.768622398376465, "lr": 0.00032, "tok_per_sec": 2756.126999840051, "grad_norm": 1.8649290800094604, "tokens_seen": 32768, "elapsed_sec": 0.7430717090028338, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:32", "step": 17, "loss": 9.880046606063843, "lr": 0.00034, "tok_per_sec": 2762.6129070979177, "grad_norm": 1.8526010513305664, "tokens_seen": 34816, "elapsed_sec": 0.7413271670229733, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:33", "step": 18, "loss": 9.652766704559326, "lr": 0.00036, "tok_per_sec": 2762.0740561225734, "grad_norm": 2.2104318141937256, "tokens_seen": 36864, "elapsed_sec": 0.74147179198917, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:34", "step": 19, "loss": 9.371065139770508, "lr": 0.00037999999999999997, "tok_per_sec": 2746.3599790187623, "grad_norm": 2.0031697750091553, "tokens_seen": 38912, "elapsed_sec": 0.7457143330248073, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:35", "step": 20, "loss": 9.290096044540405, "lr": 0.0004, "tok_per_sec": 2753.183730427524, "grad_norm": 1.9113200902938843, "tokens_seen": 40960, "elapsed_sec": 0.7438660839688964, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:26:35", "step": 20, "val_loss": 9.618017196655273, "perplexity": 15033.212463302863, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:26:47", "step": 20, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000020.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 40960}
+{"event": "train", "timestamp": "2026-03-13T14:26:48", "step": 21, "loss": 9.240976810455322, "lr": 0.00042, "tok_per_sec": 152.54166876251324, "grad_norm": 1.8867971897125244, "tokens_seen": 43008, "elapsed_sec": 13.425839749979787, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:49", "step": 22, "loss": 9.12671947479248, "lr": 0.00043999999999999996, "tok_per_sec": 2712.778906376721, "grad_norm": 1.8414427042007446, "tokens_seen": 45056, "elapsed_sec": 0.7549454160034657, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:50", "step": 23, "loss": 9.102352619171143, "lr": 0.00046, "tok_per_sec": 2708.941555033957, "grad_norm": 1.6314030885696411, "tokens_seen": 47104, "elapsed_sec": 0.7560148339834996, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:50", "step": 24, "loss": 8.811159133911133, "lr": 0.00048, "tok_per_sec": 2721.638538039314, "grad_norm": 1.8162541389465332, "tokens_seen": 49152, "elapsed_sec": 0.7524878749973141, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:51", "step": 25, "loss": 8.63177752494812, "lr": 0.0005, "tok_per_sec": 2726.8041484896708, "grad_norm": 1.7629377841949463, "tokens_seen": 51200, "elapsed_sec": 0.7510623750276864, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:52", "step": 26, "loss": 8.722702026367188, "lr": 0.0005200000000000001, "tok_per_sec": 2722.9481650192897, "grad_norm": 1.66167414188385, "tokens_seen": 53248, "elapsed_sec": 0.7521259590284899, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:53", "step": 27, "loss": 8.609044075012207, "lr": 0.00054, "tok_per_sec": 2753.0232007676027, "grad_norm": 1.4688063859939575, "tokens_seen": 55296, "elapsed_sec": 0.7439094590372406, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:53", "step": 28, "loss": 8.890318632125854, "lr": 0.0005600000000000001, "tok_per_sec": 2748.212475197479, "grad_norm": 1.3299572467803955, "tokens_seen": 57344, "elapsed_sec": 0.745211667031981, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:54", "step": 29, "loss": 8.266993045806885, "lr": 0.00058, "tok_per_sec": 2729.828001611049, "grad_norm": 1.4132530689239502, "tokens_seen": 59392, "elapsed_sec": 0.7502304170047864, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:26:55", "step": 30, "loss": 8.5481858253479, "lr": 0.0006, "tok_per_sec": 2732.970411338659, "grad_norm": 1.6723191738128662, "tokens_seen": 61440, "elapsed_sec": 0.7493677909951657, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:26:55", "step": 30, "val_loss": 8.843989849090576, "perplexity": 6932.597552057813, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:27:06", "step": 30, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000030.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 61440}
+{"event": "train", "timestamp": "2026-03-13T14:27:07", "step": 31, "loss": 8.099784016609192, "lr": 0.00062, "tok_per_sec": 168.2314308371399, "grad_norm": 1.3632475137710571, "tokens_seen": 63488, "elapsed_sec": 12.173706125002354, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:08", "step": 32, "loss": 8.278108835220337, "lr": 0.00064, "tok_per_sec": 2688.7908394379037, "grad_norm": 1.1354058980941772, "tokens_seen": 65536, "elapsed_sec": 0.7616806669975631, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:08", "step": 33, "loss": 8.04857063293457, "lr": 0.00066, "tok_per_sec": 2695.7795378460673, "grad_norm": 0.8978219032287598, "tokens_seen": 67584, "elapsed_sec": 0.7597060409607366, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:09", "step": 34, "loss": 8.70958948135376, "lr": 0.00068, "tok_per_sec": 2707.39823225129, "grad_norm": 1.751259446144104, "tokens_seen": 69632, "elapsed_sec": 0.756445791979786, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:10", "step": 35, "loss": 8.077706575393677, "lr": 0.0007000000000000001, "tok_per_sec": 2689.526034181471, "grad_norm": 0.9328188896179199, "tokens_seen": 71680, "elapsed_sec": 0.761472457961645, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:11", "step": 36, "loss": 8.057007431983948, "lr": 0.00072, "tok_per_sec": 2729.9771946058904, "grad_norm": 0.7004730701446533, "tokens_seen": 73728, "elapsed_sec": 0.7501894169836305, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:11", "step": 37, "loss": 8.280940413475037, "lr": 0.00074, "tok_per_sec": 2712.4312913681806, "grad_norm": 0.8498008251190186, "tokens_seen": 75776, "elapsed_sec": 0.7550421669729985, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:12", "step": 38, "loss": 8.420085191726685, "lr": 0.0007599999999999999, "tok_per_sec": 2728.334548153266, "grad_norm": 0.9405263662338257, "tokens_seen": 77824, "elapsed_sec": 0.7506410829955712, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:13", "step": 39, "loss": 8.040002822875977, "lr": 0.00078, "tok_per_sec": 2734.281540014069, "grad_norm": 0.8642140030860901, "tokens_seen": 79872, "elapsed_sec": 0.7490084579912946, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:14", "step": 40, "loss": 8.193370580673218, "lr": 0.0008, "tok_per_sec": 2746.819492491367, "grad_norm": 0.9126524329185486, "tokens_seen": 81920, "elapsed_sec": 0.7455895830062218, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:27:14", "step": 40, "val_loss": 8.36504077911377, "perplexity": 4294.2868516794715, "eval_batches": 2}
+{"event": "checkpoint", "timestamp": "2026-03-13T14:27:26", "step": 40, "step_checkpoint": "checkpoints/pretrain_mps_dryrun/step_0000040.pt", "last_checkpoint": "checkpoints/pretrain_mps_dryrun/last.pt", "tokens_seen": 81920}
+{"event": "train", "timestamp": "2026-03-13T14:27:27", "step": 41, "loss": 7.95119035243988, "lr": 0.00082, "tok_per_sec": 156.35387795373205, "grad_norm": 1.0787891149520874, "tokens_seen": 83968, "elapsed_sec": 13.098491874989122, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:28", "step": 42, "loss": 7.986739635467529, "lr": 0.00084, "tok_per_sec": 2722.8155804265352, "grad_norm": 0.9517979621887207, "tokens_seen": 86016, "elapsed_sec": 0.7521625829976983, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:28", "step": 43, "loss": 7.984379172325134, "lr": 0.0008600000000000001, "tok_per_sec": 2734.609061278845, "grad_norm": 0.9767814874649048, "tokens_seen": 88064, "elapsed_sec": 0.7489187500323169, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:29", "step": 44, "loss": 7.784951090812683, "lr": 0.0008799999999999999, "tok_per_sec": 2551.6733064235736, "grad_norm": 0.9357463717460632, "tokens_seen": 90112, "elapsed_sec": 0.8026105829630978, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:30", "step": 45, "loss": 8.117401361465454, "lr": 0.0009, "tok_per_sec": 2731.190447061073, "grad_norm": 0.7716737985610962, "tokens_seen": 92160, "elapsed_sec": 0.7498561670072377, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:31", "step": 46, "loss": 8.220598220825195, "lr": 0.00092, "tok_per_sec": 2729.7315806190586, "grad_norm": 0.7731218338012695, "tokens_seen": 94208, "elapsed_sec": 0.7502569170319475, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:31", "step": 47, "loss": 8.069997072219849, "lr": 0.00094, "tok_per_sec": 2718.9204771868167, "grad_norm": 1.0135213136672974, "tokens_seen": 96256, "elapsed_sec": 0.7532401249627583, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:32", "step": 48, "loss": 7.909337043762207, "lr": 0.00096, "tok_per_sec": 2717.9027900424403, "grad_norm": 1.088800072669983, "tokens_seen": 98304, "elapsed_sec": 0.7535221669822931, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:33", "step": 49, "loss": 7.95028281211853, "lr": 0.00098, "tok_per_sec": 2721.736196494166, "grad_norm": 1.794154167175293, "tokens_seen": 100352, "elapsed_sec": 0.7524608750245534, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:27:34", "step": 50, "loss": 8.580274820327759, "lr": 0.001, "tok_per_sec": 2714.6346388455563, "grad_norm": 1.2600723505020142, "tokens_seen": 102400, "elapsed_sec": 0.7544293330283836, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:27:34", "step": 50, "val_loss": 8.607748031616211, "perplexity": 5473.907720171149, "eval_batches": 2}

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142559.log

@@ -0,0 +1,69 @@
+2026-03-13 14:25:59,997 | INFO | Pretraining started
+2026-03-13 14:25:59,997 | INFO | Log file: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142559.log
+2026-03-13 14:25:59,997 | INFO | Metrics JSONL: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_142559.jsonl
+2026-03-13 14:25:59,997 | INFO | Arguments | model_config=configs/model_70m.json train_config=configs/pretrain_mps_dryrun.json max_steps_override=None
+2026-03-13 14:25:59,997 | INFO | Model config | {'vocab_size': 49152, 'max_seq_len': 8192, 'd_model': 384, 'n_layers': 32, 'n_heads': 6, 'ffn_hidden_dim': 1024, 'rope_theta': 10000.0, 'rms_norm_eps': 1e-05, 'initializer_range': 0.02, 'dropout': 0.0, 'tie_word_embeddings': True, 'bias': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2}
+2026-03-13 14:25:59,997 | INFO | Train config | {'seed': 42, 'train_dir': 'data/pretokenized/train', 'val_dir': 'data/pretokenized/val', 'output_dir': 'outputs/pretrain_mps_dryrun', 'checkpoint_dir': 'checkpoints/pretrain_mps_dryrun', 'init_from': None, 'resume_from': None, 'seq_len': 512, 'micro_batch_size': 1, 'grad_accum_steps': 4, 'max_steps': 500, 'warmup_steps': 50, 'learning_rate': 0.001, 'min_lr': 0.0001, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'precision': 'fp32', 'num_workers': 0, 'log_interval': 1, 'eval_interval': 10, 'eval_batches': 2, 'save_interval': 10, 'compile_model': False}
+2026-03-13 14:26:03,019 | INFO | Device summary | device=mps precision=fp32 compile_model=False
+2026-03-13 14:26:03,020 | INFO | Model summary | parameters=75.57M
+2026-03-13 14:26:03,020 | INFO | Batch summary | seq_len=512 micro_batch_size=1 grad_accum_steps=4 tokens_per_step=2,048
+2026-03-13 14:26:03,020 | INFO | Dataset summary | train_dir=data/pretokenized/train val_dir=data/pretokenized/val num_train_shards=100
+2026-03-13 14:26:04,222 | INFO | Train step | step=1 loss=10.8489 lr=0.000020 tok_per_sec=1,705 grad_norm=5.5737 tokens_seen=2.05K
+2026-03-13 14:26:05,183 | INFO | Train step | step=2 loss=10.8977 lr=0.000040 tok_per_sec=2,133 grad_norm=5.0279 tokens_seen=4.10K
+2026-03-13 14:26:05,936 | INFO | Train step | step=3 loss=10.7851 lr=0.000060 tok_per_sec=2,721 grad_norm=5.1142 tokens_seen=6.14K
+2026-03-13 14:26:06,680 | INFO | Train step | step=4 loss=10.6348 lr=0.000080 tok_per_sec=2,755 grad_norm=6.4229 tokens_seen=8.19K
+2026-03-13 14:26:07,424 | INFO | Train step | step=5 loss=10.7475 lr=0.000100 tok_per_sec=2,752 grad_norm=6.5803 tokens_seen=10.24K
+2026-03-13 14:26:08,173 | INFO | Train step | step=6 loss=10.6336 lr=0.000120 tok_per_sec=2,738 grad_norm=6.5251 tokens_seen=12.29K
+2026-03-13 14:26:08,918 | INFO | Train step | step=7 loss=10.4779 lr=0.000140 tok_per_sec=2,750 grad_norm=5.1896 tokens_seen=14.34K
+2026-03-13 14:26:09,663 | INFO | Train step | step=8 loss=10.3537 lr=0.000160 tok_per_sec=2,754 grad_norm=2.3572 tokens_seen=16.38K
+2026-03-13 14:26:10,408 | INFO | Train step | step=9 loss=10.3396 lr=0.000180 tok_per_sec=2,753 grad_norm=4.9875 tokens_seen=18.43K
+2026-03-13 14:26:11,152 | INFO | Train step | step=10 loss=10.2073 lr=0.000200 tok_per_sec=2,755 grad_norm=5.5540 tokens_seen=20.48K
+2026-03-13 14:26:11,326 | INFO | Eval step | step=10 val_loss=10.4182 perplexity=33464.41
+2026-03-13 14:26:27,213 | INFO | Checkpoint saved | step=10 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000010.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt
+2026-03-13 14:26:28,342 | INFO | Train step | step=11 loss=10.1618 lr=0.000220 tok_per_sec=119 grad_norm=2.1779 tokens_seen=22.53K
+2026-03-13 14:26:29,107 | INFO | Train step | step=12 loss=10.0479 lr=0.000240 tok_per_sec=2,680 grad_norm=1.9738 tokens_seen=24.58K
+2026-03-13 14:26:29,865 | INFO | Train step | step=13 loss=10.1055 lr=0.000260 tok_per_sec=2,702 grad_norm=4.0965 tokens_seen=26.62K
+2026-03-13 14:26:30,617 | INFO | Train step | step=14 loss=9.9410 lr=0.000280 tok_per_sec=2,723 grad_norm=1.8798 tokens_seen=28.67K
+2026-03-13 14:26:31,366 | INFO | Train step | step=15 loss=9.8521 lr=0.000300 tok_per_sec=2,738 grad_norm=3.7492 tokens_seen=30.72K
+2026-03-13 14:26:32,109 | INFO | Train step | step=16 loss=9.7686 lr=0.000320 tok_per_sec=2,756 grad_norm=1.8649 tokens_seen=32.77K
+2026-03-13 14:26:32,851 | INFO | Train step | step=17 loss=9.8800 lr=0.000340 tok_per_sec=2,763 grad_norm=1.8526 tokens_seen=34.82K
+2026-03-13 14:26:33,593 | INFO | Train step | step=18 loss=9.6528 lr=0.000360 tok_per_sec=2,762 grad_norm=2.2104 tokens_seen=36.86K
+2026-03-13 14:26:34,339 | INFO | Train step | step=19 loss=9.3711 lr=0.000380 tok_per_sec=2,746 grad_norm=2.0032 tokens_seen=38.91K
+2026-03-13 14:26:35,084 | INFO | Train step | step=20 loss=9.2901 lr=0.000400 tok_per_sec=2,753 grad_norm=1.9113 tokens_seen=40.96K
+2026-03-13 14:26:35,243 | INFO | Eval step | step=20 val_loss=9.6180 perplexity=15033.21
+2026-03-13 14:26:47,586 | INFO | Checkpoint saved | step=20 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000020.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt
+2026-03-13 14:26:48,510 | INFO | Train step | step=21 loss=9.2410 lr=0.000420 tok_per_sec=153 grad_norm=1.8868 tokens_seen=43.01K
+2026-03-13 14:26:49,266 | INFO | Train step | step=22 loss=9.1267 lr=0.000440 tok_per_sec=2,713 grad_norm=1.8414 tokens_seen=45.06K
+2026-03-13 14:26:50,022 | INFO | Train step | step=23 loss=9.1024 lr=0.000460 tok_per_sec=2,709 grad_norm=1.6314 tokens_seen=47.10K
+2026-03-13 14:26:50,775 | INFO | Train step | step=24 loss=8.8112 lr=0.000480 tok_per_sec=2,722 grad_norm=1.8163 tokens_seen=49.15K
+2026-03-13 14:26:51,527 | INFO | Train step | step=25 loss=8.6318 lr=0.000500 tok_per_sec=2,727 grad_norm=1.7629 tokens_seen=51.20K
+2026-03-13 14:26:52,279 | INFO | Train step | step=26 loss=8.7227 lr=0.000520 tok_per_sec=2,723 grad_norm=1.6617 tokens_seen=53.25K
+2026-03-13 14:26:53,024 | INFO | Train step | step=27 loss=8.6090 lr=0.000540 tok_per_sec=2,753 grad_norm=1.4688 tokens_seen=55.30K
+2026-03-13 14:26:53,770 | INFO | Train step | step=28 loss=8.8903 lr=0.000560 tok_per_sec=2,748 grad_norm=1.3300 tokens_seen=57.34K
+2026-03-13 14:26:54,520 | INFO | Train step | step=29 loss=8.2670 lr=0.000580 tok_per_sec=2,730 grad_norm=1.4133 tokens_seen=59.39K
+2026-03-13 14:26:55,270 | INFO | Train step | step=30 loss=8.5482 lr=0.000600 tok_per_sec=2,733 grad_norm=1.6723 tokens_seen=61.44K
+2026-03-13 14:26:55,430 | INFO | Eval step | step=30 val_loss=8.8440 perplexity=6932.60
+2026-03-13 14:27:06,544 | INFO | Checkpoint saved | step=30 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000030.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt
+2026-03-13 14:27:07,445 | INFO | Train step | step=31 loss=8.0998 lr=0.000620 tok_per_sec=168 grad_norm=1.3632 tokens_seen=63.49K
+2026-03-13 14:27:08,207 | INFO | Train step | step=32 loss=8.2781 lr=0.000640 tok_per_sec=2,689 grad_norm=1.1354 tokens_seen=65.54K
+2026-03-13 14:27:08,967 | INFO | Train step | step=33 loss=8.0486 lr=0.000660 tok_per_sec=2,696 grad_norm=0.8978 tokens_seen=67.58K
+2026-03-13 14:27:09,724 | INFO | Train step | step=34 loss=8.7096 lr=0.000680 tok_per_sec=2,707 grad_norm=1.7513 tokens_seen=69.63K
+2026-03-13 14:27:10,486 | INFO | Train step | step=35 loss=8.0777 lr=0.000700 tok_per_sec=2,690 grad_norm=0.9328 tokens_seen=71.68K
+2026-03-13 14:27:11,237 | INFO | Train step | step=36 loss=8.0570 lr=0.000720 tok_per_sec=2,730 grad_norm=0.7005 tokens_seen=73.73K
+2026-03-13 14:27:11,993 | INFO | Train step | step=37 loss=8.2809 lr=0.000740 tok_per_sec=2,712 grad_norm=0.8498 tokens_seen=75.78K
+2026-03-13 14:27:12,744 | INFO | Train step | step=38 loss=8.4201 lr=0.000760 tok_per_sec=2,728 grad_norm=0.9405 tokens_seen=77.82K
+2026-03-13 14:27:13,493 | INFO | Train step | step=39 loss=8.0400 lr=0.000780 tok_per_sec=2,734 grad_norm=0.8642 tokens_seen=79.87K
+2026-03-13 14:27:14,239 | INFO | Train step | step=40 loss=8.1934 lr=0.000800 tok_per_sec=2,747 grad_norm=0.9127 tokens_seen=81.92K
+2026-03-13 14:27:14,399 | INFO | Eval step | step=40 val_loss=8.3650 perplexity=4294.29
+2026-03-13 14:27:26,408 | INFO | Checkpoint saved | step=40 step_checkpoint=checkpoints/pretrain_mps_dryrun/step_0000040.pt last_checkpoint=checkpoints/pretrain_mps_dryrun/last.pt
+2026-03-13 14:27:27,338 | INFO | Train step | step=41 loss=7.9512 lr=0.000820 tok_per_sec=156 grad_norm=1.0788 tokens_seen=83.97K
+2026-03-13 14:27:28,091 | INFO | Train step | step=42 loss=7.9867 lr=0.000840 tok_per_sec=2,723 grad_norm=0.9518 tokens_seen=86.02K
+2026-03-13 14:27:28,841 | INFO | Train step | step=43 loss=7.9844 lr=0.000860 tok_per_sec=2,735 grad_norm=0.9768 tokens_seen=88.06K
+2026-03-13 14:27:29,644 | INFO | Train step | step=44 loss=7.7850 lr=0.000880 tok_per_sec=2,552 grad_norm=0.9357 tokens_seen=90.11K
+2026-03-13 14:27:30,394 | INFO | Train step | step=45 loss=8.1174 lr=0.000900 tok_per_sec=2,731 grad_norm=0.7717 tokens_seen=92.16K
+2026-03-13 14:27:31,145 | INFO | Train step | step=46 loss=8.2206 lr=0.000920 tok_per_sec=2,730 grad_norm=0.7731 tokens_seen=94.21K
+2026-03-13 14:27:31,898 | INFO | Train step | step=47 loss=8.0700 lr=0.000940 tok_per_sec=2,719 grad_norm=1.0135 tokens_seen=96.26K
+2026-03-13 14:27:32,652 | INFO | Train step | step=48 loss=7.9093 lr=0.000960 tok_per_sec=2,718 grad_norm=1.0888 tokens_seen=98.30K
+2026-03-13 14:27:33,406 | INFO | Train step | step=49 loss=7.9503 lr=0.000980 tok_per_sec=2,722 grad_norm=1.7942 tokens_seen=100.35K
+2026-03-13 14:27:34,161 | INFO | Train step | step=50 loss=8.5803 lr=0.001000 tok_per_sec=2,715 grad_norm=1.2601 tokens_seen=102.40K
+2026-03-13 14:27:34,323 | INFO | Eval step | step=50 val_loss=8.6077 perplexity=5473.91

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_143014.jsonl

@@ -0,0 +1,13 @@
+{"event": "run_started", "timestamp": "2026-03-13T14:30:14", "log_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_143014.log", "metrics_path": "outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_143014.jsonl", "model_config": {"vocab_size": 49152, "max_seq_len": 8192, "d_model": 384, "n_layers": 32, "n_heads": 6, "ffn_hidden_dim": 1024, "rope_theta": 10000.0, "rms_norm_eps": 1e-05, "initializer_range": 0.02, "dropout": 0.0, "tie_word_embeddings": true, "bias": false, "pad_token_id": 0, "bos_token_id": 1, "eos_token_id": 2}, "train_config": {"seed": 42, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val", "output_dir": "outputs/pretrain_mps_dryrun", "checkpoint_dir": "checkpoints/pretrain_mps_dryrun", "init_from": null, "resume_from": null, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "max_steps": 500, "warmup_steps": 50, "learning_rate": 0.001, "min_lr": 0.0001, "weight_decay": 0.1, "beta1": 0.9, "beta2": 0.95, "grad_clip": 1.0, "precision": "fp32", "num_workers": 0, "log_interval": 1, "eval_interval": 10, "eval_batches": 2, "save_interval": 10, "compile_model": false}, "args": {"model_config": "configs/model_70m.json", "train_config": "configs/pretrain_mps_dryrun.json", "max_steps_override": null}}
+{"event": "runtime_summary", "timestamp": "2026-03-13T14:30:17", "device": "mps", "precision": "fp32", "compile_model": false, "parameters": 75571584, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4, "tokens_per_step": 2048, "num_train_shards": 100, "train_dir": "data/pretokenized/train", "val_dir": "data/pretokenized/val"}
+{"event": "train", "timestamp": "2026-03-13T14:30:18", "step": 1, "loss": 10.848917245864868, "lr": 2e-05, "tok_per_sec": 1662.7320321814002, "grad_norm": 5.573695659637451, "tokens_seen": 2048, "elapsed_sec": 1.2317077919724397, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:30:19", "step": 2, "loss": 10.897652626037598, "lr": 4e-05, "tok_per_sec": 2559.899335289762, "grad_norm": 5.0279011726379395, "tokens_seen": 4096, "elapsed_sec": 0.8000314589589834, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:30:19", "step": 3, "loss": 10.785077333450317, "lr": 6e-05, "tok_per_sec": 2701.710782349432, "grad_norm": 5.1141676902771, "tokens_seen": 6144, "elapsed_sec": 0.7580382080050185, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:30:20", "step": 4, "loss": 10.634832859039307, "lr": 8e-05, "tok_per_sec": 2712.1050200194168, "grad_norm": 6.422860622406006, "tokens_seen": 8192, "elapsed_sec": 0.7551329999696463, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:30:21", "step": 5, "loss": 10.747493743896484, "lr": 0.0001, "tok_per_sec": 2709.8966635877678, "grad_norm": 6.580272674560547, "tokens_seen": 10240, "elapsed_sec": 0.7557483750279061, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:30:22", "step": 6, "loss": 10.633646488189697, "lr": 0.00012, "tok_per_sec": 2711.349354437794, "grad_norm": 6.525122165679932, "tokens_seen": 12288, "elapsed_sec": 0.7553434590226971, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:30:22", "step": 7, "loss": 10.477944374084473, "lr": 0.00014000000000000001, "tok_per_sec": 2702.4625638908824, "grad_norm": 5.189583778381348, "tokens_seen": 14336, "elapsed_sec": 0.7578273339895532, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:30:23", "step": 8, "loss": 10.35365605354309, "lr": 0.00016, "tok_per_sec": 2709.4527031511175, "grad_norm": 2.357203960418701, "tokens_seen": 16384, "elapsed_sec": 0.7558722090325318, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:30:24", "step": 9, "loss": 10.339627742767334, "lr": 0.00018, "tok_per_sec": 2711.198605030451, "grad_norm": 4.987534999847412, "tokens_seen": 18432, "elapsed_sec": 0.7553854580037296, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "train", "timestamp": "2026-03-13T14:30:25", "step": 10, "loss": 10.207262754440308, "lr": 0.0002, "tok_per_sec": 2718.8463334608755, "grad_norm": 5.554019451141357, "tokens_seen": 20480, "elapsed_sec": 0.753260666038841, "seq_len": 512, "micro_batch_size": 1, "grad_accum_steps": 4}
+{"event": "eval", "timestamp": "2026-03-13T14:30:25", "step": 10, "val_loss": 10.418238162994385, "perplexity": 33464.423318005785, "eval_batches": 2}

outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_143014.log

@@ -0,0 +1,21 @@
+2026-03-13 14:30:14,410 | INFO | Pretraining started
+2026-03-13 14:30:14,410 | INFO | Log file: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_143014.log
+2026-03-13 14:30:14,410 | INFO | Metrics JSONL: outputs/pretrain_mps_dryrun/logs/train_pretrain_20260313_143014.jsonl
+2026-03-13 14:30:14,410 | INFO | Arguments | model_config=configs/model_70m.json train_config=configs/pretrain_mps_dryrun.json max_steps_override=None
+2026-03-13 14:30:14,410 | INFO | Model config | {'vocab_size': 49152, 'max_seq_len': 8192, 'd_model': 384, 'n_layers': 32, 'n_heads': 6, 'ffn_hidden_dim': 1024, 'rope_theta': 10000.0, 'rms_norm_eps': 1e-05, 'initializer_range': 0.02, 'dropout': 0.0, 'tie_word_embeddings': True, 'bias': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2}
+2026-03-13 14:30:14,410 | INFO | Train config | {'seed': 42, 'train_dir': 'data/pretokenized/train', 'val_dir': 'data/pretokenized/val', 'output_dir': 'outputs/pretrain_mps_dryrun', 'checkpoint_dir': 'checkpoints/pretrain_mps_dryrun', 'init_from': None, 'resume_from': None, 'seq_len': 512, 'micro_batch_size': 1, 'grad_accum_steps': 4, 'max_steps': 500, 'warmup_steps': 50, 'learning_rate': 0.001, 'min_lr': 0.0001, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'precision': 'fp32', 'num_workers': 0, 'log_interval': 1, 'eval_interval': 10, 'eval_batches': 2, 'save_interval': 10, 'compile_model': False}
+2026-03-13 14:30:17,140 | INFO | Device summary | device=mps precision=fp32 compile_model=False
+2026-03-13 14:30:17,141 | INFO | Model summary | parameters=75.57M
+2026-03-13 14:30:17,141 | INFO | Batch summary | seq_len=512 micro_batch_size=1 grad_accum_steps=4 tokens_per_step=2,048
+2026-03-13 14:30:17,141 | INFO | Dataset summary | train_dir=data/pretokenized/train val_dir=data/pretokenized/val num_train_shards=100
+2026-03-13 14:30:18,374 | INFO | Train step | step=1 loss=10.8489 lr=0.000020 tok_per_sec=1,663 grad_norm=5.5737 tokens_seen=2.05K
+2026-03-13 14:30:19,174 | INFO | Train step | step=2 loss=10.8977 lr=0.000040 tok_per_sec=2,560 grad_norm=5.0279 tokens_seen=4.10K
+2026-03-13 14:30:19,933 | INFO | Train step | step=3 loss=10.7851 lr=0.000060 tok_per_sec=2,702 grad_norm=5.1142 tokens_seen=6.14K
+2026-03-13 14:30:20,688 | INFO | Train step | step=4 loss=10.6348 lr=0.000080 tok_per_sec=2,712 grad_norm=6.4229 tokens_seen=8.19K
+2026-03-13 14:30:21,445 | INFO | Train step | step=5 loss=10.7475 lr=0.000100 tok_per_sec=2,710 grad_norm=6.5803 tokens_seen=10.24K
+2026-03-13 14:30:22,201 | INFO | Train step | step=6 loss=10.6336 lr=0.000120 tok_per_sec=2,711 grad_norm=6.5251 tokens_seen=12.29K
+2026-03-13 14:30:22,959 | INFO | Train step | step=7 loss=10.4779 lr=0.000140 tok_per_sec=2,702 grad_norm=5.1896 tokens_seen=14.34K
+2026-03-13 14:30:23,715 | INFO | Train step | step=8 loss=10.3537 lr=0.000160 tok_per_sec=2,709 grad_norm=2.3572 tokens_seen=16.38K
+2026-03-13 14:30:24,471 | INFO | Train step | step=9 loss=10.3396 lr=0.000180 tok_per_sec=2,711 grad_norm=4.9875 tokens_seen=18.43K
+2026-03-13 14:30:25,225 | INFO | Train step | step=10 loss=10.2073 lr=0.000200 tok_per_sec=2,719 grad_norm=5.5540 tokens_seen=20.48K
+2026-03-13 14:30:25,394 | INFO | Eval step | step=10 val_loss=10.4182 perplexity=33464.42

outputs/pretrain_mps_dryrun/run_config.json

@@ -0,0 +1,46 @@
+{
+  "model_config": {
+    "vocab_size": 49152,
+    "max_seq_len": 8192,
+    "d_model": 384,
+    "n_layers": 32,
+    "n_heads": 6,
+    "ffn_hidden_dim": 1024,
+    "rope_theta": 10000.0,
+    "rms_norm_eps": 1e-05,
+    "initializer_range": 0.02,
+    "dropout": 0.0,
+    "tie_word_embeddings": true,
+    "bias": false,
+    "pad_token_id": 0,
+    "bos_token_id": 1,
+    "eos_token_id": 2
+  },
+  "train_config": {
+    "seed": 42,
+    "train_dir": "data/pretokenized/train",
+    "val_dir": "data/pretokenized/val",
+    "output_dir": "outputs/pretrain_mps_dryrun",
+    "checkpoint_dir": "checkpoints/pretrain_mps_dryrun",
+    "init_from": null,
+    "resume_from": null,
+    "seq_len": 512,
+    "micro_batch_size": 1,
+    "grad_accum_steps": 4,
+    "max_steps": 500,
+    "warmup_steps": 50,
+    "learning_rate": 0.001,
+    "min_lr": 0.0001,
+    "weight_decay": 0.1,
+    "beta1": 0.9,
+    "beta2": 0.95,
+    "grad_clip": 1.0,
+    "precision": "fp32",
+    "num_workers": 0,
+    "log_interval": 1,
+    "eval_interval": 10,
+    "eval_batches": 2,
+    "save_interval": 10,
+    "compile_model": false
+  }
+}

outputs/pretrain_stage1/.ipynb_checkpoints/run_config-checkpoint.json

@@ -0,0 +1,46 @@
+{
+  "model_config": {
+    "vocab_size": 49152,
+    "max_seq_len": 8192,
+    "d_model": 384,
+    "n_layers": 32,
+    "n_heads": 6,
+    "ffn_hidden_dim": 1024,
+    "rope_theta": 10000.0,
+    "rms_norm_eps": 1e-05,
+    "initializer_range": 0.02,
+    "dropout": 0.0,
+    "tie_word_embeddings": true,
+    "bias": false,
+    "pad_token_id": 0,
+    "bos_token_id": 1,
+    "eos_token_id": 2
+  },
+  "train_config": {
+    "seed": 42,
+    "train_dir": "data/pretokenized/train",
+    "val_dir": "data/pretokenized/val",
+    "output_dir": "outputs/pretrain_stage1",
+    "checkpoint_dir": "checkpoints/pretrain_stage1",
+    "init_from": null,
+    "resume_from": null,
+    "seq_len": 2048,
+    "micro_batch_size": 8,
+    "grad_accum_steps": 32,
+    "max_steps": 20000,
+    "warmup_steps": 2000,
+    "learning_rate": 0.003,
+    "min_lr": 0.0003,
+    "weight_decay": 0.1,
+    "beta1": 0.9,
+    "beta2": 0.95,
+    "grad_clip": 1.0,
+    "precision": "bf16",
+    "num_workers": 0,
+    "log_interval": 10,
+    "eval_interval": 250,
+    "eval_batches": 50,
+    "save_interval": 100,
+    "compile_model": false
+  }
+}

outputs/pretrain_stage1/logs/.ipynb_checkpoints/train_pretrain_20260313_152202-checkpoint.log

@@ -0,0 +1,82 @@
+2026-03-13 15:22:02,275 | INFO | Pretraining started
+2026-03-13 15:22:02,276 | INFO | Log file: outputs/pretrain_stage1/logs/train_pretrain_20260313_152202.log
+2026-03-13 15:22:02,276 | INFO | Metrics JSONL: outputs/pretrain_stage1/logs/train_pretrain_20260313_152202.jsonl
+2026-03-13 15:22:02,276 | INFO | Arguments | model_config=configs/model_70m.json train_config=configs/pretrain_5090_stage1.json max_steps_override=None
+2026-03-13 15:22:02,276 | INFO | Model config | {'vocab_size': 49152, 'max_seq_len': 8192, 'd_model': 384, 'n_layers': 32, 'n_heads': 6, 'ffn_hidden_dim': 1024, 'rope_theta': 10000.0, 'rms_norm_eps': 1e-05, 'initializer_range': 0.02, 'dropout': 0.0, 'tie_word_embeddings': True, 'bias': False, 'pad_token_id': 0, 'bos_token_id': 1, 'eos_token_id': 2}
+2026-03-13 15:22:02,276 | INFO | Train config | {'seed': 42, 'train_dir': 'data/pretokenized/train', 'val_dir': 'data/pretokenized/val', 'output_dir': 'outputs/pretrain_stage1', 'checkpoint_dir': 'checkpoints/pretrain_stage1', 'init_from': None, 'resume_from': None, 'seq_len': 2048, 'micro_batch_size': 8, 'grad_accum_steps': 32, 'max_steps': 20000, 'warmup_steps': 2000, 'learning_rate': 0.003, 'min_lr': 0.0003, 'weight_decay': 0.1, 'beta1': 0.9, 'beta2': 0.95, 'grad_clip': 1.0, 'precision': 'bf16', 'num_workers': 0, 'log_interval': 10, 'eval_interval': 250, 'eval_batches': 50, 'save_interval': 100, 'compile_model': False}
+2026-03-13 15:22:05,635 | INFO | Device summary | device=cuda precision=bf16 compile_model=False
+2026-03-13 15:22:05,636 | INFO | Model summary | parameters=75.57M
+2026-03-13 15:22:05,636 | INFO | Batch summary | seq_len=2048 micro_batch_size=8 grad_accum_steps=32 tokens_per_step=524,288
+2026-03-13 15:22:05,636 | INFO | Dataset summary | train_dir=data/pretokenized/train val_dir=data/pretokenized/val num_train_shards=100
+2026-03-13 15:22:48,364 | INFO | Train step | step=10 loss=10.7962 lr=0.000015 tok_per_sec=122,709 grad_norm=2.2954 tokens_seen=5.24M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:23:29,710 | INFO | Train step | step=20 loss=10.3929 lr=0.000030 tok_per_sec=126,809 grad_norm=1.6374 tokens_seen=10.49M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:24:11,004 | INFO | Train step | step=30 loss=10.1422 lr=0.000045 tok_per_sec=126,967 grad_norm=1.6471 tokens_seen=15.73M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:24:53,082 | INFO | Train step | step=40 loss=9.9494 lr=0.000060 tok_per_sec=124,605 grad_norm=1.5930 tokens_seen=20.97M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:25:36,401 | INFO | Train step | step=50 loss=9.6967 lr=0.000075 tok_per_sec=121,032 grad_norm=1.5725 tokens_seen=26.21M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:26:19,167 | INFO | Train step | step=60 loss=9.3897 lr=0.000090 tok_per_sec=122,597 grad_norm=1.5564 tokens_seen=31.46M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:27:01,620 | INFO | Train step | step=70 loss=9.0575 lr=0.000105 tok_per_sec=123,501 grad_norm=1.5012 tokens_seen=36.70M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:27:43,580 | INFO | Train step | step=80 loss=8.6948 lr=0.000120 tok_per_sec=124,954 grad_norm=1.5047 tokens_seen=41.94M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:28:26,150 | INFO | Train step | step=90 loss=8.3511 lr=0.000135 tok_per_sec=123,163 grad_norm=1.2600 tokens_seen=47.19M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:29:09,001 | INFO | Train step | step=100 loss=8.0548 lr=0.000150 tok_per_sec=122,354 grad_norm=0.9670 tokens_seen=52.43M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:29:12,440 | INFO | Checkpoint saved | step=100 step_checkpoint=checkpoints/pretrain_stage1/step_0000100.pt last_checkpoint=checkpoints/pretrain_stage1/last.pt
+2026-03-13 15:29:54,815 | INFO | Train step | step=110 loss=7.8111 lr=0.000165 tok_per_sec=114,442 grad_norm=0.7305 tokens_seen=57.67M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:30:37,782 | INFO | Train step | step=120 loss=7.6241 lr=0.000180 tok_per_sec=122,024 grad_norm=0.5833 tokens_seen=62.91M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:31:19,995 | INFO | Train step | step=130 loss=7.4835 lr=0.000195 tok_per_sec=124,205 grad_norm=1.0428 tokens_seen=68.16M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:32:03,842 | INFO | Train step | step=140 loss=7.3397 lr=0.000210 tok_per_sec=119,576 grad_norm=0.6136 tokens_seen=73.40M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:32:45,816 | INFO | Train step | step=150 loss=7.1952 lr=0.000225 tok_per_sec=124,911 grad_norm=1.2209 tokens_seen=78.64M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:33:27,516 | INFO | Train step | step=160 loss=7.0569 lr=0.000240 tok_per_sec=125,732 grad_norm=0.9325 tokens_seen=83.89M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:34:09,763 | INFO | Train step | step=170 loss=6.9308 lr=0.000255 tok_per_sec=124,102 grad_norm=1.1994 tokens_seen=89.13M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:34:51,580 | INFO | Train step | step=180 loss=6.7975 lr=0.000270 tok_per_sec=125,380 grad_norm=1.2646 tokens_seen=94.37M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:35:33,306 | INFO | Train step | step=190 loss=6.6834 lr=0.000285 tok_per_sec=125,653 grad_norm=0.9549 tokens_seen=99.61M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:36:16,369 | INFO | Train step | step=200 loss=6.5762 lr=0.000300 tok_per_sec=121,752 grad_norm=1.5983 tokens_seen=104.86M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:36:18,831 | INFO | Checkpoint saved | step=200 step_checkpoint=checkpoints/pretrain_stage1/step_0000200.pt last_checkpoint=checkpoints/pretrain_stage1/last.pt
+2026-03-13 15:37:01,458 | INFO | Train step | step=210 loss=6.4800 lr=0.000315 tok_per_sec=116,281 grad_norm=0.9575 tokens_seen=110.10M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:37:43,418 | INFO | Train step | step=220 loss=6.3799 lr=0.000330 tok_per_sec=124,955 grad_norm=1.0189 tokens_seen=115.34M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:38:25,012 | INFO | Train step | step=230 loss=6.3007 lr=0.000345 tok_per_sec=126,050 grad_norm=1.4322 tokens_seen=120.59M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:39:07,059 | INFO | Train step | step=240 loss=6.2100 lr=0.000360 tok_per_sec=124,696 grad_norm=1.4284 tokens_seen=125.83M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:39:48,188 | INFO | Train step | step=250 loss=6.1378 lr=0.000375 tok_per_sec=127,476 grad_norm=0.8502 tokens_seen=131.07M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:39:50,336 | INFO | Eval step | step=250 val_loss=6.1302 perplexity=459.54
+2026-03-13 15:40:32,241 | INFO | Train step | step=260 loss=6.0712 lr=0.000390 tok_per_sec=119,017 grad_norm=1.5691 tokens_seen=136.31M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:41:15,172 | INFO | Train step | step=270 loss=6.0020 lr=0.000405 tok_per_sec=122,129 grad_norm=1.3161 tokens_seen=141.56M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:41:57,319 | INFO | Train step | step=280 loss=5.9392 lr=0.000420 tok_per_sec=124,398 grad_norm=1.3891 tokens_seen=146.80M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:42:39,253 | INFO | Train step | step=290 loss=5.8713 lr=0.000435 tok_per_sec=125,030 grad_norm=1.1325 tokens_seen=152.04M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:43:21,127 | INFO | Train step | step=300 loss=5.8109 lr=0.000450 tok_per_sec=125,209 grad_norm=1.0078 tokens_seen=157.29M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:43:23,493 | INFO | Checkpoint saved | step=300 step_checkpoint=checkpoints/pretrain_stage1/step_0000300.pt last_checkpoint=checkpoints/pretrain_stage1/last.pt
+2026-03-13 15:44:04,852 | INFO | Train step | step=310 loss=5.7384 lr=0.000465 tok_per_sec=119,907 grad_norm=1.2581 tokens_seen=162.53M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:44:46,724 | INFO | Train step | step=320 loss=5.6798 lr=0.000480 tok_per_sec=125,216 grad_norm=0.9680 tokens_seen=167.77M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:45:29,877 | INFO | Train step | step=330 loss=5.6204 lr=0.000495 tok_per_sec=121,497 grad_norm=1.4606 tokens_seen=173.02M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:46:12,195 | INFO | Train step | step=340 loss=5.5678 lr=0.000510 tok_per_sec=123,896 grad_norm=1.2717 tokens_seen=178.26M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:46:54,104 | INFO | Train step | step=350 loss=5.5266 lr=0.000525 tok_per_sec=125,105 grad_norm=1.6313 tokens_seen=183.50M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:47:36,611 | INFO | Train step | step=360 loss=5.4781 lr=0.000540 tok_per_sec=123,343 grad_norm=1.0196 tokens_seen=188.74M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:48:18,069 | INFO | Train step | step=370 loss=5.4230 lr=0.000555 tok_per_sec=126,468 grad_norm=1.0206 tokens_seen=193.99M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:49:00,176 | INFO | Train step | step=380 loss=5.3519 lr=0.000570 tok_per_sec=124,516 grad_norm=0.7121 tokens_seen=199.23M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:49:42,492 | INFO | Train step | step=390 loss=5.3026 lr=0.000585 tok_per_sec=123,899 grad_norm=1.0407 tokens_seen=204.47M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:50:24,343 | INFO | Train step | step=400 loss=5.2721 lr=0.000600 tok_per_sec=125,278 grad_norm=0.7830 tokens_seen=209.72M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:50:27,093 | INFO | Checkpoint saved | step=400 step_checkpoint=checkpoints/pretrain_stage1/step_0000400.pt last_checkpoint=checkpoints/pretrain_stage1/last.pt
+2026-03-13 15:51:08,698 | INFO | Train step | step=410 loss=5.2136 lr=0.000615 tok_per_sec=118,206 grad_norm=0.6625 tokens_seen=214.96M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:51:50,429 | INFO | Train step | step=420 loss=5.1839 lr=0.000630 tok_per_sec=125,640 grad_norm=1.1878 tokens_seen=220.20M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:52:31,919 | INFO | Train step | step=430 loss=5.1433 lr=0.000645 tok_per_sec=126,367 grad_norm=1.0909 tokens_seen=225.44M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:53:14,261 | INFO | Train step | step=440 loss=5.0811 lr=0.000660 tok_per_sec=123,827 grad_norm=1.0818 tokens_seen=230.69M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:53:55,426 | INFO | Train step | step=450 loss=5.0691 lr=0.000675 tok_per_sec=127,367 grad_norm=0.8735 tokens_seen=235.93M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:54:36,815 | INFO | Train step | step=460 loss=5.0245 lr=0.000690 tok_per_sec=126,676 grad_norm=0.7781 tokens_seen=241.17M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:55:18,411 | INFO | Train step | step=470 loss=4.9740 lr=0.000705 tok_per_sec=126,046 grad_norm=0.8157 tokens_seen=246.42M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:55:59,935 | INFO | Train step | step=480 loss=4.9158 lr=0.000720 tok_per_sec=126,265 grad_norm=0.4327 tokens_seen=251.66M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:56:41,577 | INFO | Train step | step=490 loss=4.8794 lr=0.000735 tok_per_sec=125,907 grad_norm=0.9491 tokens_seen=256.90M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:57:23,706 | INFO | Train step | step=500 loss=4.8574 lr=0.000750 tok_per_sec=124,451 grad_norm=0.7693 tokens_seen=262.14M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:57:26,171 | INFO | Eval step | step=500 val_loss=4.8718 perplexity=130.55
+2026-03-13 15:57:30,367 | INFO | Checkpoint saved | step=500 step_checkpoint=checkpoints/pretrain_stage1/step_0000500.pt last_checkpoint=checkpoints/pretrain_stage1/last.pt
+2026-03-13 15:58:13,490 | INFO | Train step | step=510 loss=4.8269 lr=0.000765 tok_per_sec=105,314 grad_norm=0.8741 tokens_seen=267.39M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:58:56,422 | INFO | Train step | step=520 loss=4.7787 lr=0.000780 tok_per_sec=122,122 grad_norm=0.5603 tokens_seen=272.63M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 15:59:39,461 | INFO | Train step | step=530 loss=4.7081 lr=0.000795 tok_per_sec=121,823 grad_norm=0.7208 tokens_seen=277.87M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:00:23,234 | INFO | Train step | step=540 loss=4.6785 lr=0.000810 tok_per_sec=119,776 grad_norm=0.6182 tokens_seen=283.12M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:01:05,239 | INFO | Train step | step=550 loss=4.6483 lr=0.000825 tok_per_sec=124,821 grad_norm=0.8779 tokens_seen=288.36M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:01:47,156 | INFO | Train step | step=560 loss=4.6100 lr=0.000840 tok_per_sec=125,080 grad_norm=0.7765 tokens_seen=293.60M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:02:30,946 | INFO | Train step | step=570 loss=4.5568 lr=0.000855 tok_per_sec=119,733 grad_norm=0.5192 tokens_seen=298.84M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:03:13,505 | INFO | Train step | step=580 loss=4.5020 lr=0.000870 tok_per_sec=123,194 grad_norm=0.4420 tokens_seen=304.09M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:03:55,388 | INFO | Train step | step=590 loss=4.4536 lr=0.000885 tok_per_sec=125,182 grad_norm=0.4726 tokens_seen=309.33M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:04:37,233 | INFO | Train step | step=600 loss=4.4008 lr=0.000900 tok_per_sec=125,295 grad_norm=0.5401 tokens_seen=314.57M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:04:40,532 | INFO | Checkpoint saved | step=600 step_checkpoint=checkpoints/pretrain_stage1/step_0000600.pt last_checkpoint=checkpoints/pretrain_stage1/last.pt
+2026-03-13 16:05:22,263 | INFO | Train step | step=610 loss=4.3697 lr=0.000915 tok_per_sec=116,433 grad_norm=0.5282 tokens_seen=319.82M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:06:04,579 | INFO | Train step | step=620 loss=4.3184 lr=0.000930 tok_per_sec=123,903 grad_norm=0.8301 tokens_seen=325.06M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:06:47,513 | INFO | Train step | step=630 loss=4.3098 lr=0.000945 tok_per_sec=122,117 grad_norm=0.4351 tokens_seen=330.30M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53
+2026-03-13 16:07:29,171 | INFO | Train step | step=640 loss=4.2368 lr=0.000960 tok_per_sec=125,858 grad_norm=0.4222 tokens_seen=335.54M mem_alloc_gb=1.28 mem_reserved_gb=23.53 max_mem_alloc_gb=19.63 max_mem_reserved_gb=23.53

outputs/pretrain_stage1/run_config.json

@@ -0,0 +1,46 @@
+{
+  "model_config": {
+    "vocab_size": 49152,
+    "max_seq_len": 8192,
+    "d_model": 384,
+    "n_layers": 32,
+    "n_heads": 6,
+    "ffn_hidden_dim": 1024,
+    "rope_theta": 10000.0,
+    "rms_norm_eps": 1e-05,
+    "initializer_range": 0.02,
+    "dropout": 0.0,
+    "tie_word_embeddings": true,
+    "bias": false,
+    "pad_token_id": 0,
+    "bos_token_id": 1,
+    "eos_token_id": 2
+  },
+  "train_config": {
+    "seed": 42,
+    "train_dir": "data/pretokenized/train",
+    "val_dir": "data/pretokenized/val",
+    "output_dir": "outputs/pretrain_stage1",
+    "checkpoint_dir": "checkpoints/pretrain_stage1",
+    "init_from": null,
+    "resume_from": null,
+    "seq_len": 2048,
+    "micro_batch_size": 8,
+    "grad_accum_steps": 32,
+    "max_steps": 20000,
+    "warmup_steps": 2000,
+    "learning_rate": 0.003,
+    "min_lr": 0.0003,
+    "weight_decay": 0.1,
+    "beta1": 0.9,
+    "beta2": 0.95,
+    "grad_clip": 1.0,
+    "precision": "bf16",
+    "num_workers": 0,
+    "log_interval": 10,
+    "eval_interval": 250,
+    "eval_batches": 50,
+    "save_interval": 100,
+    "compile_model": false
+  }
+}

outputs/pretrain_stage2/run_config.json

@@ -0,0 +1,46 @@
+{
+  "model_config": {
+    "vocab_size": 49152,
+    "max_seq_len": 8192,
+    "d_model": 384,
+    "n_layers": 32,
+    "n_heads": 6,
+    "ffn_hidden_dim": 1024,
+    "rope_theta": 10000.0,
+    "rms_norm_eps": 1e-05,
+    "initializer_range": 0.02,
+    "dropout": 0.0,
+    "tie_word_embeddings": true,
+    "bias": false,
+    "pad_token_id": 0,
+    "bos_token_id": 1,
+    "eos_token_id": 2
+  },
+  "train_config": {
+    "seed": 42,
+    "train_dir": "data/pretokenized/train",
+    "val_dir": "data/pretokenized/val",
+    "output_dir": "outputs/pretrain_stage2",
+    "checkpoint_dir": "checkpoints/pretrain_stage2",
+    "init_from": "checkpoints/pretrain_stage1/last.pt",
+    "resume_from": null,
+    "seq_len": 8192,
+    "micro_batch_size": 2,
+    "grad_accum_steps": 16,
+    "max_steps": 1000,
+    "warmup_steps": 100,
+    "learning_rate": 0.001,
+    "min_lr": 0.0001,
+    "weight_decay": 0.1,
+    "beta1": 0.9,
+    "beta2": 0.95,
+    "grad_clip": 1.0,
+    "precision": "bf16",
+    "num_workers": 0,
+    "log_interval": 5,
+    "eval_interval": 100,
+    "eval_batches": 20,
+    "save_interval": 50,
+    "compile_model": false
+  }
+}

scripts/eval_perplexity.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import argparse
+import math
+import sys
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.append(str(ROOT / "src"))
+
+from sllm.checkpoint import load_checkpoint
+from sllm.config import ModelConfig, load_json
+from sllm.data import SequentialEvalDataset
+from sllm.model import SLLMForCausalLM
+from sllm.utils import autocast_context, get_device, resolve_runtime_precision, setup_logger
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Evaluate perplexity on validation shards.")
+    parser.add_argument("--checkpoint", required=True, help="Path to checkpoint file.")
+    parser.add_argument("--model-config", required=False, help="Optional model config JSON path.")
+    parser.add_argument("--data-dir", required=True, help="Validation root directory.")
+    parser.add_argument("--seq-len", type=int, default=2_048)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--batches", type=int, default=50)
+    parser.add_argument("--precision", default="bf16")
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    logger, log_path = setup_logger("sllm.eval_perplexity", Path("outputs/eval"), "eval_perplexity")
+    logger.info("Perplexity evaluation started")
+    logger.info("Log file: %s", log_path)
+    logger.info("Arguments | checkpoint=%s model_config=%s data_dir=%s seq_len=%s batch_size=%s batches=%s precision=%s", args.checkpoint, args.model_config, args.data_dir, args.seq_len, args.batch_size, args.batches, args.precision)
+    device = get_device()
+    runtime_precision, precision_warning = resolve_runtime_precision(device, args.precision)
+    if precision_warning is not None:
+        logger.warning(precision_warning)
+    payload = load_checkpoint(args.checkpoint, map_location=device)
+    if args.model_config:
+        model_config = ModelConfig.from_dict(load_json(args.model_config))
+    else:
+        model_config = ModelConfig.from_dict(payload["model_config"])
+
+    model = SLLMForCausalLM(model_config).to(device)
+    model.load_state_dict(payload["model"])
+    model.eval()
+
+    dataset = SequentialEvalDataset(
+        data_dir=args.data_dir,
+        split="val",
+        seq_len=args.seq_len,
+        max_batches=args.batches * args.batch_size,
+    )
+    loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=0)
+
+    losses = []
+    with torch.no_grad():
+        for batch_index, batch in enumerate(loader):
+            if batch_index >= args.batches:
+                break
+            batch = {key: value.to(device) for key, value in batch.items()}
+            with autocast_context(device, runtime_precision):
+                loss = model(**batch)["loss"]
+            losses.append(loss.detach().float().item())
+
+    mean_loss = float(sum(losses) / max(1, len(losses)))
+    perplexity = math.exp(min(mean_loss, 20))
+    logger.info("Perplexity evaluation finished | val_loss=%.4f perplexity=%.2f", mean_loss, perplexity)
+    print(f"val_loss={mean_loss:.4f}")
+    print(f"perplexity={perplexity:.2f}")
+
+
+if __name__ == "__main__":
+    main()

scripts/generate.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+from tokenizers import Tokenizer
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.append(str(ROOT / "src"))
+
+from sllm.checkpoint import load_checkpoint
+from sllm.config import ModelConfig, load_json
+from sllm.model import SLLMForCausalLM
+from sllm.utils import get_device, setup_logger
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Generate text from a trained checkpoint.")
+    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint.")
+    parser.add_argument("--tokenizer-dir", required=True, help="Directory with tokenizer.json.")
+    parser.add_argument("--prompt", required=True, help="Prompt text.")
+    parser.add_argument("--max-new-tokens", type=int, default=128)
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--top-k", type=int, default=50)
+    parser.add_argument("--model-config", required=False, help="Optional path to model config JSON.")
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    logger, log_path = setup_logger("sllm.generate", Path("outputs/generate"), "generate")
+    logger.info("Generation started")
+    logger.info("Log file: %s", log_path)
+    logger.info(
+        "Arguments | checkpoint=%s tokenizer_dir=%s max_new_tokens=%s temperature=%s top_k=%s model_config=%s",
+        args.checkpoint,
+        args.tokenizer_dir,
+        args.max_new_tokens,
+        args.temperature,
+        args.top_k,
+        args.model_config,
+    )
+    device = get_device()
+    tokenizer = Tokenizer.from_file(str(Path(args.tokenizer_dir) / "tokenizer.json"))
+    tokenizer_meta = load_json(Path(args.tokenizer_dir) / "tokenizer_meta.json")
+    specials = tokenizer_meta["special_tokens"]
+
+    payload = load_checkpoint(args.checkpoint, map_location=device)
+    if args.model_config:
+        model_config = ModelConfig.from_dict(load_json(args.model_config))
+    else:
+        model_config = ModelConfig.from_dict(payload["model_config"])
+
+    model = SLLMForCausalLM(model_config).to(device)
+    model.load_state_dict(payload["model"])
+    model.eval()
+
+    prompt_ids = [int(specials["bos_token_id"])] + tokenizer.encode(
+        args.prompt,
+        add_special_tokens=False,
+    ).ids
+    input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            input_ids=input_ids,
+            max_new_tokens=args.max_new_tokens,
+            temperature=args.temperature,
+            top_k=args.top_k,
+            eos_token_id=int(specials["eos_token_id"]),
+        )
+
+    decoded = tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=False)
+    logger.info("Generation finished | prompt_tokens=%s output_tokens=%s", len(prompt_ids), output_ids.shape[1])
+    print(decoded)
+
+
+if __name__ == "__main__":
+    main()

scripts/prepare_pretrain_data.py

@@ -0,0 +1,318 @@
+from __future__ import annotations
+
+import argparse
+import math
+import random
+import sys
+from collections import deque
+from pathlib import Path
+
+from datasets import load_dataset
+from tokenizers import Tokenizer
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.append(str(ROOT / "src"))
+
+from sllm.config import DataMixConfig, load_json, save_json
+from sllm.data import TokenShardWriter
+from sllm.utils import setup_logger
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Tokenize and shard pretraining corpora.")
+    parser.add_argument("--data-config", required=True, help="Path to data mixture JSON config.")
+    parser.add_argument("--tokenizer-dir", required=True, help="Directory with tokenizer.json.")
+    parser.add_argument("--output-dir", required=True, help="Root directory for train/val shards.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for dataset shuffling.")
+    return parser
+
+
+def load_tokenizer(tokenizer_dir: str | Path) -> tuple[Tokenizer, dict]:
+    tokenizer_dir = Path(tokenizer_dir)
+    tokenizer = Tokenizer.from_file(str(tokenizer_dir / "tokenizer.json"))
+    metadata = load_json(tokenizer_dir / "tokenizer_meta.json")
+    return tokenizer, metadata
+
+
+def iter_source_rows(source, seed: int):
+    dataset = load_dataset(
+        path=source.path,
+        name=source.config_name,
+        data_dir=source.data_dir,
+        split=source.split,
+        revision=source.revision,
+        streaming=source.streaming,
+    )
+    if source.streaming:
+        dataset = dataset.shuffle(seed=seed, buffer_size=source.shuffle_buffer)
+    return iter(dataset)
+
+
+TOKENIZE_BATCH_SIZE = 128
+
+
+def allocate_token_targets(data_config: DataMixConfig, total_tokens: int) -> dict[str, int]:
+    weights = data_config.normalized_weights()
+    raw_targets = {
+        source.name: total_tokens * weights[source.name]
+        for source in data_config.sources
+    }
+    base_targets = {
+        name: int(math.floor(value))
+        for name, value in raw_targets.items()
+    }
+    remainder = total_tokens - sum(base_targets.values())
+    ranked = sorted(
+        raw_targets.items(),
+        key=lambda item: (item[1] - math.floor(item[1]), item[0]),
+        reverse=True,
+    )
+    for index in range(remainder):
+        name = ranked[index % len(ranked)][0]
+        base_targets[name] += 1
+    return base_targets
+
+
+def make_source_state(source, seed: int) -> dict:
+    return {
+        "source": source,
+        "iterator": iter_source_rows(source, seed),
+        "documents_used": 0,
+        "train_tokens_written": 0,
+        "val_tokens_written": 0,
+        "exhausted": False,
+        "token_queue": deque(),
+    }
+
+
+def refill_token_queue(state: dict, tokenizer: Tokenizer) -> None:
+    if state["exhausted"]:
+        return
+
+    texts: list[str] = []
+    while len(texts) < TOKENIZE_BATCH_SIZE:
+        try:
+            row = next(state["iterator"])
+        except StopIteration:
+            state["exhausted"] = True
+            break
+
+        text = row.get(state["source"].text_field or "", None)
+        if not isinstance(text, str):
+            continue
+        text = text.strip()
+        if not text:
+            continue
+        texts.append(text)
+
+    if not texts:
+        return
+
+    encoded_batch = tokenizer.encode_batch(texts)
+    for encoded in encoded_batch:
+        token_ids = encoded.ids
+        if token_ids:
+            state["token_queue"].append(token_ids)
+
+
+def next_valid_token_ids(state: dict, tokenizer: Tokenizer) -> list[int] | None:
+    while True:
+        if state["token_queue"]:
+            state["documents_used"] += 1
+            return state["token_queue"].popleft()
+        if state["exhausted"]:
+            return None
+        refill_token_queue(state, tokenizer)
+
+
+def choose_source_name(states: dict[str, dict], targets: dict[str, int], split: str, rng: random.Random) -> str | None:
+    candidates = []
+    for name, state in states.items():
+        if state["exhausted"]:
+            continue
+        target = targets[name]
+        if target <= 0:
+            continue
+        written = state[f"{split}_tokens_written"]
+        if written >= target:
+            continue
+        progress = written / target
+        candidates.append((progress, rng.random(), name))
+
+    if not candidates:
+        return None
+
+    candidates.sort(key=lambda item: (item[0], item[1]))
+    return candidates[0][2]
+
+
+def interleave_split(
+    split: str,
+    writer: TokenShardWriter,
+    states: dict[str, dict],
+    targets: dict[str, int],
+    tokenizer: Tokenizer,
+    logger,
+    rng: random.Random,
+) -> int:
+    total_target = sum(targets.values())
+    total_written = 0
+    emitted_documents = 0
+
+    logger.info(
+        "Interleave start | split=%s total_target_tokens=%s strategy=weighted_progress_balancing",
+        split,
+        f"{total_target:,}",
+    )
+
+    while total_written < total_target:
+        source_name = choose_source_name(states, targets, split, rng)
+        if source_name is None:
+            raise RuntimeError(
+                f"Недостаточно данных для заполнения split={split}. "
+                "Все доступные источники исчерпаны до достижения целевого объема."
+            )
+
+        state = states[source_name]
+        token_ids = next_valid_token_ids(state, tokenizer)
+        if token_ids is None:
+            logger.warning("Source exhausted early | split=%s source=%s", split, source_name)
+            continue
+
+        source_remaining = targets[source_name] - state[f"{split}_tokens_written"]
+        split_remaining = total_target - total_written
+        chunk = token_ids[: min(len(token_ids), source_remaining, split_remaining)]
+        if not chunk:
+            continue
+
+        writer.add_tokens(chunk)
+        state[f"{split}_tokens_written"] += len(chunk)
+        total_written += len(chunk)
+        emitted_documents += 1
+
+        if emitted_documents % 10_000 == 0:
+            logger.info(
+                "Interleave progress | split=%s documents=%s total_tokens=%s/%s current_source=%s",
+                split,
+                f"{emitted_documents:,}",
+                f"{total_written:,}",
+                f"{total_target:,}",
+                source_name,
+            )
+
+    logger.info(
+        "Interleave done | split=%s documents=%s total_tokens=%s",
+        split,
+        f"{emitted_documents:,}",
+        f"{total_written:,}",
+    )
+    return total_written
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    data_config = DataMixConfig.from_dict(load_json(args.data_config))
+    tokenizer, tokenizer_meta = load_tokenizer(args.tokenizer_dir)
+    output_dir = Path(args.output_dir)
+    train_dir = output_dir / "train"
+    val_dir = output_dir / "val"
+    train_dir.mkdir(parents=True, exist_ok=True)
+    val_dir.mkdir(parents=True, exist_ok=True)
+    logger, log_path = setup_logger("sllm.prepare_pretrain_data", output_dir, "prepare_pretrain_data")
+    logger.info("Pretokenization started")
+    logger.info("Log file: %s", log_path)
+    logger.info("Arguments | data_config=%s tokenizer_dir=%s output_dir=%s seed=%s", args.data_config, args.tokenizer_dir, args.output_dir, args.seed)
+    logger.info("Tokenizer meta | vocab_size=%s special_tokens=%s", tokenizer_meta.get("vocab_size"), tokenizer_meta.get("special_tokens"))
+    logger.info("Mixing strategy | global interleaving with weighted progress balancing")
+    logger.info("Tokenization strategy | encode_batch with batch_size=%s", TOKENIZE_BATCH_SIZE)
+
+    weight_map = data_config.normalized_weights()
+    train_targets = allocate_token_targets(data_config, data_config.train_tokens)
+    val_targets = allocate_token_targets(data_config, data_config.val_tokens)
+    dataset_summary: dict[str, dict] = {}
+    states: dict[str, dict] = {}
+
+    for index, source in enumerate(data_config.sources):
+        states[source.name] = make_source_state(source, args.seed + index)
+        logger.info(
+            "Source registered | name=%s path=%s data_dir=%s split=%s text_field=%s weight=%.4f train_target=%s val_target=%s streaming=%s",
+            source.name,
+            source.path,
+            source.data_dir,
+            source.split,
+            source.text_field,
+            weight_map[source.name],
+            f"{train_targets[source.name]:,}",
+            f"{val_targets[source.name]:,}",
+            source.streaming,
+        )
+
+    rng_val = random.Random(args.seed + 10_000)
+    rng_train = random.Random(args.seed + 20_000)
+    val_writer = TokenShardWriter(
+        output_dir=val_dir,
+        prefix="val",
+        shard_size_tokens=max(1_000_000, min(data_config.shard_size_tokens, data_config.val_tokens)),
+    )
+    train_writer = TokenShardWriter(
+        output_dir=train_dir,
+        prefix="train",
+        shard_size_tokens=data_config.shard_size_tokens,
+    )
+
+    total_val = interleave_split("val", val_writer, states, val_targets, tokenizer, logger, rng_val)
+    total_train = interleave_split("train", train_writer, states, train_targets, tokenizer, logger, rng_train)
+
+    train_shards = train_writer.finalize()
+    val_shards = val_writer.finalize()
+
+    for source in data_config.sources:
+        state = states[source.name]
+        dataset_summary[source.name] = {
+            "path": source.path,
+            "data_dir": source.data_dir,
+            "split": source.split,
+            "train_target_tokens": train_targets[source.name],
+            "val_target_tokens": val_targets[source.name],
+            "train_tokens_written": state["train_tokens_written"],
+            "val_tokens_written": state["val_tokens_written"],
+            "documents_used": state["documents_used"],
+        }
+        logger.info(
+            "Source done | name=%s documents=%s train_tokens=%s/%s val_tokens=%s/%s",
+            source.name,
+            f"{state['documents_used']:,}",
+            f"{state['train_tokens_written']:,}",
+            f"{train_targets[source.name]:,}",
+            f"{state['val_tokens_written']:,}",
+            f"{val_targets[source.name]:,}",
+        )
+
+    save_json(
+        output_dir / "dataset_summary.json",
+        {
+            "tokenizer": tokenizer_meta,
+            "data_config": data_config.to_dict(),
+            "mixing_strategy": "global_interleaving_weighted_progress_balancing",
+            "train_target_tokens": data_config.train_tokens,
+            "val_target_tokens": data_config.val_tokens,
+            "train_tokens_written": total_train,
+            "val_tokens_written": total_val,
+            "train_shards": len(train_shards),
+            "val_shards": len(val_shards),
+            "sources": dataset_summary,
+        },
+    )
+    logger.info(
+        "Pretokenization finished | output_dir=%s total_train_tokens=%s total_val_tokens=%s train_shards=%s val_shards=%s",
+        output_dir,
+        f"{total_train:,}",
+        f"{total_val:,}",
+        len(train_shards),
+        len(val_shards),
+    )
+    logger.info("Dataset summary saved | path=%s", output_dir / "dataset_summary.json")
+
+
+if __name__ == "__main__":
+    main()

scripts/prepare_sft_data.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from datasets import load_dataset
+from tokenizers import Tokenizer
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.append(str(ROOT / "src"))
+
+from sllm.config import load_json, save_json
+from sllm.data import SFTShardWriter
+from sllm.utils import setup_logger
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Prepare fixed-length SFT tensors.")
+    parser.add_argument("--config", required=True, help="Path to SFT data JSON config.")
+    parser.add_argument("--tokenizer-dir", required=True, help="Directory with tokenizer.json and metadata.")
+    parser.add_argument("--output-dir", required=True, help="Directory to store processed SFT tensors.")
+    parser.add_argument("--seq-len", type=int, default=2_048, help="Packed example length.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for dataset shuffling.")
+    return parser
+
+
+def load_tokenizer(tokenizer_dir: str | Path) -> tuple[Tokenizer, dict]:
+    tokenizer_dir = Path(tokenizer_dir)
+    tokenizer = Tokenizer.from_file(str(tokenizer_dir / "tokenizer.json"))
+    metadata = load_json(tokenizer_dir / "tokenizer_meta.json")
+    return tokenizer, metadata
+
+
+def row_to_messages(row: dict, config: dict) -> list[dict[str, str]]:
+    fmt = config.get("format", "messages")
+    if fmt == "messages":
+        messages = row.get(config.get("messages_field", "messages"))
+        if not isinstance(messages, list):
+            raise ValueError("Не найден список сообщений в SFT-датасете.")
+        normalized = []
+        for message in messages:
+            if not isinstance(message, dict):
+                continue
+            role = message.get("role")
+            content = message.get("content")
+            if isinstance(content, list):
+                parts = [item.get("text", "") for item in content if isinstance(item, dict)]
+                content = "\n".join(part for part in parts if part)
+            if isinstance(role, str) and isinstance(content, str) and content.strip():
+                normalized.append({"role": role, "content": content.strip()})
+        return normalized
+
+    if fmt == "prompt_response":
+        prompt = row.get(config.get("prompt_field", "prompt"))
+        response = row.get(config.get("response_field", "response"))
+        if not isinstance(prompt, str) or not isinstance(response, str):
+            raise ValueError("Не найдены поля prompt/response в SFT-датасете.")
+        system_prompt = config.get("system_prompt")
+        messages = []
+        if isinstance(system_prompt, str) and system_prompt.strip():
+            messages.append({"role": "system", "content": system_prompt.strip()})
+        messages.append({"role": "user", "content": prompt.strip()})
+        messages.append({"role": "assistant", "content": response.strip()})
+        return messages
+
+    if fmt == "alpaca":
+        instruction = row.get(config.get("instruction_field", "instruction"))
+        input_text = row.get(config.get("input_field", "input"), "")
+        output_text = row.get(config.get("output_field", "output"))
+        if not isinstance(instruction, str) or not isinstance(output_text, str):
+            raise ValueError("Не найдены поля instruction/output в Alpaca-подобном датасете.")
+        prompt = instruction.strip()
+        if isinstance(input_text, str) and input_text.strip():
+            prompt = f"{prompt}\n\n{input_text.strip()}"
+        return [
+            {"role": "user", "content": prompt},
+            {"role": "assistant", "content": output_text.strip()},
+        ]
+
+    raise ValueError(f"Unsupported SFT format: {fmt}")
+
+
+def tokenize_messages(
+    tokenizer: Tokenizer,
+    messages: list[dict[str, str]],
+    bos_id: int,
+    eos_id: int,
+) -> tuple[list[int], list[int]]:
+    input_ids = [bos_id]
+    labels = [-100]
+
+    for message in messages:
+        role = message["role"].strip().lower()
+        content = message["content"].strip()
+        if not content:
+            continue
+        text = f"<|{role}|>\n{content}\n"
+        piece = tokenizer.encode(text, add_special_tokens=False).ids
+        if not piece:
+            continue
+        input_ids.extend(piece)
+        if role == "assistant":
+            labels.extend(piece)
+        else:
+            labels.extend([-100] * len(piece))
+
+    input_ids.append(eos_id)
+    labels.append(eos_id)
+    return input_ids, labels
+
+
+def pad_or_truncate(
+    input_ids: list[int],
+    labels: list[int],
+    seq_len: int,
+    pad_id: int,
+) -> tuple[list[int], list[int]]:
+    input_ids = input_ids[:seq_len]
+    labels = labels[:seq_len]
+    if len(input_ids) < seq_len:
+        pad_length = seq_len - len(input_ids)
+        input_ids = input_ids + [pad_id] * pad_length
+        labels = labels + [-100] * pad_length
+    return input_ids, labels
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    config = load_json(args.config)
+    tokenizer, tokenizer_meta = load_tokenizer(args.tokenizer_dir)
+    specials = tokenizer_meta["special_tokens"]
+    bos_id = int(specials["bos_token_id"])
+    eos_id = int(specials["eos_token_id"])
+    pad_id = int(specials["pad_token_id"])
+
+    dataset = load_dataset(
+        path=config["path"],
+        name=config.get("config_name"),
+        split=config.get("split", "train"),
+        revision=config.get("revision"),
+        streaming=bool(config.get("streaming", False)),
+    )
+    if config.get("shuffle", True):
+        dataset = dataset.shuffle(seed=args.seed)
+
+    val_examples = int(config.get("val_examples", 1_000))
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    logger, log_path = setup_logger("sllm.prepare_sft_data", output_dir, "prepare_sft_data")
+    logger.info("SFT data preparation started")
+    logger.info("Log file: %s", log_path)
+    logger.info(
+        "Arguments | config=%s tokenizer_dir=%s output_dir=%s seq_len=%s seed=%s",
+        args.config,
+        args.tokenizer_dir,
+        args.output_dir,
+        args.seq_len,
+        args.seed,
+    )
+    logger.info(
+        "SFT source config | path=%s config_name=%s split=%s format=%s streaming=%s val_examples=%s max_train_examples=%s",
+        config.get("path"),
+        config.get("config_name"),
+        config.get("split", "train"),
+        config.get("format", "messages"),
+        bool(config.get("streaming", False)),
+        val_examples,
+        config.get("max_train_examples"),
+    )
+    train_writer = SFTShardWriter(output_dir, prefix="train", seq_len=args.seq_len)
+    val_writer = SFTShardWriter(output_dir, prefix="val", seq_len=args.seq_len)
+
+    train_count = 0
+    val_count = 0
+    max_train_examples = config.get("max_train_examples")
+
+    for row in dataset:
+        messages = row_to_messages(row, config)
+        if not messages:
+            continue
+        input_ids, labels = tokenize_messages(tokenizer, messages, bos_id=bos_id, eos_id=eos_id)
+        input_ids, labels = pad_or_truncate(input_ids, labels, args.seq_len, pad_id=pad_id)
+
+        if val_count < val_examples:
+            val_writer.add_example(input_ids, labels)
+            val_count += 1
+        else:
+            train_writer.add_example(input_ids, labels)
+            train_count += 1
+
+        total_examples = train_count + val_count
+        if total_examples % 5_000 == 0:
+            logger.info(
+                "SFT progress | processed=%s train_examples=%s val_examples=%s",
+                f"{total_examples:,}",
+                f"{train_count:,}",
+                f"{val_count:,}",
+            )
+
+        if max_train_examples is not None and train_count >= int(max_train_examples):
+            break
+
+    train_metadata = train_writer.finalize()
+    val_metadata = val_writer.finalize()
+    save_json(
+        output_dir / "dataset_summary.json",
+        {
+            "config": config,
+            "tokenizer_meta": tokenizer_meta,
+            "train": train_metadata,
+            "val": val_metadata,
+        },
+    )
+    logger.info("SFT dataset saved | output_dir=%s", output_dir)
+    logger.info("SFT summary | train_examples=%s val_examples=%s", f"{train_count:,}", f"{val_count:,}")
+    logger.info("SFT metadata saved | path=%s", output_dir / "dataset_summary.json")
+
+
+if __name__ == "__main__":
+    main()

scripts/train_pretrain.py

@@ -0,0 +1,405 @@
+from __future__ import annotations
+
+import argparse
+import math
+import sys
+import time
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.append(str(ROOT / "src"))
+
+from sllm.checkpoint import load_checkpoint, save_checkpoint
+from sllm.config import ModelConfig, TrainConfig, load_json, save_json
+from sllm.data import RandomTokenDataset, SequentialEvalDataset
+from sllm.model import SLLMForCausalLM
+from sllm.utils import (
+    append_jsonl,
+    autocast_context,
+    cosine_lr,
+    cuda_memory_snapshot,
+    ensure_dir,
+    format_number,
+    get_device,
+    iso_timestamp,
+    maybe_enable_tf32,
+    model_parameter_count,
+    resolve_runtime_precision,
+    set_optimizer_lr,
+    set_seed,
+    setup_logger,
+    timestamp,
+    tokens_per_step,
+)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Pretrain the small causal LM.")
+    parser.add_argument("--model-config", required=True, help="Path to model JSON config.")
+    parser.add_argument("--train-config", required=True, help="Path to pretraining JSON config.")
+    parser.add_argument(
+        "--max-steps",
+        type=int,
+        default=None,
+        help="Optional override for debugging or dry runs.",
+    )
+    return parser
+
+
+def build_optimizer(model: torch.nn.Module, config: TrainConfig, device: torch.device):
+    decay_params = []
+    no_decay_params = []
+    for name, parameter in model.named_parameters():
+        if not parameter.requires_grad:
+            continue
+        if parameter.ndim <= 1 or name.endswith("bias"):
+            no_decay_params.append(parameter)
+        else:
+            decay_params.append(parameter)
+
+    fused_supported = device.type == "cuda"
+    return torch.optim.AdamW(
+        [
+            {"params": decay_params, "weight_decay": config.weight_decay},
+            {"params": no_decay_params, "weight_decay": 0.0},
+        ],
+        lr=config.learning_rate,
+        betas=(config.beta1, config.beta2),
+        fused=fused_supported,
+    )
+
+
+@torch.no_grad()
+def evaluate(
+    model: SLLMForCausalLM,
+    config: TrainConfig,
+    device: torch.device,
+) -> tuple[float, float]:
+    model.eval()
+    dataset = SequentialEvalDataset(
+        data_dir=config.val_dir,
+        split="val",
+        seq_len=config.seq_len,
+        max_batches=config.eval_batches * config.micro_batch_size,
+    )
+    loader = DataLoader(dataset, batch_size=config.micro_batch_size, num_workers=0)
+
+    losses = []
+    for batch_index, batch in enumerate(loader):
+        if batch_index >= config.eval_batches:
+            break
+        batch = {key: value.to(device) for key, value in batch.items()}
+        with autocast_context(device, config.precision):
+            loss = model(**batch)["loss"]
+        losses.append(loss.detach().float().item())
+
+    mean_loss = float(sum(losses) / max(1, len(losses)))
+    perplexity = math.exp(min(mean_loss, 20))
+    model.train()
+    return mean_loss, perplexity
+
+
+def maybe_load_weights(
+    model: SLLMForCausalLM,
+    optimizer: torch.optim.Optimizer,
+    config: TrainConfig,
+    device: torch.device,
+    logger,
+) -> int:
+    step = 0
+    checkpoint_path = config.resume_from or config.init_from
+    if checkpoint_path is None:
+        return step
+
+    payload = load_checkpoint(checkpoint_path, map_location=device)
+    model.load_state_dict(payload["model"])
+    if config.resume_from and payload.get("optimizer") is not None:
+        optimizer.load_state_dict(payload["optimizer"])
+        step = int(payload.get("step", 0))
+        logger.info("Resumed training | step=%s checkpoint=%s", step, checkpoint_path)
+    else:
+        logger.info("Loaded model weights | checkpoint=%s", checkpoint_path)
+    return step
+
+
+def save_run_config(output_dir: Path, model_config: ModelConfig, train_config: TrainConfig) -> None:
+    save_json(
+        output_dir / "run_config.json",
+        {
+            "model_config": model_config.to_dict(),
+            "train_config": train_config.to_dict(),
+        },
+    )
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    model_config = ModelConfig.from_dict(load_json(args.model_config))
+    train_config = TrainConfig.from_dict(load_json(args.train_config))
+    if args.max_steps is not None:
+        train_config.max_steps = args.max_steps
+
+    set_seed(train_config.seed)
+    device = get_device()
+    maybe_enable_tf32(device)
+    runtime_precision, precision_warning = resolve_runtime_precision(device, train_config.precision)
+    train_config.precision = runtime_precision
+
+    output_dir = ensure_dir(train_config.output_dir)
+    checkpoint_dir = ensure_dir(train_config.checkpoint_dir)
+    logger, log_path = setup_logger("sllm.train_pretrain", output_dir, "train_pretrain")
+    metrics_path = Path(output_dir) / "logs" / f"{log_path.stem}.jsonl"
+    logger.info("Pretraining started")
+    logger.info("Log file: %s", log_path)
+    logger.info("Metrics JSONL: %s", metrics_path)
+    logger.info("Arguments | model_config=%s train_config=%s max_steps_override=%s", args.model_config, args.train_config, args.max_steps)
+    if precision_warning is not None:
+        logger.warning(precision_warning)
+    logger.info("Model config | %s", model_config.to_dict())
+    logger.info("Train config | %s", train_config.to_dict())
+    append_jsonl(
+        metrics_path,
+        {
+            "event": "run_started",
+            "timestamp": iso_timestamp(),
+            "log_path": str(log_path),
+            "metrics_path": str(metrics_path),
+            "model_config": model_config.to_dict(),
+            "train_config": train_config.to_dict(),
+            "args": {
+                "model_config": args.model_config,
+                "train_config": args.train_config,
+                "max_steps_override": args.max_steps,
+            },
+        },
+    )
+    save_run_config(output_dir, model_config, train_config)
+
+    dataset = RandomTokenDataset(
+        data_dir=train_config.train_dir,
+        split="train",
+        seq_len=train_config.seq_len,
+        seed=train_config.seed,
+    )
+    loader = DataLoader(
+        dataset,
+        batch_size=train_config.micro_batch_size,
+        num_workers=train_config.num_workers,
+        pin_memory=device.type == "cuda",
+    )
+    data_iter = iter(loader)
+
+    model = SLLMForCausalLM(model_config).to(device)
+    if train_config.compile_model and hasattr(torch, "compile"):
+        model = torch.compile(model)  # type: ignore[assignment]
+
+    optimizer = build_optimizer(model, train_config, device)
+    scaler = torch.amp.GradScaler(
+        "cuda",
+        enabled=device.type == "cuda" and train_config.precision.lower() == "fp16",
+    )
+    start_step = maybe_load_weights(model, optimizer, train_config, device, logger)
+    if start_step > 0:
+        append_jsonl(
+            metrics_path,
+            {
+                "event": "resumed",
+                "timestamp": iso_timestamp(),
+                "step": start_step,
+                "checkpoint": train_config.resume_from,
+            },
+        )
+    model.train()
+
+    tokens_step = tokens_per_step(
+        train_config.micro_batch_size,
+        train_config.grad_accum_steps,
+        train_config.seq_len,
+    )
+
+    logger.info("Device summary | device=%s precision=%s compile_model=%s", device, train_config.precision, train_config.compile_model)
+    logger.info("Model summary | parameters=%s", format_number(model_parameter_count(model)))
+    logger.info(
+        "Batch summary | seq_len=%s micro_batch_size=%s grad_accum_steps=%s tokens_per_step=%s",
+        train_config.seq_len,
+        train_config.micro_batch_size,
+        train_config.grad_accum_steps,
+        f"{tokens_step:,}",
+    )
+    logger.info("Dataset summary | train_dir=%s val_dir=%s num_train_shards=%s", train_config.train_dir, train_config.val_dir, len(dataset.shards))
+    append_jsonl(
+        metrics_path,
+        {
+            "event": "runtime_summary",
+            "timestamp": iso_timestamp(),
+            "device": str(device),
+            "precision": train_config.precision,
+            "compile_model": train_config.compile_model,
+            "parameters": model_parameter_count(model),
+            "seq_len": train_config.seq_len,
+            "micro_batch_size": train_config.micro_batch_size,
+            "grad_accum_steps": train_config.grad_accum_steps,
+            "tokens_per_step": tokens_step,
+            "num_train_shards": len(dataset.shards),
+            "train_dir": train_config.train_dir,
+            "val_dir": train_config.val_dir,
+        },
+    )
+
+    running_loss = 0.0
+    log_start_time = time.perf_counter()
+    last_grad_norm = float("nan")
+
+    for step in range(start_step, train_config.max_steps):
+        lr = cosine_lr(
+            step=step,
+            warmup_steps=train_config.warmup_steps,
+            max_steps=train_config.max_steps,
+            max_lr=train_config.learning_rate,
+            min_lr=train_config.min_lr,
+        )
+        set_optimizer_lr(optimizer, lr)
+        optimizer.zero_grad(set_to_none=True)
+
+        step_loss = 0.0
+        for micro_step in range(train_config.grad_accum_steps):
+            batch = next(data_iter)
+            batch = {key: value.to(device, non_blocking=device.type == "cuda") for key, value in batch.items()}
+
+            with autocast_context(device, train_config.precision):
+                loss = model(**batch)["loss"] / train_config.grad_accum_steps
+
+            step_loss += loss.detach().float().item()
+            if scaler.is_enabled():
+                scaler.scale(loss).backward()
+            else:
+                loss.backward()
+
+        if train_config.grad_clip is not None and train_config.grad_clip > 0:
+            if scaler.is_enabled():
+                scaler.unscale_(optimizer)
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), train_config.grad_clip)
+            last_grad_norm = float(grad_norm)
+
+        if scaler.is_enabled():
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            optimizer.step()
+
+        running_loss += step_loss
+
+        if (step + 1) % train_config.log_interval == 0:
+            elapsed = time.perf_counter() - log_start_time
+            avg_loss = running_loss / train_config.log_interval
+            tok_per_sec = (tokens_step * train_config.log_interval) / max(elapsed, 1e-6)
+            memory = cuda_memory_snapshot(device)
+            memory_suffix = ""
+            if memory:
+                memory_suffix = (
+                    f" mem_alloc_gb={memory['allocated_gb']:.2f}"
+                    f" mem_reserved_gb={memory['reserved_gb']:.2f}"
+                    f" max_mem_alloc_gb={memory['max_allocated_gb']:.2f}"
+                    f" max_mem_reserved_gb={memory['max_reserved_gb']:.2f}"
+                )
+            logger.info(
+                "Train step | step=%s loss=%.4f lr=%.6f tok_per_sec=%s grad_norm=%.4f tokens_seen=%s%s",
+                step + 1,
+                avg_loss,
+                lr,
+                f"{tok_per_sec:,.0f}",
+                last_grad_norm,
+                format_number((step + 1) * tokens_step),
+                memory_suffix,
+            )
+            append_jsonl(
+                metrics_path,
+                {
+                    "event": "train",
+                    "timestamp": iso_timestamp(),
+                    "step": step + 1,
+                    "loss": avg_loss,
+                    "lr": lr,
+                    "tok_per_sec": tok_per_sec,
+                    "grad_norm": last_grad_norm,
+                    "tokens_seen": (step + 1) * tokens_step,
+                    "elapsed_sec": elapsed,
+                    "seq_len": train_config.seq_len,
+                    "micro_batch_size": train_config.micro_batch_size,
+                    "grad_accum_steps": train_config.grad_accum_steps,
+                    **memory,
+                },
+            )
+            running_loss = 0.0
+            log_start_time = time.perf_counter()
+
+        if (step + 1) % train_config.eval_interval == 0:
+            val_loss, perplexity = evaluate(model, train_config, device)
+            logger.info("Eval step | step=%s val_loss=%.4f perplexity=%.2f", step + 1, val_loss, perplexity)
+            append_jsonl(
+                metrics_path,
+                {
+                    "event": "eval",
+                    "timestamp": iso_timestamp(),
+                    "step": step + 1,
+                    "val_loss": val_loss,
+                    "perplexity": perplexity,
+                    "eval_batches": train_config.eval_batches,
+                },
+            )
+
+        if (step + 1) % train_config.save_interval == 0 or (step + 1) == train_config.max_steps:
+            step_checkpoint_path = checkpoint_dir / f"step_{step + 1:07d}.pt"
+            last_checkpoint_path = checkpoint_dir / "last.pt"
+            save_checkpoint(
+                step_checkpoint_path,
+                model=model,
+                optimizer=optimizer,
+                step=step + 1,
+                model_config=model_config.to_dict(),
+                train_config=train_config.to_dict(),
+                extra_state={"tokens_seen": (step + 1) * tokens_step},
+            )
+            save_checkpoint(
+                last_checkpoint_path,
+                model=model,
+                optimizer=optimizer,
+                step=step + 1,
+                model_config=model_config.to_dict(),
+                train_config=train_config.to_dict(),
+                extra_state={"tokens_seen": (step + 1) * tokens_step},
+            )
+            logger.info(
+                "Checkpoint saved | step=%s step_checkpoint=%s last_checkpoint=%s",
+                step + 1,
+                step_checkpoint_path,
+                last_checkpoint_path,
+            )
+            append_jsonl(
+                metrics_path,
+                {
+                    "event": "checkpoint",
+                    "timestamp": iso_timestamp(),
+                    "step": step + 1,
+                    "step_checkpoint": str(step_checkpoint_path),
+                    "last_checkpoint": str(last_checkpoint_path),
+                    "tokens_seen": (step + 1) * tokens_step,
+                },
+            )
+
+    append_jsonl(
+        metrics_path,
+        {
+            "event": "run_finished",
+            "timestamp": iso_timestamp(),
+            "final_step": train_config.max_steps,
+            "tokens_seen": train_config.max_steps * tokens_step,
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/train_sft.py

@@ -0,0 +1,394 @@
+from __future__ import annotations
+
+import argparse
+import math
+import sys
+import time
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.append(str(ROOT / "src"))
+
+from sllm.checkpoint import load_checkpoint, save_checkpoint
+from sllm.config import ModelConfig, SFTConfig, load_json, save_json
+from sllm.data import FixedSFTDataset
+from sllm.model import SLLMForCausalLM
+from sllm.utils import (
+    append_jsonl,
+    autocast_context,
+    cosine_lr,
+    cuda_memory_snapshot,
+    ensure_dir,
+    format_number,
+    get_device,
+    iso_timestamp,
+    maybe_enable_tf32,
+    model_parameter_count,
+    resolve_runtime_precision,
+    set_optimizer_lr,
+    set_seed,
+    setup_logger,
+    timestamp,
+    tokens_per_step,
+)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run supervised fine-tuning for the sLLM.")
+    parser.add_argument("--model-config", required=True, help="Path to model JSON config.")
+    parser.add_argument("--train-config", required=True, help="Path to SFT JSON config.")
+    parser.add_argument("--max-steps", type=int, default=None, help="Optional debug override.")
+    return parser
+
+
+def build_optimizer(model: torch.nn.Module, config: SFTConfig, device: torch.device):
+    decay_params = []
+    no_decay_params = []
+    for name, parameter in model.named_parameters():
+        if not parameter.requires_grad:
+            continue
+        if parameter.ndim <= 1 or name.endswith("bias"):
+            no_decay_params.append(parameter)
+        else:
+            decay_params.append(parameter)
+    return torch.optim.AdamW(
+        [
+            {"params": decay_params, "weight_decay": config.weight_decay},
+            {"params": no_decay_params, "weight_decay": 0.0},
+        ],
+        lr=config.learning_rate,
+        betas=(config.beta1, config.beta2),
+        fused=device.type == "cuda",
+    )
+
+
+@torch.no_grad()
+def evaluate(model: SLLMForCausalLM, loader: DataLoader, device: torch.device, precision: str, max_batches: int):
+    model.eval()
+    losses = []
+    for batch_index, batch in enumerate(loader):
+        if batch_index >= max_batches:
+            break
+        batch = {key: value.to(device) for key, value in batch.items()}
+        with autocast_context(device, precision):
+            loss = model(**batch)["loss"]
+        losses.append(loss.detach().float().item())
+    model.train()
+    mean_loss = float(sum(losses) / max(1, len(losses)))
+    return mean_loss, math.exp(min(mean_loss, 20))
+
+
+def save_run_config(output_dir: Path, model_config: ModelConfig, train_config: SFTConfig) -> None:
+    save_json(
+        output_dir / "run_config.json",
+        {
+            "model_config": model_config.to_dict(),
+            "train_config": train_config.to_dict(),
+        },
+    )
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    model_config = ModelConfig.from_dict(load_json(args.model_config))
+    train_config = SFTConfig.from_dict(load_json(args.train_config))
+    if args.max_steps is not None:
+        train_config.max_steps = args.max_steps
+
+    set_seed(train_config.seed)
+    device = get_device()
+    maybe_enable_tf32(device)
+    runtime_precision, precision_warning = resolve_runtime_precision(device, train_config.precision)
+    train_config.precision = runtime_precision
+
+    output_dir = ensure_dir(train_config.output_dir)
+    checkpoint_dir = ensure_dir(train_config.checkpoint_dir)
+    logger, log_path = setup_logger("sllm.train_sft", output_dir, "train_sft")
+    metrics_path = Path(output_dir) / "logs" / f"{log_path.stem}.jsonl"
+    logger.info("SFT training started")
+    logger.info("Log file: %s", log_path)
+    logger.info("Metrics JSONL: %s", metrics_path)
+    logger.info("Arguments | model_config=%s train_config=%s max_steps_override=%s", args.model_config, args.train_config, args.max_steps)
+    if precision_warning is not None:
+        logger.warning(precision_warning)
+    logger.info("Model config | %s", model_config.to_dict())
+    logger.info("SFT config | %s", train_config.to_dict())
+    append_jsonl(
+        metrics_path,
+        {
+            "event": "run_started",
+            "timestamp": iso_timestamp(),
+            "log_path": str(log_path),
+            "metrics_path": str(metrics_path),
+            "model_config": model_config.to_dict(),
+            "train_config": train_config.to_dict(),
+            "args": {
+                "model_config": args.model_config,
+                "train_config": args.train_config,
+                "max_steps_override": args.max_steps,
+            },
+        },
+    )
+    save_run_config(output_dir, model_config, train_config)
+
+    train_dataset = FixedSFTDataset(train_config.dataset_path, split="train")
+    val_dataset = FixedSFTDataset(train_config.dataset_path, split="val")
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=train_config.micro_batch_size,
+        shuffle=True,
+        num_workers=train_config.num_workers,
+        pin_memory=device.type == "cuda",
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=train_config.micro_batch_size,
+        shuffle=False,
+        num_workers=0,
+        pin_memory=device.type == "cuda",
+    )
+
+    model = SLLMForCausalLM(model_config).to(device)
+    if train_config.compile_model and hasattr(torch, "compile"):
+        model = torch.compile(model)  # type: ignore[assignment]
+
+    optimizer = build_optimizer(model, train_config, device)
+    scaler = torch.amp.GradScaler(
+        "cuda",
+        enabled=device.type == "cuda" and train_config.precision.lower() == "fp16",
+    )
+
+    start_step = 0
+    checkpoint_path = train_config.resume_from or train_config.init_from
+    if checkpoint_path:
+        payload = load_checkpoint(checkpoint_path, map_location=device)
+        model.load_state_dict(payload["model"])
+        if train_config.resume_from and payload.get("optimizer") is not None:
+            optimizer.load_state_dict(payload["optimizer"])
+            start_step = int(payload.get("step", 0))
+            logger.info("Resumed SFT | step=%s checkpoint=%s", start_step, checkpoint_path)
+            append_jsonl(
+                metrics_path,
+                {
+                    "event": "resumed",
+                    "timestamp": iso_timestamp(),
+                    "step": start_step,
+                    "checkpoint": checkpoint_path,
+                },
+            )
+        else:
+            logger.info("Loaded initialization weights | checkpoint=%s", checkpoint_path)
+            append_jsonl(
+                metrics_path,
+                {
+                    "event": "initialized_from_checkpoint",
+                    "timestamp": iso_timestamp(),
+                    "checkpoint": checkpoint_path,
+                },
+            )
+
+    model.train()
+    tokens_step = tokens_per_step(
+        train_config.micro_batch_size,
+        train_config.grad_accum_steps,
+        train_config.seq_len,
+    )
+    logger.info("Device summary | device=%s precision=%s compile_model=%s", device, train_config.precision, train_config.compile_model)
+    logger.info("Model summary | parameters=%s", format_number(model_parameter_count(model)))
+    logger.info(
+        "Batch summary | seq_len=%s micro_batch_size=%s grad_accum_steps=%s tokens_per_step=%s",
+        train_config.seq_len,
+        train_config.micro_batch_size,
+        train_config.grad_accum_steps,
+        f"{tokens_step:,}",
+    )
+    logger.info(
+        "Dataset summary | dataset_path=%s train_examples=%s val_examples=%s",
+        train_config.dataset_path,
+        len(train_dataset),
+        len(val_dataset),
+    )
+    append_jsonl(
+        metrics_path,
+        {
+            "event": "runtime_summary",
+            "timestamp": iso_timestamp(),
+            "device": str(device),
+            "precision": train_config.precision,
+            "compile_model": train_config.compile_model,
+            "parameters": model_parameter_count(model),
+            "seq_len": train_config.seq_len,
+            "micro_batch_size": train_config.micro_batch_size,
+            "grad_accum_steps": train_config.grad_accum_steps,
+            "tokens_per_step": tokens_step,
+            "dataset_path": train_config.dataset_path,
+            "train_examples": len(train_dataset),
+            "val_examples": len(val_dataset),
+        },
+    )
+    running_loss = 0.0
+    log_start_time = time.perf_counter()
+    train_iterator = iter(train_loader)
+    last_grad_norm = float("nan")
+
+    for step in range(start_step, train_config.max_steps):
+        lr = cosine_lr(
+            step=step,
+            warmup_steps=train_config.warmup_steps,
+            max_steps=train_config.max_steps,
+            max_lr=train_config.learning_rate,
+            min_lr=train_config.min_lr,
+        )
+        set_optimizer_lr(optimizer, lr)
+        optimizer.zero_grad(set_to_none=True)
+
+        step_loss = 0.0
+        for _ in range(train_config.grad_accum_steps):
+            try:
+                batch = next(train_iterator)
+            except StopIteration:
+                train_iterator = iter(train_loader)
+                batch = next(train_iterator)
+
+            batch = {key: value.to(device, non_blocking=device.type == "cuda") for key, value in batch.items()}
+            with autocast_context(device, train_config.precision):
+                loss = model(**batch)["loss"] / train_config.grad_accum_steps
+            step_loss += loss.detach().float().item()
+            if scaler.is_enabled():
+                scaler.scale(loss).backward()
+            else:
+                loss.backward()
+
+        if train_config.grad_clip and train_config.grad_clip > 0:
+            if scaler.is_enabled():
+                scaler.unscale_(optimizer)
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), train_config.grad_clip)
+            last_grad_norm = float(grad_norm)
+
+        if scaler.is_enabled():
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            optimizer.step()
+
+        running_loss += step_loss
+
+        if (step + 1) % train_config.log_interval == 0:
+            elapsed = time.perf_counter() - log_start_time
+            avg_loss = running_loss / train_config.log_interval
+            tok_per_sec = (tokens_step * train_config.log_interval) / max(elapsed, 1e-6)
+            memory = cuda_memory_snapshot(device)
+            memory_suffix = ""
+            if memory:
+                memory_suffix = (
+                    f" mem_alloc_gb={memory['allocated_gb']:.2f}"
+                    f" mem_reserved_gb={memory['reserved_gb']:.2f}"
+                    f" max_mem_alloc_gb={memory['max_allocated_gb']:.2f}"
+                    f" max_mem_reserved_gb={memory['max_reserved_gb']:.2f}"
+                )
+            logger.info(
+                "Train step | step=%s loss=%.4f lr=%.6f tok_per_sec=%s grad_norm=%.4f%s",
+                step + 1,
+                avg_loss,
+                lr,
+                f"{tok_per_sec:,.0f}",
+                last_grad_norm,
+                memory_suffix,
+            )
+            append_jsonl(
+                metrics_path,
+                {
+                    "event": "train",
+                    "timestamp": iso_timestamp(),
+                    "step": step + 1,
+                    "loss": avg_loss,
+                    "lr": lr,
+                    "tok_per_sec": tok_per_sec,
+                    "grad_norm": last_grad_norm,
+                    "tokens_seen": (step + 1) * tokens_step,
+                    "elapsed_sec": elapsed,
+                    "seq_len": train_config.seq_len,
+                    "micro_batch_size": train_config.micro_batch_size,
+                    "grad_accum_steps": train_config.grad_accum_steps,
+                    **memory,
+                },
+            )
+            running_loss = 0.0
+            log_start_time = time.perf_counter()
+
+        if (step + 1) % train_config.eval_interval == 0:
+            val_loss, val_ppl = evaluate(
+                model=model,
+                loader=val_loader,
+                device=device,
+                precision=train_config.precision,
+                max_batches=train_config.eval_batches,
+            )
+            logger.info("Eval step | step=%s val_loss=%.4f perplexity=%.2f", step + 1, val_loss, val_ppl)
+            append_jsonl(
+                metrics_path,
+                {
+                    "event": "eval",
+                    "timestamp": iso_timestamp(),
+                    "step": step + 1,
+                    "val_loss": val_loss,
+                    "perplexity": val_ppl,
+                    "eval_batches": train_config.eval_batches,
+                },
+            )
+
+        if (step + 1) % train_config.save_interval == 0 or (step + 1) == train_config.max_steps:
+            step_checkpoint_path = checkpoint_dir / f"step_{step + 1:07d}.pt"
+            last_checkpoint_path = checkpoint_dir / "last.pt"
+            save_checkpoint(
+                step_checkpoint_path,
+                model=model,
+                optimizer=optimizer,
+                step=step + 1,
+                model_config=model_config.to_dict(),
+                train_config=train_config.to_dict(),
+                extra_state={"tokens_seen": (step + 1) * tokens_step},
+            )
+            save_checkpoint(
+                last_checkpoint_path,
+                model=model,
+                optimizer=optimizer,
+                step=step + 1,
+                model_config=model_config.to_dict(),
+                train_config=train_config.to_dict(),
+                extra_state={"tokens_seen": (step + 1) * tokens_step},
+            )
+            logger.info(
+                "Checkpoint saved | step=%s step_checkpoint=%s last_checkpoint=%s",
+                step + 1,
+                step_checkpoint_path,
+                last_checkpoint_path,
+            )
+            append_jsonl(
+                metrics_path,
+                {
+                    "event": "checkpoint",
+                    "timestamp": iso_timestamp(),
+                    "step": step + 1,
+                    "step_checkpoint": str(step_checkpoint_path),
+                    "last_checkpoint": str(last_checkpoint_path),
+                    "tokens_seen": (step + 1) * tokens_step,
+                },
+            )
+
+    append_jsonl(
+        metrics_path,
+        {
+            "event": "run_finished",
+            "timestamp": iso_timestamp(),
+            "final_step": train_config.max_steps,
+            "tokens_seen": train_config.max_steps * tokens_step,
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/train_tokenizer.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Iterator
+
+from datasets import load_dataset
+from tokenizers import Tokenizer, decoders, models, pre_tokenizers, processors, trainers
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.append(str(ROOT / "src"))
+
+from sllm.config import DataMixConfig, load_json, save_json
+from sllm.utils import setup_logger
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Train a BPE tokenizer for the sLLM pipeline.")
+    parser.add_argument("--data-config", required=True, help="Path to data mixture JSON config.")
+    parser.add_argument("--output-dir", required=True, help="Directory where tokenizer files will be stored.")
+    parser.add_argument("--vocab-size", type=int, default=49_152, help="Target tokenizer vocabulary size.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for dataset shuffling.")
+    return parser
+
+
+def iter_source_texts(source, seed: int, limit: int) -> Iterator[str]:
+    dataset = load_dataset(
+        path=source.path,
+        name=source.config_name,
+        data_dir=source.data_dir,
+        split=source.split,
+        revision=source.revision,
+        streaming=source.streaming,
+    )
+    if source.streaming:
+        dataset = dataset.shuffle(seed=seed, buffer_size=source.shuffle_buffer)
+
+    yielded = 0
+    for row in dataset:
+        text = row.get(source.text_field or "", None)
+        if not isinstance(text, str):
+            continue
+        text = text.strip()
+        if not text:
+            continue
+        yield text
+        yielded += 1
+        if yielded >= limit:
+            return
+
+
+def mixed_iterator(config: DataMixConfig, seed: int, logger) -> Iterator[str]:
+    weight_map = config.normalized_weights()
+    total_docs = config.tokenizer_sample_documents
+    per_source = {
+        source.name: max(1, int(total_docs * weight_map[source.name]))
+        for source in config.sources
+    }
+
+    for index, source in enumerate(config.sources):
+        limit = source.sample_documents or per_source[source.name]
+        logger.info(
+            "Tokenizer source start | name=%s path=%s data_dir=%s split=%s text_field=%s limit_docs=%s streaming=%s",
+            source.name,
+            source.path,
+            source.data_dir,
+            source.split,
+            source.text_field,
+            f"{limit:,}",
+            source.streaming,
+        )
+        yield from iter_source_texts(source, seed + index, limit)
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    data_config = DataMixConfig.from_dict(load_json(args.data_config))
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    logger, log_path = setup_logger("sllm.train_tokenizer", output_dir, "train_tokenizer")
+    logger.info("Tokenizer training started")
+    logger.info("Log file: %s", log_path)
+    logger.info("Arguments | data_config=%s output_dir=%s vocab_size=%s seed=%s", args.data_config, args.output_dir, args.vocab_size, args.seed)
+    logger.info("Tokenizer config | sample_documents=%s min_frequency=%s special_tokens=%s num_sources=%s", f"{data_config.tokenizer_sample_documents:,}", data_config.tokenizer_min_frequency, data_config.tokenizer_special_tokens, len(data_config.sources))
+
+    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    tokenizer.decoder = decoders.ByteLevel()
+    trainer = trainers.BpeTrainer(
+        vocab_size=args.vocab_size,
+        min_frequency=data_config.tokenizer_min_frequency,
+        special_tokens=data_config.tokenizer_special_tokens,
+        show_progress=True,
+    )
+    tokenizer.train_from_iterator(mixed_iterator(data_config, args.seed, logger), trainer=trainer)
+
+    bos_id = tokenizer.token_to_id("<bos>")
+    eos_id = tokenizer.token_to_id("<eos>")
+    pad_id = tokenizer.token_to_id("<pad>")
+    if bos_id is None or eos_id is None or pad_id is None:
+        raise RuntimeError("Tokenizer special tokens were not created correctly.")
+
+    tokenizer.post_processor = processors.TemplateProcessing(
+        single="<bos> $A <eos>",
+        pair="<bos> $A <eos> $B:1 <eos>:1",
+        special_tokens=[
+            ("<bos>", bos_id),
+            ("<eos>", eos_id),
+        ],
+    )
+
+    tokenizer_path = output_dir / "tokenizer.json"
+    tokenizer.save(str(tokenizer_path))
+
+    metadata = {
+        "vocab_size": tokenizer.get_vocab_size(),
+        "special_tokens": {
+            "pad_token": "<pad>",
+            "bos_token": "<bos>",
+            "eos_token": "<eos>",
+            "unk_token": "<unk>",
+            "pad_token_id": pad_id,
+            "bos_token_id": bos_id,
+            "eos_token_id": eos_id,
+            "unk_token_id": tokenizer.token_to_id("<unk>"),
+        },
+        "data_config": data_config.to_dict(),
+    }
+    save_json(output_dir / "tokenizer_meta.json", metadata)
+
+    with (output_dir / "tokenizer_summary.json").open("w", encoding="utf-8") as handle:
+        json.dump(metadata, handle, ensure_ascii=False, indent=2)
+
+    logger.info("Tokenizer saved | path=%s", tokenizer_path)
+    logger.info(
+        "Tokenizer summary | vocab_size=%s pad_id=%s bos_id=%s eos_id=%s unk_id=%s",
+        tokenizer.get_vocab_size(),
+        pad_id,
+        bos_id,
+        eos_id,
+        tokenizer.token_to_id("<unk>"),
+    )
+    logger.info("Tokenizer metadata saved | path=%s", output_dir / "tokenizer_meta.json")
+
+
+if __name__ == "__main__":
+    main()