xLSTM is ready for image generation

.gitattributes

@@ -33,3 +33,7 @@ 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
+logs/10000/default/ocdbt.process_0/d/399bbb6a26e3af7391ef4ee2c557d20e filter=lfs diff=lfs merge=lfs -text
+logs/10000/default/ocdbt.process_0/d/f957dd57345bd22cb0916bf0f123bcac filter=lfs diff=lfs merge=lfs -text
+logs/8000/default/ocdbt.process_0/d/2f9047fa57561a8b64e87e5c02a78301 filter=lfs diff=lfs merge=lfs -text
+logs/8000/default/ocdbt.process_0/d/b07075d27dbfe48c28f624fcfc889e48 filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,81 @@
+{
+    "mlstm_block": {
+        "mlstm": {
+            "proj_factor": 2.0,
+            "round_proj_up_dim_up": true,
+            "round_proj_up_to_multiple_of": 64,
+            "_proj_up_dim": 768,
+            "conv1d_kernel_size": 4,
+            "qkv_proj_blocksize": 32,
+            "num_heads": 4,
+            "embedding_dim": 384,
+            "bias": false,
+            "dropout": 0.0,
+            "context_length": 256,
+            "_num_blocks": 6,
+            "_inner_embedding_dim": 768
+        },
+        "_num_blocks": 6,
+        "_block_idx": null
+    },
+    "slstm_block": {
+        "slstm": {
+            "hidden_size": 384,
+            "num_heads": 4,
+            "num_states": 4,
+            "backend": "vanilla",
+            "function": "slstm",
+            "bias_init": "powerlaw_blockdependent",
+            "recurrent_weight_init": "zeros",
+            "_block_idx": null,
+            "_num_blocks": 6,
+            "num_gates": 4,
+            "gradient_recurrent_clipval": null,
+            "forward_clipval": null,
+            "batch_size": 8,
+            "input_shape": "BSGNH",
+            "internal_input_shape": "SBNGH",
+            "output_shape": "BNSH",
+            "dtype": "bfloat16",
+            "dtype_b": "float32",
+            "dtype_r": "bfloat16",
+            "dtype_w": "bfloat16",
+            "dtype_g": "bfloat16",
+            "dtype_s": "bfloat16",
+            "dtype_a": "float32",
+            "initial_val": 0.0,
+            "enable_automatic_mixed_precision": true,
+            "embedding_dim": 384,
+            "conv1d_kernel_size": 4,
+            "group_norm_weight": true,
+            "dropout": 0.0
+        },
+        "feedforward": {
+            "proj_factor": 1.3,
+            "round_proj_up_dim_up": true,
+            "round_proj_up_to_multiple_of": 64,
+            "_proj_up_dim": 0,
+            "act_fn": "swish",
+            "embedding_dim": -1,
+            "dropout": 0.0,
+            "bias": false,
+            "ff_type": "ffn_gated",
+            "_num_blocks": 1
+        },
+        "_num_blocks": 6,
+        "_block_idx": null
+    },
+    "context_length": 256,
+    "num_blocks": 6,
+    "embedding_dim": 384,
+    "add_post_blocks_norm": true,
+    "bias": false,
+    "dropout": 0.0,
+    "slstm_at": [],
+    "_block_map": "0,0,0,0,0,0",
+    "vocab_size": 49152,
+    "tie_weights": false,
+    "weight_decay_on_embedding": false,
+    "add_embedding_dropout": false,
+    "pad_token_id": 0
+}

logs/10000/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1761148034520662925, "commit_timestamp_nsecs": 1761148035000077197, "custom_metadata": {}}

