xLSTM is ready for image generation

.gitattributes

@@ -33,3 +33,9 @@ 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/56000/default/ocdbt.process_0/d/589bbb1911ef77a68bd39d66cb5ad987 filter=lfs diff=lfs merge=lfs -text
+logs/56000/default/ocdbt.process_0/d/b7cccee99dacbfc00deb918b7696e346 filter=lfs diff=lfs merge=lfs -text
+logs/58000/default/ocdbt.process_0/d/ec9e13c2502c3b87b5d5b37860c03ad2 filter=lfs diff=lfs merge=lfs -text
+logs/58000/default/ocdbt.process_0/d/f72b9e7aa4ac1fe48d1dc7bfa695d52e filter=lfs diff=lfs merge=lfs -text
+logs/60000/default/ocdbt.process_0/d/a649a7f03f371810ad4a4d0072986f38 filter=lfs diff=lfs merge=lfs -text
+logs/60000/default/ocdbt.process_0/d/bf2745114a156076b53b8b5635b0aeac 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": 1024,
+            "conv1d_kernel_size": 4,
+            "qkv_proj_blocksize": 32,
+            "num_heads": 4,
+            "embedding_dim": 512,
+            "bias": false,
+            "dropout": 0.0,
+            "context_length": 384,
+            "_num_blocks": 12,
+            "_inner_embedding_dim": 1024
+        },
+        "_num_blocks": 12,
+        "_block_idx": null
+    },
+    "slstm_block": {
+        "slstm": {
+            "hidden_size": 512,
+            "num_heads": 4,
+            "num_states": 4,
+            "backend": "vanilla",
+            "function": "slstm",
+            "bias_init": "powerlaw_blockdependent",
+            "recurrent_weight_init": "zeros",
+            "_block_idx": null,
+            "_num_blocks": 12,
+            "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": 512,
+            "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": 12,
+        "_block_idx": null
+    },
+    "context_length": 384,
+    "num_blocks": 12,
+    "embedding_dim": 512,
+    "add_post_blocks_norm": true,
+    "bias": false,
+    "dropout": 0.0,
+    "slstm_at": [],
+    "_block_map": "0,0,0,0,0,0,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/56000/_CHECKPOINT_METADATA

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

logs/56000/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": [64, 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": [6144, 512]}}, "('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": [12, 128]}}, "('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": [12, 4, 1, 128]}}, "('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": [12, 4, 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": [12, 128]}}, "('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": [12, 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": [12, 384, 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": [12, 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": [12, 384, 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": [12, 128]}}, "('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": [12, 128, 512]}}, "('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": [12, 64, 2048]}}, "('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": [12, 4, 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": [12, 4, 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": [12, 64]}}, "('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": [64]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

logs/56000/default/_sharding

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

logs/56000/default/array_metadatas/process_0

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

logs/56000/default/ocdbt.process_0/d/589bbb1911ef77a68bd39d66cb5ad987

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

logs/56000/default/ocdbt.process_0/d/b7cccee99dacbfc00deb918b7696e346

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

logs/58000/_CHECKPOINT_METADATA

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

logs/58000/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": [64, 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": [6144, 512]}}, "('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": [12, 128]}}, "('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": [12, 4, 1, 128]}}, "('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": [12, 4, 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": [12, 128]}}, "('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": [12, 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": [12, 384, 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": [12, 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": [12, 384, 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": [12, 128]}}, "('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": [12, 128, 512]}}, "('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": [12, 64, 2048]}}, "('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": [12, 4, 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": [12, 4, 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": [12, 64]}}, "('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": [64]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

logs/58000/default/_sharding

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

logs/58000/default/array_metadatas/process_0

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

logs/58000/default/ocdbt.process_0/d/ec9e13c2502c3b87b5d5b37860c03ad2

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

logs/58000/default/ocdbt.process_0/d/f72b9e7aa4ac1fe48d1dc7bfa695d52e

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

logs/60000/_CHECKPOINT_METADATA

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

logs/60000/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": [64, 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": [6144, 512]}}, "('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": [12, 128]}}, "('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": [12, 4, 1, 128]}}, "('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": [12, 4, 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": [12, 128]}}, "('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": [12, 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": [12, 384, 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": [12, 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": [12, 384, 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": [12, 128]}}, "('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": [12, 128, 512]}}, "('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": [12, 64, 2048]}}, "('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": [12, 4, 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": [12, 4, 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": [12, 64]}}, "('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": [64]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

logs/60000/default/_sharding

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

logs/60000/default/array_metadatas/process_0

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

logs/60000/default/ocdbt.process_0/d/a649a7f03f371810ad4a4d0072986f38

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

logs/60000/default/ocdbt.process_0/d/bf2745114a156076b53b8b5635b0aeac

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

logs/xLSTM-TPU/events.out.tfevents.1761148059.5b24db4313ba.16459.0.v2

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

timing_summary.json

@@ -0,0 +1,7 @@
+{
+    "total_training_duration_seconds": 2785.887582896,
+    "total_training_duration_hours": 0.7738576619155555,
+    "average_epoch_duration_seconds": 91.71301750610003,
+    "num_epochs_completed": 30,
+    "num_evaluations_completed": 30
+}

train_history.json

@@ -0,0 +1,11 @@
+{
+    "total_training_duration": 2785.887582896,
+    "avg_epoch_duration": 91.71301750610003,
+    "num_epochs_completed": 30,
+    "global_steps": 60000,
+    "global_optimizer_steps": 12000,
+    "params": {
+        "millions": 70.77,
+        "billions": 0.07
+    }
+}

trainer_config.json

@@ -0,0 +1,47 @@
+{
+    "tokenizer": "HuggingFaceTB/SmolLM2-135M",
+    "dtype": "fp32",
+    "param_dtype": "bf16",
+    "num_train_epochs": 30,
+    "per_device_train_batch_size": 32,
+    "per_device_eval_batch_size": 32,
+    "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-TPU",
+    "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": 64000,
+    "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": [
+        8,
+        1
+    ],
+    "axis_names": [
+        "dp",
+        "tp"
+    ]
+}