Upload transformer model after 3120 steps (perplexity: 6.170201301574707)

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model_checkpoint/default/ocdbt.process_0/d/7cda073904d823e88c6221b1b1753d92 filter=lfs diff=lfs merge=lfs -text
+model_checkpoint/default/ocdbt.process_0/d/8eb7ace34ec179a3f15dc99655709dee filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,97 @@
+---
+license: mit
+tags:
+- transformer
+- autoregressive
+- language-modeling
+- tinystories
+- jax
+- flax
+datasets:
+- roneneldan/TinyStories
+language:
+- en
+pipeline_tag: text-generation
+---
+
+# Autoregressive Transformer trained on TinyStories
+
+This is an autoregressive decoder-only transformer model trained on the TinyStories dataset using JAX and Flax NNX.
+
+## Model Details
+
+- **Model Type**: Autoregressive Decoder-only Transformer
+- **Framework**: JAX + Flax NNX
+- **Dataset**: TinyStories
+- **Parameters**: ~85.0M
+- **Precision**: Mixed (FP32 parameters, BF16 computation)
+
+## Architecture
+
+```
+- Hidden Size: 512
+- Number of Layers: 8
+- Attention Heads: 8
+- Intermediate Size: 2048
+- Max Position Embeddings: 256
+- Vocab Size: 50257
+- Rotary Position Embeddings: True
+```
+
+## Training Details
+
+- **Training Steps**: 3,120
+- **Batch Size**: 32
+- **Gradient Accumulation**: 4
+- **Learning Rate**: 0.0003
+- **Training Duration**: 0.43 hours
+- **Final Eval Loss**: 1.7965960502624512
+- **Final Eval Perplexity**: 6.170201301574707
+
+## Usage
+
+```python
+# This model was trained with JAX/Flax and requires the custom transformer implementation
+# to load and use. See the repository for implementation details.
+
+from transformers import AutoTokenizer
+import jax.numpy as jnp
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# Example text generation (requires custom model loading)
+prompt = "Once upon a time, there was a little"
+# ... (model loading and generation code)
+```
+
+## Training Configuration
+
+```yaml
+model:
+  hidden_size: 512
+  num_layers: 8  
+  num_attention_heads: 8
+  intermediate_size: 2048
+  max_position_embeddings: 256
+
+training:
+  learning_rate: 0.0003
+  batch_size: 32
+  epochs: 10
+  warmup_ratio: 0.1
+```
+
+## Files
+
+- `config.json`: Model configuration
+- `train_history.json`: Training metrics and duration
+- `tokenizer/`: GPT-2 tokenizer files  
+- `model_checkpoint/`: Best model checkpoint
+- `tensorboard_logs/`: Training logs for TensorBoard
+
+## License
+
+MIT License - see LICENSE file for details.

config.json

@@ -0,0 +1,19 @@
+{
+  "vocab_size": 50257,
+  "hidden_size": 512,
+  "num_layers": 8,
+  "num_attention_heads": 8,
+  "intermediate_size": 2048,
+  "max_position_embeddings": 256,
+  "attention_dropout": 0.1,
+  "use_rotary_pos_emb": true,
+  "rope_theta": 10000.0,
+  "hidden_dropout": 0.1,
+  "layer_norm_eps": 1e-05,
+  "use_bias": false,
+  "initializer_range": 0.02,
+  "pad_token_id": 50256,
+  "bos_token_id": 50256,
+  "eos_token_id": 50256,
+  "use_cache": true
+}

model_checkpoint/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"default": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "metrics": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1754355577302929106, "commit_timestamp_nsecs": 1754355580447202487, "custom_metadata": {}}