logs/10000/default/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('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": [384, 49152]}}, "('token_embedding', 'embedding', 'value')": {"key_metadata": [{"key": "token_embedding", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [49152, 384]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'conv1d', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "conv1d", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 768]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'conv1d', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "conv1d", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 4, 1, 768]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "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": [6, 24, 32, 32]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'learnable_skip', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "learnable_skip", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 768]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'fgate', 'bias', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "fgate", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 4]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'fgate', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "fgate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 2304, 4]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'igate', 'bias', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "igate", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 4]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'igate', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "igate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 2304, 4]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'outnorm', 'scale', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "outnorm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 768]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'proj_down', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "proj_down", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 768, 384]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'proj_up', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "proj_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 384, 1536]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "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": [6, 24, 32, 32]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "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": [6, 24, 32, 32]}}, "('xlstm_block_stack', 'blocks', 'xlstm_norm', 'scale', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 384]}}, "('xlstm_block_stack', 'post_blocks_norm', 'scale', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "post_blocks_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

logs/10000/default/_sharding

@@ -0,0 +1 @@
+{"bG1faGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dG9rZW5fZW1iZWRkaW5nLmVtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwuZmdhdGUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwuZmdhdGUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwuaWdhdGUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwuaWdhdGUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwub3V0bm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLmxlYXJuYWJsZV9za2lwLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLnByb2pfZG93bi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLnByb2pfdXAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2sucG9zdF9ibG9ja3Nfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"}

logs/10000/default/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "lm_head.kernel.value", "write_shape": [384, 49152], "chunk_shape": [384, 49152], "ext_metadata": null}}, {"array_metadata": {"param_name": "token_embedding.embedding.value", "write_shape": [49152, 384], "chunk_shape": [49152, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.conv1d.conv.bias.value", "write_shape": [6, 768], "chunk_shape": [6, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.conv1d.conv.kernel.value", "write_shape": [6, 4, 1, 768], "chunk_shape": [6, 4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.k_proj.kernel.value", "write_shape": [6, 24, 32, 32], "chunk_shape": [6, 24, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.learnable_skip.value", "write_shape": [6, 768], "chunk_shape": [6, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.fgate.bias.value", "write_shape": [6, 4], "chunk_shape": [6, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.fgate.kernel.value", "write_shape": [6, 2304, 4], "chunk_shape": [6, 2304, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.igate.bias.value", "write_shape": [6, 4], "chunk_shape": [6, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.igate.kernel.value", "write_shape": [6, 2304, 4], "chunk_shape": [6, 2304, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.outnorm.scale.value", "write_shape": [6, 768], "chunk_shape": [6, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.proj_down.kernel.value", "write_shape": [6, 768, 384], "chunk_shape": [6, 768, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.proj_up.kernel.value", "write_shape": [6, 384, 1536], "chunk_shape": [6, 384, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.q_proj.kernel.value", "write_shape": [6, 24, 32, 32], "chunk_shape": [6, 24, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.v_proj.kernel.value", "write_shape": [6, 24, 32, 32], "chunk_shape": [6, 24, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm_norm.scale.value", "write_shape": [6, 384], "chunk_shape": [6, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.post_blocks_norm.scale.value", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}]}

logs/10000/default/ocdbt.process_0/d/399bbb6a26e3af7391ef4ee2c557d20e

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

logs/10000/default/ocdbt.process_0/d/f957dd57345bd22cb0916bf0f123bcac

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

logs/8000/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"default": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1761147357293181298, "commit_timestamp_nsecs": 1761147357751493968, "custom_metadata": {}}

logs/8000/default/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('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": [384, 49152]}}, "('token_embedding', 'embedding', 'value')": {"key_metadata": [{"key": "token_embedding", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [49152, 384]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'conv1d', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "conv1d", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 768]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'conv1d', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "conv1d", "key_type": 2}, {"key": "conv", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 4, 1, 768]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "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": [6, 24, 32, 32]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'learnable_skip', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "learnable_skip", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 768]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'fgate', 'bias', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "fgate", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 4]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'fgate', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "fgate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 2304, 4]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'igate', 'bias', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "igate", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 4]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'igate', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "igate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 2304, 4]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'mlstm_cell', 'outnorm', 'scale', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "mlstm_cell", "key_type": 2}, {"key": "outnorm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 768]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'proj_down', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "proj_down", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 768, 384]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'proj_up', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "key_type": 2}, {"key": "proj_up", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 384, 1536]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "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": [6, 24, 32, 32]}}, "('xlstm_block_stack', 'blocks', 'xlstm', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm", "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": [6, 24, 32, 32]}}, "('xlstm_block_stack', 'blocks', 'xlstm_norm', 'scale', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "xlstm_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 384]}}, "('xlstm_block_stack', 'post_blocks_norm', 'scale', 'value')": {"key_metadata": [{"key": "xlstm_block_stack", "key_type": 2}, {"key": "post_blocks_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [384]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

logs/8000/default/_sharding

@@ -0,0 +1 @@
+{"bG1faGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","dG9rZW5fZW1iZWRkaW5nLmVtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwuZmdhdGUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwuZmdhdGUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwuaWdhdGUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwuaWdhdGUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLm1sc3RtX2NlbGwub3V0bm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLmNvbnYxZC5jb252LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLmNvbnYxZC5jb252Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLmxlYXJuYWJsZV9za2lwLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLnByb2pfZG93bi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLnByb2pfdXAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2suYmxvY2tzLnhsc3RtX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [null, \"tp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}","eGxzdG1fYmxvY2tfc3RhY2sucG9zdF9ibG9ja3Nfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"dp\", \"tp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}]]}}"}

logs/8000/default/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "lm_head.kernel.value", "write_shape": [384, 49152], "chunk_shape": [384, 49152], "ext_metadata": null}}, {"array_metadata": {"param_name": "token_embedding.embedding.value", "write_shape": [49152, 384], "chunk_shape": [49152, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.conv1d.conv.bias.value", "write_shape": [6, 768], "chunk_shape": [6, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.conv1d.conv.kernel.value", "write_shape": [6, 4, 1, 768], "chunk_shape": [6, 4, 1, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.k_proj.kernel.value", "write_shape": [6, 24, 32, 32], "chunk_shape": [6, 24, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.learnable_skip.value", "write_shape": [6, 768], "chunk_shape": [6, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.fgate.bias.value", "write_shape": [6, 4], "chunk_shape": [6, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.fgate.kernel.value", "write_shape": [6, 2304, 4], "chunk_shape": [6, 2304, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.igate.bias.value", "write_shape": [6, 4], "chunk_shape": [6, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.igate.kernel.value", "write_shape": [6, 2304, 4], "chunk_shape": [6, 2304, 4], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.mlstm_cell.outnorm.scale.value", "write_shape": [6, 768], "chunk_shape": [6, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.proj_down.kernel.value", "write_shape": [6, 768, 384], "chunk_shape": [6, 768, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.proj_up.kernel.value", "write_shape": [6, 384, 1536], "chunk_shape": [6, 384, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.q_proj.kernel.value", "write_shape": [6, 24, 32, 32], "chunk_shape": [6, 24, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm.v_proj.kernel.value", "write_shape": [6, 24, 32, 32], "chunk_shape": [6, 24, 32, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.blocks.xlstm_norm.scale.value", "write_shape": [6, 384], "chunk_shape": [6, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "xlstm_block_stack.post_blocks_norm.scale.value", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}]}

logs/8000/default/ocdbt.process_0/d/2f9047fa57561a8b64e87e5c02a78301

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

logs/8000/default/ocdbt.process_0/d/b07075d27dbfe48c28f624fcfc889e48

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

logs/xLSTM-tiny-stories/events.out.tfevents.1761144595.3abd4fca6fd6.6217.0.v2

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

timing_summary.json

@@ -0,0 +1,7 @@
+{
+    "total_training_duration_seconds": 3404.960225803,
+    "total_training_duration_hours": 0.9458222849452778,
+    "average_epoch_duration_seconds": 673.6284810999998,
+    "num_epochs_completed": 5,
+    "num_evaluations_completed": 5
+}

train_history.json

@@ -0,0 +1,11 @@
+{
+    "total_training_duration": 3404.960225803,
+    "avg_epoch_duration": 673.6284810999998,
+    "num_epochs_completed": 5,
+    "global_steps": 10000,
+    "global_optimizer_steps": 2000,
+    "params": {
+        "millions": 43.65,
+        "billions": 0.04
+    }
+}

trainer_config.json

@@ -0,0 +1,47 @@
+{
+    "tokenizer": "HuggingFaceTB/SmolLM2-135M",
+    "dtype": "fp32",
+    "param_dtype": "bf16",
+    "num_train_epochs": 5,
+    "per_device_train_batch_size": 16,
+    "per_device_eval_batch_size": 16,
+    "gradient_accumulation_steps": 5,
+    "seed": 42,
+    "learning_rate": 0.0003,
+    "weight_decay": 0.01,
+    "adam_beta1": 0.9,
+    "adam_beta2": 0.999,
+    "warmup_ratio": 0.2,
+    "max_grad_norm": 1.0,
+    "logging_steps": 200,
+    "output_dir": "./artifacts/",
+    "logging_dir": "./artifacts/logs/",
+    "run_name": "train",
+    "best_metric_key": "perplexity",
+    "best_n_to_keep": 3,
+    "hub_model_id": "thiomajid/xLSTM-tiny-stories",
+    "hub_private_repo": false,
+    "upload_message": "xLSTM is ready for image generation",
+    "train_dataset_url": "roneneldan/TinyStories",
+    "train_subset": null,
+    "train_split": "train",
+    "train_samples": 32000,
+    "eval_dataset_url": "roneneldan/TinyStories",
+    "eval_subset": null,
+    "eval_split": "validation",
+    "eval_samples": 3200,
+    "dataloader_drop_last": true,
+    "dataloader_num_workers": 4,
+    "worker_buffer_size": 2,
+    "text_column": "text",
+    "use_dataset_cache": true,
+    "dataset_cache_dir": "./.hf_data_cache",
+    "mesh_shape": [
+        1,
+        1
+    ],
+    "axis_names": [
+        "dp",
+        "tp"
+    ]
+}