Publish inference-only terminal weights for e2e1b-natgeo8k-bin-eff050-z0001-gtcf-soft-s1024

.gitattributes

@@ -33,3 +33,8 @@ 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
+checkpoint/1024/model_weights/d/c5c2301d2045b8a597be95f0d0a44eb7 filter=lfs diff=lfs merge=lfs -text
+checkpoint/1024/model_weights/ocdbt.process_0/d/2e48c18f3051ddbba5af052a6f903cbd filter=lfs diff=lfs merge=lfs -text
+checkpoint/1024/model_weights/ocdbt.process_0/d/6369ceab31f301404d45022c5f5c48be filter=lfs diff=lfs merge=lfs -text
+checkpoint/1024/model_weights/ocdbt.process_0/d/6d506e6662bd1f79fa15ff8ab90f6a18 filter=lfs diff=lfs merge=lfs -text
+checkpoint/1024/model_weights/ocdbt.process_0/d/eeb4ce0c6ffec3d63346a9f5308953fa filter=lfs diff=lfs merge=lfs -text

LICENSE_NOTICE.md

@@ -0,0 +1,5 @@
+# License and redistribution notice
+
+This repository contains derived research-model weights. The project owner has confirmed permission to redistribute these derived weights publicly. No standalone permissive software license is asserted for the weights by this package.
+
+Users remain responsible for complying with all applicable terms for upstream model components, tokenizer assets, datasets, and code. No tokenizer or dataset content is distributed here.

README.md

@@ -0,0 +1,26 @@
+---
+license: other
+library_name: jax
+tags:
+- jax
+- orbax
+- test-time-training
+- early-exit
+---
+
+# E2E 1B p24-r1: full_network_binary_router
+
+This is a weights-only, inference-ready Orbax release of the terminal step-1024 E2E checkpoint.
+
+- Hidden size: 2048
+- Layers: 24
+- Training sequence length: 8192
+- Router z coefficient: 0.0001
+- GTCF: true
+- Evaluation: NIAH-1: complete, NIAH-2: complete, NIAH-3: complete, LongBenchV2: complete, PG-19: complete
+
+`PG-19: complete` is an explicit status, not a numeric score.
+
+Use `checkpoint/` as the manager root and restore step `1024` through the project's read-only routed inference loader. The package includes only `model_weights`; it cannot resume optimizer training.
+
+No tokenizer bytes, optimizer state, data-iterator state, tracking metadata, internal service paths, or raw evaluation samples are included. Use a separately obtained compatible Llama-3-base tokenizer. See `LICENSE_NOTICE.md` before use or redistribution.

SHA256SUMS