model_checkpoint/default/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('embed_tokens', 'embedding', 'value')": {"key_metadata": [{"key": "embed_tokens", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [50257, 64]}}, "('layers', 'input_layernorm', 'weight', 'value')": {"key_metadata": [{"key": "layers", "key_type": 2}, {"key": "input_layernorm", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 512]}}, "('layers', 'mlp', 'down_proj', 'kernel', 'value')": {"key_metadata": [{"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "down_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 2048, 512]}}, "('layers', 'mlp', 'gate_proj', 'kernel', 'value')": {"key_metadata": [{"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gate_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 2048]}}, "('layers', 'mlp', 'up_proj', 'kernel', 'value')": {"key_metadata": [{"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "up_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 2048]}}, "('layers', 'post_attention_layernorm', 'weight', 'value')": {"key_metadata": [{"key": "layers", "key_type": 2}, {"key": "post_attention_layernorm", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 512]}}, "('layers', 'self_attn', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "layers", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('layers', 'self_attn', 'o_proj', 'kernel', 'value')": {"key_metadata": [{"key": "layers", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "o_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 512, 512]}}, "('layers', 'self_attn', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "layers", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('layers', 'self_attn', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "layers", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 64, 512]}}, "('lm_head', 'kernel', 'value')": {"key_metadata": [{"key": "lm_head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 50257]}}, "('norm', 'weight', 'value')": {"key_metadata": [{"key": "norm", "key_type": 2}, {"key": "weight", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

model_checkpoint/default/_sharding

@@ -0,0 +1 @@
+{"ZW1iZWRfdG9rZW5zLmVtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bG1faGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [\"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bGF5ZXJzLm1scC51cF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bGF5ZXJzLm1scC5kb3duX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [\"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bGF5ZXJzLm1scC5nYXRlX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bGF5ZXJzLmlucHV0X2xheWVybm9ybS53ZWlnaHQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [\"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bGF5ZXJzLnBvc3RfYXR0ZW50aW9uX2xheWVybm9ybS53ZWlnaHQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [\"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bGF5ZXJzLnNlbGZfYXR0bi52X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bGF5ZXJzLnNlbGZfYXR0bi5rX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bGF5ZXJzLnNlbGZfYXR0bi5vX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [\"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bGF5ZXJzLnNlbGZfYXR0bi5xX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}","bm9ybS53ZWlnaHQudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [\"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 6}, {\"id\": 7}, {\"id\": 4}, {\"id\": 5}]]}}"}

model_checkpoint/default/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "embed_tokens.embedding.value", "write_shape": [50257, 64], "chunk_shape": [50257, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "layers.input_layernorm.weight.value", "write_shape": [1, 512], "chunk_shape": [1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "layers.mlp.down_proj.kernel.value", "write_shape": [1, 2048, 512], "chunk_shape": [1, 2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "layers.mlp.gate_proj.kernel.value", "write_shape": [8, 64, 2048], "chunk_shape": [8, 64, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "layers.mlp.up_proj.kernel.value", "write_shape": [8, 64, 2048], "chunk_shape": [8, 64, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "layers.post_attention_layernorm.weight.value", "write_shape": [1, 512], "chunk_shape": [1, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "layers.self_attn.k_proj.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "layers.self_attn.o_proj.kernel.value", "write_shape": [1, 512, 512], "chunk_shape": [1, 512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "layers.self_attn.q_proj.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "layers.self_attn.v_proj.kernel.value", "write_shape": [8, 64, 512], "chunk_shape": [8, 64, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "lm_head.kernel.value", "write_shape": [64, 50257], "chunk_shape": [64, 50257], "ext_metadata": null}}, {"array_metadata": {"param_name": "norm.weight.value", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}]}

model_checkpoint/default/ocdbt.process_0/d/7cda073904d823e88c6221b1b1753d92

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

model_checkpoint/default/ocdbt.process_0/d/8eb7ace34ec179a3f15dc99655709dee

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

model_checkpoint/metrics/metrics

@@ -0,0 +1 @@
+{"eval_loss": 1.7965960502624512}

tensorboard_logs/events.out.tfevents.1754354033.047df4d1a242.2109.0.v2

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

tokenizer/special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer/tokenizer_config.json

@@ -0,0 +1,21 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

train_history.json

@@ -0,0 +1,8 @@
+{
+  "total_training_duration": 1547.211970743,
+  "avg_epoch_duration": 154.72109574149997,
+  "num_epochs_completed": 10,
+  "global_steps": 3120,
+  "global_optimizer_steps": 780,
+  "total_parameters": 85026304
+}