@@ -0,0 +1,17 @@
+8f89543b80cafe92c77882f0afc7f15529c7f22ca73f7081e86657789d4bf6c8  LICENSE_NOTICE.md
+8721f69c4755c71c849e96a0fad7f504a49d78b092121436853112b6e969230a  README.md
+8abfe691d17bcca2c5e5b4969c1cbd2700c3ade9ba41d2876af670691c15fccf  checkpoint/1024/_CHECKPOINT_METADATA
+6c43e542a981b80f40460ee2b676927fcfcd439df4ba93a79261d76cb228126e  checkpoint/1024/model_weights/_METADATA
+6cc4522e48b1eb9e3b458677877757e395fb1705ebfaee8816ced5c04c3d8a92  checkpoint/1024/model_weights/_sharding
+8c33a10fe9da5e7807720801e3b2c1cb781976b6220425282c85a760ee0f1e93  checkpoint/1024/model_weights/array_metadatas/process_0
+f695b1e5f01a5ae7a1f62403b6f692ab3886825806ccd122e063ad1bb56aaefe  checkpoint/1024/model_weights/d/c5c2301d2045b8a597be95f0d0a44eb7
+1a0f88626bf90cf0c2053676336476c2d98c36cf371fe7b8941c060371bb7a20  checkpoint/1024/model_weights/manifest.ocdbt
+21fd8196c2798250442ee478d35b061e4c0ec2074e9f548ff026368a41df3709  checkpoint/1024/model_weights/ocdbt.process_0/d/2e48c18f3051ddbba5af052a6f903cbd
+1d8d84d4ecacebf5baa683f2ffcf9d2ad63c63f9adc55eacca547a3205d38a29  checkpoint/1024/model_weights/ocdbt.process_0/d/34b917bd8c5e654907b52df045b67ade
+824182bb46833461d66de199b509f7ea8a4834041d51f9d1d45d1b0b3e866240  checkpoint/1024/model_weights/ocdbt.process_0/d/6369ceab31f301404d45022c5f5c48be
+82d9ade6ae24fb6c761163f3f4998104499924e3f4d351cb51dcf6f4fda8c386  checkpoint/1024/model_weights/ocdbt.process_0/d/6d506e6662bd1f79fa15ff8ab90f6a18
+21f788e3cb18e50e665788f03314badc9ec89f891221d77fb2e790aca14fe225  checkpoint/1024/model_weights/ocdbt.process_0/d/7dbea4f7155bd691f0fbb54c8c66d905
+deb37fcd41df46f64a61397e30855cf226199ac27fe408391246bc36d9c64c82  checkpoint/1024/model_weights/ocdbt.process_0/d/eeb4ce0c6ffec3d63346a9f5308953fa
+0c678290b25f4330dc84819f8a274821c4c5dae97d227ec307dd2a3de3b58a8e  checkpoint/1024/model_weights/ocdbt.process_0/manifest.ocdbt
+cd2395f75177d9458a31bb867abf766f34308168bcbff8ac4c42fcef68d2c8ef  checkpoint/router_contract.json
+2def913c9b9b4017399bc1d4162002f20e2bb41f236bc3266b23c229f8ef0a65  provenance/release.json

checkpoint/1024/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers":{"model_weights":"orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"},"metrics":{},"performance_metrics":{},"init_timestamp_nsecs":1786902717399031174,"commit_timestamp_nsecs":1786902760802642758,"custom_metadata":{}}

checkpoint/1024/model_weights/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('step_index', 'init')": {"key_metadata": [{"key": "step_index", "key_type": 2}, {"key": "init", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'wte', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "wte", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128256, 256]}}, "('language_model', 'model', 'ln_f', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "ln_f", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'ln_f', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "ln_f", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'wq', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "wq", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256, 2048]}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'wk', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "wk", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256, 2048]}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'wv', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "wv", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256, 2048]}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'wo', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "wo", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 2048, 256]}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'q_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'q_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "q_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'k_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 64]}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'k_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "k_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'kv_cache_index', 'init')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "kv_cache_index", "key_type": 2}, {"key": "init", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'kv_valid_index', 'init')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "kv_valid_index", "key_type": 2}, {"key": "init", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'seq_modeling_block', 'chunk_index', 'init')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_modeling_block", "key_type": 2}, {"key": "chunk_index", "key_type": 2}, {"key": "init", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'feed_forward', 'w1', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "feed_forward", "key_type": 2}, {"key": "w1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256, 4224]}}, "('language_model', 'model', 'h', 'blocks', 'feed_forward', 'w2', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "feed_forward", "key_type": 2}, {"key": "w2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 4224, 256]}}, "('language_model', 'model', 'h', 'blocks', 'feed_forward', 'w3', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "feed_forward", "key_type": 2}, {"key": "w3", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256, 4224]}}, "('language_model', 'model', 'h', 'blocks', 'seq_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256]}}, "('language_model', 'model', 'h', 'blocks', 'seq_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'ffn_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256]}}, "('language_model', 'model', 'h', 'blocks', 'ffn_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'seq_post_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_post_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256]}}, "('language_model', 'model', 'h', 'blocks', 'seq_post_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "seq_post_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'ffn_post_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "ffn_post_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [24, 256]}}, "('language_model', 'model', 'h', 'blocks', 'ffn_post_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "ffn_post_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'ffn_prime_norm')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "ffn_prime_norm", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'ffn_prime_post_norm')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "ffn_prime_post_norm", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'blocks', 'feed_forward_prime')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "feed_forward_prime", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'prime_storage', 'ffn_prime_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "prime_storage", "key_type": 2}, {"key": "ffn_prime_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 256]}}, "('language_model', 'model', 'h', 'prime_storage', 'ffn_prime_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "prime_storage", "key_type": 2}, {"key": "ffn_prime_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'prime_storage', 'ffn_prime_post_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "prime_storage", "key_type": 2}, {"key": "ffn_prime_post_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 256]}}, "('language_model', 'model', 'h', 'prime_storage', 'ffn_prime_post_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "prime_storage", "key_type": 2}, {"key": "ffn_prime_post_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'prime_storage', 'feed_forward_prime', 'w1', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "prime_storage", "key_type": 2}, {"key": "feed_forward_prime", "key_type": 2}, {"key": "w1", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 256, 4224]}}, "('language_model', 'model', 'h', 'prime_storage', 'feed_forward_prime', 'w2', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "prime_storage", "key_type": 2}, {"key": "feed_forward_prime", "key_type": 2}, {"key": "w2", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 4224, 256]}}, "('language_model', 'model', 'h', 'prime_storage', 'feed_forward_prime', 'w3', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "prime_storage", "key_type": 2}, {"key": "feed_forward_prime", "key_type": 2}, {"key": "w3", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [6, 256, 4224]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '0', 'mean_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '0', 'mean_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '0', 'max_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '0', 'max_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '0', 'projection', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 512]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '0', 'projection', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '1', 'mean_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '1', 'mean_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '1', 'max_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '1', 'max_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '1', 'projection', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 512]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '1', 'projection', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '2', 'mean_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '2', 'mean_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '2', 'max_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '2', 'max_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '2', 'projection', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 512]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '2', 'projection', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '3', 'mean_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '3', 'mean_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '3', 'max_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '3', 'max_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '3', 'projection', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 512]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '3', 'projection', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '4', 'mean_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '4', 'mean_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "mean_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '4', 'max_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '4', 'max_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "max_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '4', 'projection', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 512]}}, "('language_model', 'model', 'h', 'router_storage', 'router', 'heads', '4', 'projection', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "router", "key_type": 2}, {"key": "heads", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "projection", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '0', 'adapter', 'down', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '0', 'adapter', 'down', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '0', 'adapter', 'up', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '0', 'adapter', 'up', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '0', 'exit_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '0', 'exit_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "0", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '1', 'adapter', 'down', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '1', 'adapter', 'down', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '1', 'adapter', 'up', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '1', 'adapter', 'up', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '1', 'exit_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '1', 'exit_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '2', 'adapter', 'down', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '2', 'adapter', 'down', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '2', 'adapter', 'up', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '2', 'adapter', 'up', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '2', 'exit_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '2', 'exit_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "2", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '3', 'adapter', 'down', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '3', 'adapter', 'down', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '3', 'adapter', 'up', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '3', 'adapter', 'up', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '3', 'exit_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '3', 'exit_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "3", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '4', 'adapter', 'down', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '4', 'adapter', 'down', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "down", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '4', 'adapter', 'up', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 128]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '4', 'adapter', 'up', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "adapter", "key_type": 2}, {"key": "up", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '4', 'exit_norm', 'weight')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "weight", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('language_model', 'model', 'h', 'router_storage', 'readouts', 'readouts', '4', 'exit_norm', 'bias')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "model", "key_type": 2}, {"key": "h", "key_type": 2}, {"key": "router_storage", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "readouts", "key_type": 2}, {"key": "4", "key_type": 1}, {"key": "exit_norm", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('language_model', 'lm_head')": {"key_metadata": [{"key": "language_model", "key_type": 2}, {"key": "lm_head", "key_type": 2}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

checkpoint/1024/model_weights/_sharding

@@ -0,0 +1 @@
+{"bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3MuZmVlZF9mb3J3YXJkLncxLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3MuZmVlZF9mb3J3YXJkLncyLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3MuZmVlZF9mb3J3YXJkLnczLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3MuZmZuX25vcm0ud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3MuZmZuX3Bvc3Rfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3Muc2VxX21vZGVsaW5nX2Jsb2NrLmtfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3Muc2VxX21vZGVsaW5nX2Jsb2NrLnFfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3Muc2VxX21vZGVsaW5nX2Jsb2NrLnd2LndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3Muc2VxX21vZGVsaW5nX2Jsb2NrLndrLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3Muc2VxX21vZGVsaW5nX2Jsb2NrLndvLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3Muc2VxX21vZGVsaW5nX2Jsb2NrLndxLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3Muc2VxX25vcm0ud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5ibG9ja3Muc2VxX3Bvc3Rfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5wcmltZV9zdG9yYWdlLmZlZWRfZm9yd2FyZF9wcmltZS53MS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5wcmltZV9zdG9yYWdlLmZlZWRfZm9yd2FyZF9wcmltZS53Mi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5wcmltZV9zdG9yYWdlLmZlZWRfZm9yd2FyZF9wcmltZS53My53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5wcmltZV9zdG9yYWdlLmZmbl9wcmltZV9ub3JtLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5wcmltZV9zdG9yYWdlLmZmbl9wcmltZV9wb3N0X25vcm0ud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy40LmFkYXB0ZXIuZG93bi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy40LmFkYXB0ZXIudXAud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy40LmV4aXRfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4wLmFkYXB0ZXIuZG93bi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4wLmFkYXB0ZXIudXAud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4wLmV4aXRfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4xLmFkYXB0ZXIuZG93bi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4xLmFkYXB0ZXIudXAud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4xLmV4aXRfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4yLmFkYXB0ZXIuZG93bi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4yLmFkYXB0ZXIudXAud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4yLmV4aXRfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4zLmFkYXB0ZXIuZG93bi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4zLmFkYXB0ZXIudXAud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yZWFkb3V0cy5yZWFkb3V0cy4zLmV4aXRfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMC5tYXhfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMC5tZWFuX25vcm0ud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMC5wcm9qZWN0aW9uLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMC5wcm9qZWN0aW9uLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMS5tYXhfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMS5tZWFuX25vcm0ud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMS5wcm9qZWN0aW9uLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMS5wcm9qZWN0aW9uLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMi5tYXhfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMi5tZWFuX25vcm0ud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMi5wcm9qZWN0aW9uLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMi5wcm9qZWN0aW9uLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMy5tYXhfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMy5tZWFuX25vcm0ud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMy5wcm9qZWN0aW9uLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuMy5wcm9qZWN0aW9uLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuNC5tYXhfbm9ybS53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuNC5tZWFuX25vcm0ud2VpZ2h0":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuNC5wcm9qZWN0aW9uLmJpYXM=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwuaC5yb3V0ZXJfc3RvcmFnZS5yb3V0ZXIuaGVhZHMuNC5wcm9qZWN0aW9uLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwubG5fZi53ZWlnaHQ=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [\"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","bGFuZ3VhZ2VfbW9kZWwubW9kZWwud3RlLndlaWdodA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"data\", \"state\"], \"partition_spec\": [null, \"state\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}"}

checkpoint/1024/model_weights/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "language_model.model.wte.weight", "write_shape": [128256, 256], "chunk_shape": [128256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.ln_f.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.seq_modeling_block.wq.weight", "write_shape": [24, 256, 2048], "chunk_shape": [24, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.seq_modeling_block.wk.weight", "write_shape": [24, 256, 2048], "chunk_shape": [24, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.seq_modeling_block.wv.weight", "write_shape": [24, 256, 2048], "chunk_shape": [24, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.seq_modeling_block.wo.weight", "write_shape": [24, 2048, 256], "chunk_shape": [24, 2048, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.seq_modeling_block.q_norm.weight", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.seq_modeling_block.k_norm.weight", "write_shape": [3, 64], "chunk_shape": [3, 64], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.feed_forward.w1.weight", "write_shape": [24, 256, 4224], "chunk_shape": [24, 256, 4224], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.feed_forward.w2.weight", "write_shape": [24, 4224, 256], "chunk_shape": [24, 4224, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.feed_forward.w3.weight", "write_shape": [24, 256, 4224], "chunk_shape": [24, 256, 4224], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.seq_norm.weight", "write_shape": [24, 256], "chunk_shape": [24, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.ffn_norm.weight", "write_shape": [24, 256], "chunk_shape": [24, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.seq_post_norm.weight", "write_shape": [24, 256], "chunk_shape": [24, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.blocks.ffn_post_norm.weight", "write_shape": [24, 256], "chunk_shape": [24, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.prime_storage.ffn_prime_norm.weight", "write_shape": [6, 256], "chunk_shape": [6, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.prime_storage.ffn_prime_post_norm.weight", "write_shape": [6, 256], "chunk_shape": [6, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.prime_storage.feed_forward_prime.w1.weight", "write_shape": [6, 256, 4224], "chunk_shape": [6, 256, 4224], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.prime_storage.feed_forward_prime.w2.weight", "write_shape": [6, 4224, 256], "chunk_shape": [6, 4224, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.prime_storage.feed_forward_prime.w3.weight", "write_shape": [6, 256, 4224], "chunk_shape": [6, 256, 4224], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.0.mean_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.0.max_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.0.projection.weight", "write_shape": [2, 512], "chunk_shape": [2, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.0.projection.bias", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.1.mean_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.1.max_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.1.projection.weight", "write_shape": [2, 512], "chunk_shape": [2, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.1.projection.bias", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.2.mean_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.2.max_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.2.projection.weight", "write_shape": [2, 512], "chunk_shape": [2, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.2.projection.bias", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.3.mean_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.3.max_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.3.projection.weight", "write_shape": [2, 512], "chunk_shape": [2, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.3.projection.bias", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.4.mean_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.4.max_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.4.projection.weight", "write_shape": [2, 512], "chunk_shape": [2, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.router.heads.4.projection.bias", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.0.adapter.down.weight", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.0.adapter.up.weight", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.0.exit_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.1.adapter.down.weight", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.1.adapter.up.weight", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.1.exit_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.2.adapter.down.weight", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.2.adapter.up.weight", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.2.exit_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.3.adapter.down.weight", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.3.adapter.up.weight", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.3.exit_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.4.adapter.down.weight", "write_shape": [128, 256], "chunk_shape": [128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.4.adapter.up.weight", "write_shape": [256, 128], "chunk_shape": [256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "language_model.model.h.router_storage.readouts.readouts.4.exit_norm.weight", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}]}

checkpoint/1024/model_weights/d/c5c2301d2045b8a597be95f0d0a44eb7

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

checkpoint/1024/model_weights/ocdbt.process_0/d/2e48c18f3051ddbba5af052a6f903cbd

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

checkpoint/1024/model_weights/ocdbt.process_0/d/6369ceab31f301404d45022c5f5c48be

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

checkpoint/1024/model_weights/ocdbt.process_0/d/6d506e6662bd1f79fa15ff8ab90f6a18

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

checkpoint/1024/model_weights/ocdbt.process_0/d/eeb4ce0c6ffec3d63346a9f5308953fa

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

checkpoint/router_contract.json

@@ -0,0 +1,138 @@
+{
+  "architecture": {
+    "attn_pdrop": 0.0,
+    "bos_token_id": 128000,
+    "compute_dtype": "bf16",
+    "embd_pdrop": 0.0,
+    "eos_token_id": 128001,
+    "feed_forward_prime": "swiglu",
+    "force_flash": false,
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 4224,
+    "mini_batch_size": 1024,
+    "num_attention_heads": 32,
+    "num_hidden_layers": 24,
+    "output_size": 128256,
+    "param_dtype": "fp32",
+    "post_norm": true,
+    "pre_norm": true,
+    "prime": true,
+    "prime_start": 18,
+    "qk_norm": true,
+    "resid_pdrop": 0.0,
+    "rms_norm_eps": 1e-06,
+    "rope_theta": 500000.0,
+    "seq_modeling_block": "SWA",
+    "sliding_window_size": 8192,
+    "state_dtype": "fp32",
+    "suffix_len": 6,
+    "swa_attention_backend": "auto",
+    "tie_word_embeddings": true,
+    "vocab_size": 128256
+  },
+  "format": "e2e-routed-router-v7",
+  "router": {
+    "adapter_rank": 128,
+    "adapter_silu": true,
+    "cadence": 4,
+    "dead_branch_skip": false,
+    "exit_ce_sites_per_chunk": 0,
+    "exit_ce_vocab_chunk_size": 4096,
+    "exit_norm": true,
+    "input_norm_eps": 0.0078125,
+    "kv_policy": "calculate",
+    "masked_exit_cost": 0.0,
+    "mode": "exit_only",
+    "sites": [
+      {
+        "end": 8,
+        "schema": "binary",
+        "start": 4
+      },
+      {
+        "end": 12,
+        "schema": "binary",
+        "start": 8
+      },
+      {
+        "end": 16,
+        "schema": "binary",
+        "start": 12
+      },
+      {
+        "end": 20,
+        "schema": "binary",
+        "start": 16
+      },
+      {
+        "end": 24,
+        "schema": "binary",
+        "start": 20
+      }
+    ],
+    "start_layer": 4,
+    "temperature": 1.0,
+    "ttt_only": false,
+    "update_bias": 1.0
+  },
+  "training": {
+    "accum_steps": 32,
+    "data_seed": 42,
+    "data_split": "train",
+    "dataset_name": "books3",
+    "dummy_dataset": false,
+    "dummy_num_tokens": 1048576,
+    "efficiency_coefficient": 0.5,
+    "exit_ce_coefficient": 0.5,
+    "gate_teacher_counterfactual": true,
+    "global_batch_size": 32,
+    "inner_learning_rate": 1.0,
+    "inner_warmup_steps": 0,
+    "mask_bos_targets": false,
+    "mask_final_target": true,
+    "model_seed": 42,
+    "optimizer_inner": {
+      "b1": null,
+      "b2": null,
+      "bf16_momentum": null,
+      "clip_gradient": 1.0,
+      "emb_wd": null,
+      "end_lr": null,
+      "init_lr": null,
+      "lr": 1.0,
+      "lr_decay_steps": null,
+      "lr_warmup_steps": null,
+      "optimizer_type": "sgd",
+      "weight_decay": null
+    },
+    "optimizer_outer": {
+      "b1": 0.9,
+      "b2": 0.95,
+      "bf16_momentum": false,
+      "clip_gradient": 1.0,
+      "emb_wd": true,
+      "end_lr": 0.0001,
+      "init_lr": 0.0,
+      "lr": 0.001,
+      "lr_decay_steps": 1024,
+      "lr_warmup_steps": 256,
+      "optimizer_type": "adamw",
+      "weight_decay": 0.1
+    },
+    "pure_min_realized": false,
+    "router_warmup_steps": 256,
+    "seq_length": 8192,
+    "spec_inner": [
+      "language_model.**.suffix_blocks.feed_forward_prime.**"
+    ],
+    "spec_outer": [
+      "language_model.**.router_storage.**"
+    ],
+    "target_cost": null,
+    "target_cost_realized": true,
+    "total_steps": 1024,
+    "train_mode": "meta",
+    "z_coefficient": 0.0001
+  }
+}

provenance/release.json

@@ -0,0 +1,105 @@
+{
+  "checkpoint_step": 1024,
+  "evaluation_status": {
+    "LongBenchV2": "complete",
+    "NIAH-1": "complete",
+    "NIAH-2": "complete",
+    "NIAH-3": "complete",
+    "PG-19": "complete"
+  },
+  "exclusions": [
+    "optimizer_state",
+    "training_data_iterator",
+    "checkpoint_metrics",
+    "raw_evaluation_outputs",
+    "evaluation_inputs_or_targets",
+    "tokenizer_bytes",
+    "tracking_metadata",
+    "internal_paths"
+  ],
+  "family": "E2E 1B p24-r1",
+  "format": "e2e-public-model-weights-release-v1",
+  "model_weights": {
+    "bytes": 5357239266,
+    "file_records": [
+      {
+        "bytes": 46772,
+        "path": "_METADATA",
+        "sha256": "6c43e542a981b80f40460ee2b676927fcfcd439df4ba93a79261d76cb228126e"
+      },
+      {
+        "bytes": 19934,
+        "path": "_sharding",
+        "sha256": "6cc4522e48b1eb9e3b458677877757e395fb1705ebfaee8816ced5c04c3d8a92"
+      },
+      {
+        "bytes": 9809,
+        "path": "array_metadatas/process_0",
+        "sha256": "8c33a10fe9da5e7807720801e3b2c1cb781976b6220425282c85a760ee0f1e93"
+      },
+      {
+        "bytes": 125777,
+        "path": "d/c5c2301d2045b8a597be95f0d0a44eb7",
+        "sha256": "f695b1e5f01a5ae7a1f62403b6f692ab3886825806ccd122e063ad1bb56aaefe"
+      },
+      {
+        "bytes": 120,
+        "path": "manifest.ocdbt",
+        "sha256": "1a0f88626bf90cf0c2053676336476c2d98c36cf371fe7b8941c060371bb7a20"
+      },
+      {
+        "bytes": 975194870,
+        "path": "ocdbt.process_0/d/2e48c18f3051ddbba5af052a6f903cbd",
+        "sha256": "21fd8196c2798250442ee478d35b061e4c0ec2074e9f548ff026368a41df3709"
+      },
+      {
+        "bytes": 741,
+        "path": "ocdbt.process_0/d/34b917bd8c5e654907b52df045b67ade",
+        "sha256": "1d8d84d4ecacebf5baa683f2ffcf9d2ad63c63f9adc55eacca547a3205d38a29"
+      },
+      {
+        "bytes": 2152925938,
+        "path": "ocdbt.process_0/d/6369ceab31f301404d45022c5f5c48be",
+        "sha256": "824182bb46833461d66de199b509f7ea8a4834041d51f9d1d45d1b0b3e866240"
+      },
+      {
+        "bytes": 96132635,
+        "path": "ocdbt.process_0/d/6d506e6662bd1f79fa15ff8ab90f6a18",
+        "sha256": "82d9ade6ae24fb6c761163f3f4998104499924e3f4d351cb51dcf6f4fda8c386"
+      },
+      {
+        "bytes": 207,
+        "path": "ocdbt.process_0/d/7dbea4f7155bd691f0fbb54c8c66d905",
+        "sha256": "21f788e3cb18e50e665788f03314badc9ec89f891221d77fb2e790aca14fe225"
+      },
+      {
+        "bytes": 2132782235,
+        "path": "ocdbt.process_0/d/eeb4ce0c6ffec3d63346a9f5308953fa",
+        "sha256": "deb37fcd41df46f64a61397e30855cf226199ac27fe408391246bc36d9c64c82"
+      },
+      {
+        "bytes": 228,
+        "path": "ocdbt.process_0/manifest.ocdbt",
+        "sha256": "0c678290b25f4330dc84819f8a274821c4c5dae97d227ec307dd2a3de3b58a8e"
+      }
+    ],
+    "files": 12,
+    "schema": "orbax-model-item-framed-byte-graph-v1",
+    "sha256": "4cd3eabe835be799b53fba8fe4382fe3437214bc4749449f9d563ec6a34cd97e"
+  },
+  "redistribution_basis": "project_owner_confirmed_permission",
+  "repo_id": "paulzy/e2e-1b-bin-eff050-z0001-gtcf-soft-step-1024",
+  "router_contract": {
+    "canonical_sha256": "5536e290353b18aa4ded1cd64f452b7f44949a708668a28ce0c0faa60de68fc4",
+    "sha256": "cd2395f75177d9458a31bb867abf766f34308168bcbff8ac4c42fcef68d2c8ef"
+  },
+  "run_label": "e2e1b-natgeo8k-bin-eff050-z0001-gtcf-soft-s1024",
+  "source_checkpoint": {
+    "schema": "orbax-full-byte-file-graph-v1",
+    "sha256": "3a931d6d97f7bc98243924ffe7691484ddb0f538e3f15e31de6a07dfce06a00b"
+  },
+  "source_code_sha256": "93bd363070c1e1bcc89b5ad54697da8a118271e9d87f9bd34d1b284e97b58de3",
+  "training_resume_supported": false,
+  "training_scope": "full_network_binary_router",
+  "weights_only": true
+}