diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..b7a74f04f5ebe62741a6b3de0f21aa8fa3781426 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,57 @@ 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 +simulation/rmbench/full_attention/stage2/step-100000/params/d/4e59861641970db34a401178d94abccb filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/7dc0a88b6d199ccbb016750c98012cf1 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/recurrent/stage2/step-100000/params/d/cbf2fef1100ba6c460caea5d0023e183 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa7cea5bceef1307d13d9a61ea91a35e filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/3a109d42459c2150f8b988ce2442e2b5 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fb91f39ce5a047cdcaebc0bad36caad6 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/b7327531c46c5cc244cc7981e9f4da90 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/054944d2cae72ec6dcde8844c294284c filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/88cd20aa3cd4faa14d8123fbe93a44b9 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa71f45e6411d633e6827c04a989a92f filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/recurrent/stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 filter=lfs diff=lfs merge=lfs -text +simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/pi0/stage2/step-100000/params/d/3bd5e7a25312cf2bc1713db9a1d29fa5 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/3c822efa587f8375db55af5ad6ff4923 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/5c003ca5eacb776438fd966f8fe51bdf filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/24d3ba3b910a02ad89c9e9c5ee773b7d filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/bc0fe1585dbbc9796ccd89ef1cb53d7e filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/e364328538b1fe570bf0733fb7086ce5 filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/181a1cb7822e6e4ae74e5d52fd4c3e2c filter=lfs diff=lfs merge=lfs -text +simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/6f0b57f972221d7c08e1c6e0542b0f59 filter=lfs diff=lfs merge=lfs -text +real_world/dust_table_100/full_attention/stage1/step-50000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text +real_world/dust_table_100/full_attention/stage2/step-50000/checkpoint/.metadata filter=lfs diff=lfs merge=lfs -text +real_world/dust_table_100/full_attention/stage2/step-50000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text +real_world/multitask_400/full_attention/balanced/stage1/step-200000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text +real_world/multitask_400/full_attention/linear_reward/stage1/step-200000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text +real_world/multitask_400/full_attention/balanced/stage2/step-200000/checkpoint/.metadata filter=lfs diff=lfs merge=lfs -text +real_world/multitask_400/full_attention/balanced/stage2/step-200000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text +real_world/multitask_400/full_attention/linear_reward/stage2/step-200000/checkpoint/.metadata filter=lfs diff=lfs merge=lfs -text +real_world/multitask_400/full_attention/linear_reward/stage2/step-200000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..2acb16e46cbb87b51dba7fae265fdfb77ecb7140 --- /dev/null +++ b/README.md @@ -0,0 +1,37 @@ +--- +license: other +tags: + - robotics + - vision-language-action + - memory + - imprint +--- + +# IMPRINT Models + +This repository contains the public model checkpoints for IMPRINT. + +## Layout + +| Path | Benchmark | Memory encoder | +| --- | --- | --- | +| `simulation/rmbench/full_attention/` | RMBench | Full attention | +| `simulation/rmbench/recurrent/` | RMBench | Recurrent gated delta | +| `simulation/rmbench/pi0/` | RMBench | Full attention + native Pi0 | +| `simulation/rmbench/molmoact2/` | RMBench | Full attention + MolmoAct2 | +| `simulation/robocasa/full_attention/` | Customized RoboCasa | Full attention | +| `simulation/robocasa/recurrent/` | Customized RoboCasa | Recurrent gated delta | + +The Pi0.5 directories preserve the files from their original release +repositories. The MolmoAct2 directory contains a model-only conversion of the +paper checkpoint. Optimizer, dataloader, scheduler, and trainer state are +omitted. See +`release_manifest.json` for immutable provenance, file counts, byte counts, +and checksums. + +Code, configuration, and evaluation instructions are published in the +`Amazon-FAR/imprint` GitHub repository. + +The checkpoints remain subject to their applicable upstream model licenses +and terms, including Gemma/PaliGemma where used. The repository metadata does +not relicense upstream weights. diff --git a/release_manifest.json b/release_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..eb5ed4ff29044f4c6b05aba1ab59107fdd38cd55 --- /dev/null +++ b/release_manifest.json @@ -0,0 +1,94 @@ +{ + "format_version": 1, + "release": "imprint-v1", + "migration_revisions": { + "datasets": "refs/pr/1", + "models": "refs/pr/1" + }, + "destinations": { + "datasets": { + "repo_id": "Amazon-FAR/imprint_datasets", + "repo_type": "dataset" + }, + "models": { + "repo_id": "Amazon-FAR/imprint_models", + "repo_type": "model" + } + }, + "artifacts": { + "simulation/rmbench/450": { + "source": "yilin-wu/imprint-rmbench-450", + "source_revision": "0875f60a1a0fb79751cde5c502b5e1923a157512", + "files": 350, + "bytes": 411250782303, + "destination": "Amazon-FAR/imprint_datasets" + }, + "simulation/robocasa/600": { + "source": "yilin-wu/imprint-robocasa-600", + "source_revision": "979c8e896855f393d560ef43aabc4e68bc4e6248", + "files": 465, + "bytes": 590046837296, + "destination": "Amazon-FAR/imprint_datasets" + }, + "simulation/robocasa/environment": { + "source": "yilin-wu/imprint-robocasa-environment", + "source_revision": "86b8c2f042ad03776b0d23d3cd582e62b7fedacf", + "files": 6, + "bytes": 41746512, + "destination": "Amazon-FAR/imprint_datasets" + }, + "simulation/rmbench/full_attention": { + "source": "yilin-wu/imprint-rmbench-full-attention", + "source_revision": "7be585a21fa03b9d114d771c0bc67f5e244e7b06", + "files": 29, + "bytes": 11163847579, + "destination": "Amazon-FAR/imprint_models" + }, + "simulation/rmbench/recurrent": { + "source": "yilin-wu/imprint-rmbench-recurrent", + "source_revision": "f4e20a5d54bdbb562769175625c0926369ab1d46", + "files": 31, + "bytes": 12406277703, + "destination": "Amazon-FAR/imprint_models" + }, + "simulation/robocasa/full_attention": { + "source": "yilin-wu/imprint-robocasa-full-attention", + "source_revision": "ce51086f892028fb5bf669aa3e7debe1b75998cc", + "files": 32, + "bytes": 12291932686, + "destination": "Amazon-FAR/imprint_models" + }, + "simulation/robocasa/recurrent": { + "source": "yilin-wu/imprint-robocasa-recurrent", + "source_revision": "54b23b08dc766db75827defaafe718bb18b40b73", + "files": 29, + "bytes": 12406163068, + "destination": "Amazon-FAR/imprint_models" + }, + "simulation/rmbench/molmoact2": { + "source": "W&B run mit0jkx7, model-only conversion at step 100000", + "source_revision": "mit0jkx7:100000", + "files": 4, + "bytes": 22747370042, + "destination": "Amazon-FAR/imprint_models", + "sha256": { + "checkpoint/model.pt": "ae0dacfcb2a45e82d18266ff5ed49edc26c35c41e038a43065c9a5671773360c", + "checkpoint/config.yaml": "ed4573be786a951ccb3319a090b8443db03acb1a204c827327eb6164ecfb905d" + } + }, + "simulation/rmbench/molmoact2_eval_assets": { + "source": "W&B run mit0jkx7 evaluation metadata and stored language embeddings", + "source_revision": "mit0jkx7:100000", + "files": 454, + "bytes": 3899744, + "destination": "Amazon-FAR/imprint_datasets" + }, + "simulation/rmbench/pi0": { + "source": "W&B run s746sl7k, native Pi0 model-only checkpoint", + "source_revision": "s746sl7k:100000", + "files": 24, + "bytes": 11534442762, + "destination": "Amazon-FAR/imprint_models" + } + } +} diff --git a/simulation/rmbench/full_attention/.gitattributes b/simulation/rmbench/full_attention/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..a2d93eb719e379de37214d6596a6cd84f422c49e --- /dev/null +++ b/simulation/rmbench/full_attention/.gitattributes @@ -0,0 +1,44 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz 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 +stage2/step-100000/params/d/4e59861641970db34a401178d94abccb filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261 filter=lfs diff=lfs merge=lfs -text diff --git a/simulation/rmbench/full_attention/LICENSE_GEMMA.txt b/simulation/rmbench/full_attention/LICENSE_GEMMA.txt new file mode 100644 index 0000000000000000000000000000000000000000..e0f8930d83fe9eef924c0eed50ee0e572eb471fa --- /dev/null +++ b/simulation/rmbench/full_attention/LICENSE_GEMMA.txt @@ -0,0 +1,77 @@ +Gemma Terms of Use + +Last modified: February 21, 2024 + +By using, reproducing, modifying, distributing, performing or displaying any portion or element of Gemma, Model Derivatives including via any Hosted Service, (each as defined below) (collectively, the "Gemma Services") or otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement. + +Section 1: DEFINITIONS +1.1 Definitions +(a) "Agreement" or "Gemma Terms of Use" means these terms and conditions that govern the use, reproduction, Distribution or modification of the Gemma Services and any terms and conditions incorporated by reference. + +(b) "Distribution" or "Distribute" means any transmission, publication, or other sharing of Gemma or Model Derivatives to a third party, including by providing or making Gemma or its functionality available as a hosted service via API, web access, or any other electronic or remote means ("Hosted Service"). + +(c) "Gemma" means the set of machine learning language models, trained model weights and parameters identified at ai.google.dev/gemma, regardless of the source that you obtained it from. + +(d) "Google" means Google LLC. + +(e) "Model Derivatives" means all (i) modifications to Gemma, (ii) works based on Gemma, or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Gemma, to that model in order to cause that model to perform similarly to Gemma, including distillation methods that use intermediate data representations or methods based on the generation of synthetic data Outputs by Gemma for training that model. For clarity, Outputs are not deemed Model Derivatives. + +(f) "Output" means the information content output of Gemma or a Model Derivative that results from operating or otherwise using Gemma or the Model Derivative, including via a Hosted Service. + +1.2 +As used in this Agreement, "including" means "including without limitation". + +Section 2: ELIGIBILITY AND USAGE +2.1 Eligibility +You represent and warrant that you have the legal capacity to enter into this Agreement (including being of sufficient age of consent). If you are accessing or using any of the Gemma Services for or on behalf of a legal entity, (a) you are entering into this Agreement on behalf of yourself and that legal entity, (b) you represent and warrant that you have the authority to act on behalf of and bind that entity to this Agreement and (c) references to "you" or "your" in the remainder of this Agreement refers to both you (as an individual) and that entity. + +2.2 Use +You may use, reproduce, modify, Distribute, perform or display any of the Gemma Services only in accordance with the terms of this Agreement, and must not violate (or encourage or permit anyone else to violate) any term of this Agreement. + +Section 3: DISTRIBUTION AND RESTRICTIONS +3.1 Distribution and Redistribution +You may reproduce or Distribute copies of Gemma or Model Derivatives if you meet all of the following conditions: + +You must include the use restrictions referenced in Section 3.2 as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Gemma or Model Derivatives and you must provide notice to subsequent users you Distribute to that Gemma or Model Derivatives are subject to the use restrictions in Section 3.2. +You must provide all third party recipients of Gemma or Model Derivatives a copy of this Agreement. +You must cause any modified files to carry prominent notices stating that you modified the files. +All Distributions (other than through a Hosted Service) must be accompanied by a "Notice" text file that contains the following notice: "Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms". +You may add your own intellectual property statement to your modifications and, except as set forth in this Section, may provide additional or different terms and conditions for use, reproduction, or Distribution of your modifications, or for any such Model Derivatives as a whole, provided your use, reproduction, modification, Distribution, performance, and display of Gemma otherwise complies with the terms and conditions of this Agreement. Any additional or different terms and conditions you impose must not conflict with the terms of this Agreement. + +3.2 Use Restrictions +You must not use any of the Gemma Services: + +for the restricted uses set forth in the Gemma Prohibited Use Policy at ai.google.dev/gemma/prohibited_use_policy ("Prohibited Use Policy"), which is hereby incorporated by reference into this Agreement; or +in violation of applicable laws and regulations. +To the maximum extent permitted by law, Google reserves the right to restrict (remotely or otherwise) usage of any of the Gemma Services that Google reasonably believes are in violation of this Agreement. + +3.3 Generated Output +Google claims no rights in Outputs you generate using Gemma. You and your users are solely responsible for Outputs and their subsequent uses. + +Section 4: ADDITIONAL PROVISIONS +4.1 Updates +Google may update Gemma from time to time, and you must make reasonable efforts to use the latest version of Gemma. + +4.2 Trademarks +Nothing in this Agreement grants you any rights to use Google's trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between you and Google. Google reserves any rights not expressly granted herein. + +4.3 DISCLAIMER OF WARRANTY +UNLESS REQUIRED BY APPLICABLE LAW, THE GEMMA SERVICES, AND OUTPUTS, ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY WARRANTIES OR CONDITIONS OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR OR DISTRIBUTING ANY OF THE GEMMA SERVICES OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR USE OR DISTRIBUTION OF ANY OF THE GEMMA SERVICES OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT. + +4.4 LIMITATION OF LIABILITY +TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY, CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW, SHALL GOOGLE OR ITS AFFILIATES BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL, OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO, ANY OF THE GEMMA SERVICES OR OUTPUTS EVEN IF GOOGLE OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. + +4.5 Term, Termination, and Survival +The term of this Agreement will commence upon your acceptance of this Agreement (including acceptance by your use, modification, or Distribution, reproduction, performance or display of any portion or element of the Gemma Services) and will continue in full force and effect until terminated in accordance with the terms of this Agreement. Google may terminate this Agreement if you are in breach of any term of this Agreement. Upon termination of this Agreement, you must delete and cease use and Distribution of all copies of Gemma and Model Derivatives in your possession or control. Sections 1, 2.1, 3.3, 4.2 to 4.9 shall survive the termination of this Agreement. + +4.6 Governing Law and Jurisdiction +This Agreement will be governed by the laws of the State of California without regard to choice of law principles. The UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The state and federal courts of Santa Clara County, California shall have exclusive jurisdiction of any dispute arising out of this Agreement. + +4.7 Severability +If any provision of this Agreement is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein. + +4.8 Entire Agreement +This Agreement states all the terms agreed between the parties and supersedes all other agreements between the parties as of the date of acceptance relating to its subject matter. + +4.9 No Waiver +Google will not be treated as having waived any rights by not exercising (or delaying the exercise of) any rights under this Agreement. \ No newline at end of file diff --git a/simulation/rmbench/full_attention/NOTICE b/simulation/rmbench/full_attention/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..1b9cb6541d7135c4c78a542ac4a3e04a971c96db --- /dev/null +++ b/simulation/rmbench/full_attention/NOTICE @@ -0,0 +1,11 @@ +IMPRINT +Copyright 2026 + +This distribution contains modifications to the OpenPI project: +https://github.com/Physical-Intelligence/openpi + +OpenPI is distributed under the Apache License 2.0. Gemma components and +derivatives are subject to the Gemma Terms of Use in LICENSE_GEMMA.txt. + +Gemma is provided under and subject to the Gemma Terms of Use found at +ai.google.dev/gemma/terms. diff --git a/simulation/rmbench/full_attention/README.md b/simulation/rmbench/full_attention/README.md new file mode 100644 index 0000000000000000000000000000000000000000..975ef797ebd4edaeaf3bba8dc3227d093510acdd --- /dev/null +++ b/simulation/rmbench/full_attention/README.md @@ -0,0 +1,36 @@ +--- +license: gemma +library_name: jax +tags: + - robotics + - vision-language-action + - memory + - rmbench + - dinov3 +--- + +# IMPRINT RMBench Full-Attention Checkpoint + +Canonical Pi0.5 IMPRINT checkpoint for RMBench using the full-attention +memory encoder. + +- Stage 1: 160,000 updates, W&B `ylinwu/Robocasa/j8t4ebpo` +- Stage 2: 100,000 updates, W&B `liqhuiha/Robocasa/grruxd5p` +- memory tokens: 256 x 1024 +- policy horizon: 50 +- visual memory: three-view pooled DINOv3 features +- Stage 2: memory encoder finetuned; image tower and language embedding table frozen + +The release contains the Stage-1 encoder, the complete Stage-2 Orbax parameter +checkpoint (including the PaliGemma/Gemma backbone), normalization statistics, +portable history configuration, and provenance. The checkpoint is a Gemma +model derivative and is distributed subject to the included Gemma Terms of +Use. + +Use the matching `yilin-wu/imprint-rmbench-450` dataset and the evaluation +contract in the IMPRINT GitHub repository. Evaluation requires strict causal +actions, online recent DINOv3 features, `EVAL_CONSOLIDATE_ALL=1`, BGR policy +inputs, and 50-step prediction/execution chunks. + +The matching training store and reproducibility code are publicly available +under the IMPRINT release. diff --git a/simulation/rmbench/full_attention/SHA256SUMS b/simulation/rmbench/full_attention/SHA256SUMS new file mode 100644 index 0000000000000000000000000000000000000000..88df321c5695e2ff71c7967d85a65ffe1f828676 --- /dev/null +++ b/simulation/rmbench/full_attention/SHA256SUMS @@ -0,0 +1,27 @@ +3e2c24001f9ef57bf7ec959a3658fbb49cdad113cdf394c264da9d16f9bdd132 ./LICENSE_GEMMA.txt +5b21d7116ac42a0ab6141488cb654c2411557ac977e2276accb88ef774f28c19 ./NOTICE +32a0a8a3422765c68c10bbd944cbb0d96d4781fab507b061f02f2a4099085a05 ./README.md +a4c8e95585b546442350a3e3e7b88939569a69a092409fc32ef1203e9589d73b ./config/history_config.yaml +6569bbc55da67ba7f89a5250ae7261f192754bb27a7fc5db8b579b013452d61a ./config/input_provenance.json +372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994 ./config/norm_stats.json +5ea9c49f1b975f0aa1c4d7275381be406954d7d5e9bf6c359c9ee6dbe32359a9 ./stage1/stage1_step160000.npz +3280d8311fd54fd78143091908db4051671ef7dd6de4665e0298438e428b5cac ./stage2/step-100000/_CHECKPOINT_METADATA +372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994 ./stage2/step-100000/assets/robomme/norm_stats.json +cc86ddddfcb438e841a5533b5f22f33c86d4e29c6268084313bcee24f6fb918c ./stage2/step-100000/params/_METADATA +f126ea3bbd65b911ba0a0e3d27dc3c92c6c4b072fe97807cefe8b947a0e8e66e ./stage2/step-100000/params/_sharding +f666d8e56c6199af6d45f6cc1c46765cb71cf3583647a13d67c8d56ee4b22356 ./stage2/step-100000/params/array_metadatas/process_0 +21bd96f6985b94e4a1521163ca6bf5e2008f322a58c9fd0e949ba6008b08d8a2 ./stage2/step-100000/params/d/4e59861641970db34a401178d94abccb +69bddfd565fec7dcc3f500bca4fbe862a0615908f25363b7053e02b74319b962 ./stage2/step-100000/params/manifest.ocdbt +91b19940b50c87336da9008ed29e69dfaed8cf1a83d3caff93d3664e5ac9c5be ./stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1 +3b5861efcd7b43d842b14554dbafb062a087edbe0f7f8a61b275bfa6d23a0e9a ./stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae +5f18bdcbb4fb31730dbf44b59cec875e876788e1fe3e814e9adfb69663df3da0 ./stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261 +43046316feb9752714eb4ea27d8cb758279120aeb5ed646374206af3983f4975 ./stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9 +60681c50b5dcc37dfa61b64e6a0eb5de17b290463e3bebf82ddb39cde4323a7f ./stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164 +7b03262932b220351d793247960e57eadd53cd4b1e1be9cf7f54f5de35600ea6 ./stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa +034c510b6b4a7b7a01bff53cd076593b5f8e993f6497015f311517b1feced6d7 ./stage2/step-100000/params/ocdbt.process_0/d/48e92482d01b438d4442cd8ccbcb5df5 +72b78b1724e3ec3b4eb9f23e9ce9986349691e72e3ad863b839533f3a96f922e ./stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798 +29c1fdaf3193fee770906cd10bd8b155c1a124a6b3008f364e4e56d41367c77d ./stage2/step-100000/params/ocdbt.process_0/d/70e9c830e0a291cdb934ad90c5ca3dc1 +455068761755068aae533bc9f2f7c8f556736ad27128f90a4bfea3f63c9f003c ./stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055 +7b85974983b9719ecf40991f70066431672edc8a4b870e766ea41b67ac4c8993 ./stage2/step-100000/params/ocdbt.process_0/d/9940789f52fa3ebe1b791f1d439deb33 +bcf80d7c5753d9a4e7c040f4695512be62299c2653395f56ced3797e1c1dc4f1 ./stage2/step-100000/params/ocdbt.process_0/d/a9d841920262430040bbe87d60eeeb29 +2b2be2c66199fcd1776de3aea5b3e2c601388e74766bbaf5b767b6ed02384c72 ./stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt diff --git a/simulation/rmbench/full_attention/config/history_config.yaml b/simulation/rmbench/full_attention/config/history_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..511dfcde0816cf6df7fc8cfbb8be6d22a379b021 --- /dev/null +++ b/simulation/rmbench/full_attention/config/history_config.yaml @@ -0,0 +1,68 @@ +budget: 256 +num_views: 3 +streaming_obs_horizon: 16 +integration_type: modulation +memory_token_dim: 1024 +representation_type: cosmo_bottleneck +memory_feature: + img: + net: identity + input_dim: 768 + pos: + input_dim: 768 + hidden_dim: 768 + state: + input_dim: 16 + hidden_dim: 512 +cosmo_bottleneck_memory: + episode_cache_size: 450 + visual_feature_stream: dinov3 + feature_store_path: data/rmbench-450/stage1 + latent_source: offline + n_variants: 8 + num_query_samples: 16 + variant: Wan2.2-TI2V-5B + temporal_factor: 4 + downsample: 2 + num_latents: 256 + num_cross_blocks: 2 + num_heads: 8 + cosmos_latent_dim: 48 + siglip_dim: 768 + proprio_in_dim: 16 + proprio_encoding: raw + action_in_dim: 16 + action_encoding: raw + lang_emb_dim: 2048 + action_tokenizer_hidden: 512 + proprio_tokenizer_hidden: 512 + pool_h: 2 + pool_w: 2 + wan_h: 256 + wan_w: 256 + latent_cap: 4096 + max_recent_tokens: 768 + max_step_tokens: 1244 +cosmo_decoder: + num_heads: 8 + num_blocks: 2 + align_dim: 768 + align_num_patches: 192 + ffn_inner: 8192 + action_dim: 16 + proprio_dim: 16 + max_query_steps: 2048 +frame_jitter: + enabled: false + start_jitter: 3 + stride_choices: + - 1 + - 2 + - 3 + interval_jitter: 1 + canonical_downsample: 2 +stream_aug: + enabled: false + gauss_sigma: 0.05 + stream_drop_p: 0.1 + feat_drop_p: 0.0 diff --git a/simulation/rmbench/full_attention/config/input_provenance.json b/simulation/rmbench/full_attention/config/input_provenance.json new file mode 100644 index 0000000000000000000000000000000000000000..960801fdf7306acf86186947faed7eb0f308699a --- /dev/null +++ b/simulation/rmbench/full_attention/config/input_provenance.json @@ -0,0 +1,25 @@ +{ + "action_causality": "strictly_causal", + "architecture": "perceiver_v1", + "base_policy": "pi0.5", + "checkpoint_metadata_sha256": "3280d8311fd54fd78143091908db4051671ef7dd6de4665e0298438e428b5cac", + "checkpoint_path": "stage2/step-100000", + "checkpoint_step": 100000, + "dataset": "yilin-wu/imprint-rmbench-450", + "eval_consolidate_all": 1, + "execution_horizon": 50, + "freeze_variant": "image_and_language_embedding", + "latent_cap": 4096, + "max_recent_tokens": 768, + "max_step_tokens": 1244, + "memory_language_source": "live", + "num_latents": 256, + "prediction_horizon": 50, + "stage1_checkpoint": "stage1/stage1_step160000.npz", + "stage1_sha256": "5ea9c49f1b975f0aa1c4d7275381be406954d7d5e9bf6c359c9ee6dbe32359a9", + "stage1_wandb_run": "ylinwu/Robocasa/j8t4ebpo", + "stage2_auxiliary_loss_weight": 0.0, + "stage2_memory_encoder": "finetuned", + "stage2_wandb_run": "liqhuiha/Robocasa/grruxd5p", + "visual_feature_stream": "dinov3" +} diff --git a/simulation/rmbench/full_attention/config/norm_stats.json b/simulation/rmbench/full_attention/config/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..631b84bab53fc384b33fa5468192db069a535b6b --- /dev/null +++ b/simulation/rmbench/full_attention/config/norm_stats.json @@ -0,0 +1,152 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + -0.2884466350078583, + 1.116854190826416, + 0.8176121115684509, + -0.6286301016807556, + 0.0027539930306375027, + 0.3858006000518799, + 0.0, + 0.29012763500213623, + 1.1434277296066284, + 0.834911048412323, + -0.6164064407348633, + 0.0033871293999254704, + 0.1835649162530899, + 0.0, + 0.44485998153686523, + 0.693847119808197 + ], + "std": [ + 0.2924744486808777, + 0.9123055338859558, + 0.6921584606170654, + 0.5167360901832581, + 0.024822130799293518, + 0.7730752229690552, + 0.0, + 0.3022353947162628, + 0.9363216161727905, + 0.7197534441947937, + 0.5159713625907898, + 0.042254358530044556, + 0.4932859241962433, + 0.0, + 0.4852803945541382, + 0.43509986996650696 + ], + "q01": [ + -1.1566282510757446, + -0.00001813583730836399, + -0.000035173223295714706, + -1.493809461593628, + -0.044773634523153305, + -2.2744300365448, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.4539425373077393, + -0.11438962072134018, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410673141479492, + 2.2643818855285645, + 0.000017526803276268765, + 0.09094329178333282, + 1.4207451343536377, + 0.0, + 0.9992868304252625, + 2.5619561672210693, + 2.3963286876678467, + 0.000027463351216283627, + 0.12538659572601318, + 1.50703763961792, + 0.0, + 1.0, + 1.0 + ] + }, + "state": { + "mean": [ + -0.28619107604026794, + 1.1170974969863892, + 0.8187518119812012, + -0.6292704939842224, + 0.002742639509961009, + 0.38321176171302795, + 0.0, + 0.287190318107605, + 1.1249585151672363, + 0.8228866457939148, + -0.6072733998298645, + 0.003468270180746913, + 0.18374495208263397, + 0.0, + 0.47009503841400146, + 0.700372040271759 + ], + "std": [ + 0.28622451424598694, + 0.9085981845855713, + 0.6900886297225952, + 0.5145737528800964, + 0.024824900552630424, + 0.7737277746200562, + 0.0, + 0.3019956946372986, + 0.9365604519844055, + 0.720547616481781, + 0.5160077214241028, + 0.04248249530792236, + 0.4887017011642456, + 0.0, + 0.4873206615447998, + 0.4321308135986328 + ], + "q01": [ + -0.991405725479126, + -0.00001813583730836399, + -0.00003476843994576484, + -1.4938021898269653, + -0.04476987197995186, + -2.274427652359009, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.455243706703186, + -0.11395290493965149, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410663366317749, + 2.2642898559570312, + 0.000016971260265563615, + 0.09090737998485565, + 1.4203816652297974, + 0.0, + 1.0001307725906372, + 2.56237530708313, + 2.397764205932617, + 0.000027463351216283627, + 0.12876388430595398, + 1.4705802202224731, + 0.0, + 1.0, + 1.0 + ] + } + } +} \ No newline at end of file diff --git a/simulation/rmbench/full_attention/stage1/stage1_step160000.npz b/simulation/rmbench/full_attention/stage1/stage1_step160000.npz new file mode 100644 index 0000000000000000000000000000000000000000..7b517351ef49bf422d24afa04b2292894ede296b --- /dev/null +++ b/simulation/rmbench/full_attention/stage1/stage1_step160000.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ea9c49f1b975f0aa1c4d7275381be406954d7d5e9bf6c359c9ee6dbe32359a9 +size 331937012 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/_CHECKPOINT_METADATA b/simulation/rmbench/full_attention/stage2/step-100000/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..cdb9eaf23ce998efdaffcdee260ce324a1997224 --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1788097695402693623, "commit_timestamp_nsecs": 1788097725293618197, "custom_metadata": {}} \ No newline at end of file diff --git a/simulation/rmbench/full_attention/stage2/step-100000/assets/robomme/norm_stats.json b/simulation/rmbench/full_attention/stage2/step-100000/assets/robomme/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..631b84bab53fc384b33fa5468192db069a535b6b --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/assets/robomme/norm_stats.json @@ -0,0 +1,152 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + -0.2884466350078583, + 1.116854190826416, + 0.8176121115684509, + -0.6286301016807556, + 0.0027539930306375027, + 0.3858006000518799, + 0.0, + 0.29012763500213623, + 1.1434277296066284, + 0.834911048412323, + -0.6164064407348633, + 0.0033871293999254704, + 0.1835649162530899, + 0.0, + 0.44485998153686523, + 0.693847119808197 + ], + "std": [ + 0.2924744486808777, + 0.9123055338859558, + 0.6921584606170654, + 0.5167360901832581, + 0.024822130799293518, + 0.7730752229690552, + 0.0, + 0.3022353947162628, + 0.9363216161727905, + 0.7197534441947937, + 0.5159713625907898, + 0.042254358530044556, + 0.4932859241962433, + 0.0, + 0.4852803945541382, + 0.43509986996650696 + ], + "q01": [ + -1.1566282510757446, + -0.00001813583730836399, + -0.000035173223295714706, + -1.493809461593628, + -0.044773634523153305, + -2.2744300365448, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.4539425373077393, + -0.11438962072134018, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410673141479492, + 2.2643818855285645, + 0.000017526803276268765, + 0.09094329178333282, + 1.4207451343536377, + 0.0, + 0.9992868304252625, + 2.5619561672210693, + 2.3963286876678467, + 0.000027463351216283627, + 0.12538659572601318, + 1.50703763961792, + 0.0, + 1.0, + 1.0 + ] + }, + "state": { + "mean": [ + -0.28619107604026794, + 1.1170974969863892, + 0.8187518119812012, + -0.6292704939842224, + 0.002742639509961009, + 0.38321176171302795, + 0.0, + 0.287190318107605, + 1.1249585151672363, + 0.8228866457939148, + -0.6072733998298645, + 0.003468270180746913, + 0.18374495208263397, + 0.0, + 0.47009503841400146, + 0.700372040271759 + ], + "std": [ + 0.28622451424598694, + 0.9085981845855713, + 0.6900886297225952, + 0.5145737528800964, + 0.024824900552630424, + 0.7737277746200562, + 0.0, + 0.3019956946372986, + 0.9365604519844055, + 0.720547616481781, + 0.5160077214241028, + 0.04248249530792236, + 0.4887017011642456, + 0.0, + 0.4873206615447998, + 0.4321308135986328 + ], + "q01": [ + -0.991405725479126, + -0.00001813583730836399, + -0.00003476843994576484, + -1.4938021898269653, + -0.04476987197995186, + -2.274427652359009, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.455243706703186, + -0.11395290493965149, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410663366317749, + 2.2642898559570312, + 0.000016971260265563615, + 0.09090737998485565, + 1.4203816652297974, + 0.0, + 1.0001307725906372, + 2.56237530708313, + 2.397764205932617, + 0.000027463351216283627, + 0.12876388430595398, + 1.4705802202224731, + 0.0, + 1.0, + 1.0 + ] + } + } +} \ No newline at end of file diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/_METADATA b/simulation/rmbench/full_attention/stage2/step-100000/params/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..35e00c4cdd795b15ba973b4b0a873d98dc2ab5d0 --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_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": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1152, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 16, 72, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [14, 14, 3, 288]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "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": [1152, 512]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64288, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_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": [512]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'kv_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "kv_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'mem_rms_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "mem_rms_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": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'out_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "out_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'q_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "q_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 2048, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4096, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_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": [256, 32]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'lang_film', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'lang_film', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('params', 'mem_encoder', 'latents', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "latents", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1024]}}, "('params', 'mem_encoder', 'proj_cosmo', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_cosmo', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "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, 1024]}}, "('params', 'mem_encoder', 'proj_recent', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_recent', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 1024]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/_sharding b/simulation/rmbench/full_attention/stage2/step-100000/params/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..4b5d211606dee38d346158df308661fd70a6e695 --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/_sharding @@ -0,0 +1 @@ +{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm1lbV9ybXNfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm91dF9laW5zdW1fbWVtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLmt2X2VpbnN1bV9tZW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLnFfZWluc3VtX21lbS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhdGVudHMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}"} \ No newline at end of file diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/array_metadatas/process_0 b/simulation/rmbench/full_attention/stage2/step-100000/params/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..b851940b8671eadcb003b52cf463492509c92d25 --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 1152, 1076], "chunk_shape": [27, 1152, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 16, 72, 288], "chunk_shape": [27, 16, 72, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [1152, 512], "chunk_shape": [1152, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [1024, 768], "chunk_shape": [1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 8, 256, 512], "chunk_shape": [18, 8, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 8, 512, 256], "chunk_shape": [18, 8, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.kv_einsum_mem.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.mem_rms_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.out_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.q_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.bias.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.kernel.value", "write_shape": [18, 1024, 512], "chunk_shape": [18, 1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 2048, 4096], "chunk_shape": [18, 2, 2048, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 1024, 1024], "chunk_shape": [18, 2, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.latents.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.kernel.value", "write_shape": [12, 1024], "chunk_shape": [12, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.kernel.value", "write_shape": [192, 1024], "chunk_shape": [192, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}]} \ No newline at end of file diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/d/4e59861641970db34a401178d94abccb b/simulation/rmbench/full_attention/stage2/step-100000/params/d/4e59861641970db34a401178d94abccb new file mode 100644 index 0000000000000000000000000000000000000000..285fdaa752b8ccae348aaf7f2e4445f8a42d8789 --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/d/4e59861641970db34a401178d94abccb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:21bd96f6985b94e4a1521163ca6bf5e2008f322a58c9fd0e949ba6008b08d8a2 +size 150486 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/manifest.ocdbt b/simulation/rmbench/full_attention/stage2/step-100000/params/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..f691003c29d67f27727ce497c7c3aacfa2935be5 Binary files /dev/null and b/simulation/rmbench/full_attention/stage2/step-100000/params/manifest.ocdbt differ diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1 new file mode 100644 index 0000000000000000000000000000000000000000..fb557acfbf518d0226ce4b4a9faed83e3da8769b --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:91b19940b50c87336da9008ed29e69dfaed8cf1a83d3caff93d3664e5ac9c5be +size 2240468512 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae new file mode 100644 index 0000000000000000000000000000000000000000..b7b29deeac80cc8c33ebaf0ac4c20ee10975c0ec --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3b5861efcd7b43d842b14554dbafb062a087edbe0f7f8a61b275bfa6d23a0e9a +size 2215372366 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261 new file mode 100644 index 0000000000000000000000000000000000000000..68c352ffeb3dfc3d6a980fbac9943564d3166b76 --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5f18bdcbb4fb31730dbf44b59cec875e876788e1fe3e814e9adfb69663df3da0 +size 1248712019 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9 new file mode 100644 index 0000000000000000000000000000000000000000..0a94bf44bb6b33a438a8969179ddb9ebb4a47359 --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:43046316feb9752714eb4ea27d8cb758279120aeb5ed646374206af3983f4975 +size 1120252540 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164 new file mode 100644 index 0000000000000000000000000000000000000000..7189084d126bba682a96e4b32b4f5d9a2a190ccd --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:60681c50b5dcc37dfa61b64e6a0eb5de17b290463e3bebf82ddb39cde4323a7f +size 2240199200 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa new file mode 100644 index 0000000000000000000000000000000000000000..b988fb2589a62896ec88eb56154dcd4b5d1955a2 --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7b03262932b220351d793247960e57eadd53cd4b1e1be9cf7f54f5de35600ea6 +size 646214199 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/48e92482d01b438d4442cd8ccbcb5df5 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/48e92482d01b438d4442cd8ccbcb5df5 new file mode 100644 index 0000000000000000000000000000000000000000..2a0b58f2c74a2f8b9539cde943dce7b6fad34243 Binary files /dev/null and b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/48e92482d01b438d4442cd8ccbcb5df5 differ diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798 new file mode 100644 index 0000000000000000000000000000000000000000..19779930ef2a577945b3c222e1b34f87591fbfc9 --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:72b78b1724e3ec3b4eb9f23e9ce9986349691e72e3ad863b839533f3a96f922e +size 1120250955 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/70e9c830e0a291cdb934ad90c5ca3dc1 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/70e9c830e0a291cdb934ad90c5ca3dc1 new file mode 100644 index 0000000000000000000000000000000000000000..117e340c8db0a9debe76a627cee3562807fdd1b9 Binary files /dev/null and b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/70e9c830e0a291cdb934ad90c5ca3dc1 differ diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055 new file mode 100644 index 0000000000000000000000000000000000000000..dd6eb9cda15d2cd779663bcad4a924195348ddea --- /dev/null +++ b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:455068761755068aae533bc9f2f7c8f556736ad27128f90a4bfea3f63c9f003c +size 150102 diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/9940789f52fa3ebe1b791f1d439deb33 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/9940789f52fa3ebe1b791f1d439deb33 new file mode 100644 index 0000000000000000000000000000000000000000..fd81898ba1c0f7c4ca1720a23fc59c648a36e352 Binary files /dev/null and b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/9940789f52fa3ebe1b791f1d439deb33 differ diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/a9d841920262430040bbe87d60eeeb29 b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/a9d841920262430040bbe87d60eeeb29 new file mode 100644 index 0000000000000000000000000000000000000000..1e6fa1d3d16759afb968a0b029b46e52a1584330 Binary files /dev/null and b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/a9d841920262430040bbe87d60eeeb29 differ diff --git a/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..0b754ca558715583f05764086432604e8e58792f Binary files /dev/null and b/simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt differ diff --git a/simulation/rmbench/molmoact2/README.md b/simulation/rmbench/molmoact2/README.md new file mode 100644 index 0000000000000000000000000000000000000000..56a95fc9a6b17a18fd3438032f49073c9b364b27 --- /dev/null +++ b/simulation/rmbench/molmoact2/README.md @@ -0,0 +1,22 @@ +--- +library_name: pytorch +tags: + - robotics + - rmbench + - molmoact2 + - vision-language-action +--- + +# IMPRINT RMBench MolmoAct2 + +This directory contains the model-only 100k Stage-2 checkpoint used for the +IMPRINT MolmoAct2 RMBench experiment (`mit0jkx7`). Optimizer and trainer state +are excluded. + +The checkpoint uses the full-attention IMPRINT Stage-1 memory encoder with 256 +memory tokens. Evaluation uses online WAN and DINOv3 history features and the +stored Pi0.5 language embeddings under +`Amazon-FAR/imprint_datasets/simulation/rmbench/molmoact2_eval_assets/`. + +Use the `molmoact2` branch of the IMPRINT repository for training and +closed-loop evaluation. diff --git a/simulation/rmbench/molmoact2/checkpoint/config.yaml b/simulation/rmbench/molmoact2/checkpoint/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0c01d6182152ea1f34a43e295f9040ee3e6696c3 --- /dev/null +++ b/simulation/rmbench/molmoact2/checkpoint/config.yaml @@ -0,0 +1,3172 @@ +run_name: molmoact2-stage2-dinov3-memory-reference-fixed-film-100k-v3 +model: + model_name: molmoact2 + data_formatter: + prompt_templates: uber_model_v2 + message_format: qwen3 + system_prompt: demo_or_style_v2 + always_start_with_space: false + default_inference_len: 65 + select_answer: best + debug: false + image_last: false + format_message_list: null + p_one_message: 0.0 + eval_system_prompt_mapping: null + p_choice_content_in_mc: 1.0 + template_video_mc_questions: true + pointing_format: html-v2 + points_decimal_places: 1 + use_seperate_non_pointing_qa_style: false + timestamp_mode: 50-percent-seconds + output_timestamp_mode: seconds + seconds_decimal_places: 1 + p_multi_point_all_image: 0.5 + add_setup_tokens: true + add_control_tokens: true + use_seperate_count_without_pointing_style: false + sample_random_initial_point: true + llm: + d_model: 2560 + n_heads: 32 + n_kv_heads: 8 + head_dim: 128 + qkv_bias: false + clip_qkv: null + n_layers: 36 + mlp_ratio: 4 + mlp_hidden_size: 19456 + activation_type: swiglu + block_type: sequential + rope: true + rope_full_precision: true + rope_theta: 5000000.0 + rope_type: default + rope_factor: null + rope_high_freq_factor: null + rope_low_freq_factor: null + rope_original_max_position_embeddings: null + rope_attention_factor: null + rope_beta_fast: null + rope_beta_slow: null + rope_mscale: null + rope_mscale_all_dim: null + rope_truncate: null + attention_type: sdpa + full_attention_layers: null + sliding_attention_rope_scaling: false + float32_attention: true + attention_dropout: 0.0 + attention_layer_norm: true + attention_layer_norm_type: qwen3 + residual_dropout: 0.1 + response_residual_dropout: 0.0 + layer_norm_type: rms + layer_norm_with_affine: true + layer_norm_eps: 1.0e-06 + attention_layer_norm_with_affine: true + max_sequence_length: 16384 + max_position_embeddings: null + include_bias: false + bias_for_layer_norm: null + norm_after: false + moe_num_experts: 8 + moe_top_k: 2 + moe_mlp_impl: sparse + moe_log_expert_assignment: false + moe_shared_expert: false + moe_lbl_in_fp32: false + moe_interleave: false + moe_loss_weight: 0.1 + moe_zloss_weight: null + moe_dropless: true + moe_capacity_factor: 1.25 + embedding_dropout: 0.0 + scale_logits: false + vocab_size: 151669 + additional_vocab_size: 128 + weight_tying: false + embedding_size: 155648 + use_position_ids: true + tokenizer: + identifier: Qwen/Qwen3-4B + tokenizer_dir: null + new_tokens_for_both_input_and_output: + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + - + init_path: null + init_incremental: null + new_embedding_init_range: 0.02 + initializer_range: 0.02 + normalize_input_embeds: false + activation_checkpoint: whole_layer + compile: blocks + fix_pad_tokenizer: true + init_std: 0.02 + init_fn: normal + init_cutoff_factor: null + vision_backbone: + vit: + image_model_type: siglip + image_default_input_size: + - 378 + - 378 + image_patch_size: 14 + image_pos_patch_size: 14 + image_emb_dim: 1152 + image_num_heads: 16 + image_num_key_value_heads: 16 + image_num_layers: 27 + image_head_dim: 72 + image_mlp_dim: 4304 + image_mlp_activations: gelu_pytorch_tanh + image_dropout_rate: 0.0 + image_num_pos: 729 + image_norm_eps: 1.0e-06 + attention_dropout: 0.0 + residual_dropout: 0.0 + initializer_range: 0.02 + float32_attention: true + attention_type: sdpa + sdpa_backend: all + activation_checkpointing: true + init_path: null + resize_mode: siglip + pad_value: 0.0 + normalize: siglip + image_pooling_2d: attention_meanq + pooling_attention_mask: true + image_projector: mlp + image_padding_embed: null + vit_layers: + - -3 + - -9 + skip_unused_layers: true + use_deepstack: false + share_connector: false + image_feature_dropout: 0.0 + connector_activation_checkpointing: true + compile_vit: blocks + pool_size_embeds: null + compile_connector: dynamic + normalize_on_gpu: true + use_image_augmentation: full + mm_preprocessor: + max_answer_len: null + last_message_loss_only: false + max_text_tokens: null + loss_token_weighting: root_subsegments_root_tokens + max_frames: 8 + frame_sample_mode: uniform_last_frame + candidate_sampling_fps: + - 0.25 + - 0.5 + - 1.0 + - 2.0 + - 4.0 + - 6.0 + - 8.0 + - 16.0 + cache_videos: true + loading_method: av + max_fps: + - 2.0 + time_sampling: true + time_mode: per-frame-compact + subtitle_mode: frame_1 + max_crops: 1 + overlap_margins: + - 4.0 + - 4.0 + use_col_tokens: false + periodic_high_res_frame: null + high_low_train_mode: local_rnd + high_res_frame_sample_options: null + periodic_sample_rate_training: + 4: + - 0.9 + - 0.03 + - 0.03 + - 0.04 + 3: + - 0.6 + - 0.2 + - 0.2 + skip_low_res_in_high_low: false + pooling_w: 3 + pooling_h: 3 + high_res_pooling_w: null + high_res_pooling_h: null + query_based_resolution_selection: false + max_queries_for_resolution_selection: 8 + use_frame_special_tokens: true + frame_sel_clip_identifier: google/siglip2-so400m-patch14-384 + image_padding_mask: false + max_subtitle_tokens: null + image: + crop_mode: resize + use_col_tokens: true + max_crops: 8 + high_res_max_crops: 24 + p_high_res: 0.0 + pooling_w: 2 + pooling_h: 2 + overlap_margins: + - 4 + - 4 + max_images: 5 + max_multi_image_crops: 8 + multi_image_pooling_w: 2 + multi_image_pooling_h: 2 + use_single_crop_col_tokens: false + use_single_crop_start_token: true + topk: null + prune_from_frame: 0 + bi_directional_attn: null + shared_low_high_embedding: true + debug: null + cp_enabled: false + apply_cp_to_vision_backbone: false + lora_enable: false + lora_rank: 64 + lora_alpha: 16 + lora_dropout: 0.05 + lora_bias: none + max_action_dim: 16 + action_horizon: 50 + n_action_steps: null + n_obs_steps: 1 + action_expert: + max_horizon: 50 + max_action_dim: 16 + hidden_size: 768 + num_layers: 36 + num_heads: 8 + mlp_ratio: 4.0 + ffn_multiple_of: 256 + timestep_embed_dim: 256 + dropout: 0.0 + attn_dropout: 0.0 + context_layer_norm: true + qk_norm: true + qk_norm_eps: 1.0e-06 + rope: true + rope_on_cross_attention: false + causal_attn: false + memory_enabled: true + compile: null + memory: + enabled: true + mode: finetune + num_latents: 256 + memory_dim: 1024 + cosmos_latent_dim: 48 + recent_dim: 768 + action_dim: 16 + proprio_dim: 16 + language_dim: 2048 + tokenizer_hidden_dim: 512 + num_cross_blocks: 2 + num_heads: 8 + stage1_checkpoint: artifacts/checkpoints/rmbench-full-attention/stage1/stage1_step160000.npz + add_action_expert: true + action_expert_detach_vlm: false + action_expert_depth_gate: false + action_expert_depth_gate_per_layer: false + action_expert_depth_gate_init_bias: -4.0 + action_format: continuous + state_format: discrete + flow_matching_num_steps: 10 + flow_matching_cutoff: 1.0 + flow_matching_time_offset: 0.001 + flow_matching_time_scale: 0.999 + flow_matching_beta_alpha: 1.0 + flow_matching_beta_beta: 1.5 + num_flow_timesteps: 1 + mask_action_chunk_padding: true + mask_action_dim_padding: true + enable_depth_reasoning: false + num_depth_codes: 100 + depth_code_input_noise_rate: 0.0 + robot_processor: + metadata_by_tag: + rmbench_stage2: + action_key: action + state_keys: + - observation.state + camera_keys: + - observation.images.image + - observation.images.wrist_image + - observation.images.right_image + normalize_gripper: true + action_dim: 16 + action_horizon: 50 + n_action_steps: 50 + setup_type: three-camera dual-arm robot in RMBench + control_mode: absolute joint pose + action_stats: + mean: + - -0.2884466350078583 + - 1.116854190826416 + - 0.8176121115684509 + - -0.6286301016807556 + - 0.0027539930306375027 + - 0.3858006000518799 + - 0.0 + - 0.29012763500213623 + - 1.1434277296066284 + - 0.834911048412323 + - -0.6164064407348633 + - 0.0033871293999254704 + - 0.1835649162530899 + - 0.0 + - 0.44485998153686523 + - 0.693847119808197 + q01: + - -1.1566282510757446 + - -1.813583730836399e-05 + - -3.5173223295714706e-05 + - -1.493809461593628 + - -0.044773634523153305 + - -2.2744300365448 + - 0.0 + - -2.3844564566388726e-05 + - -1.9354150936123915e-05 + - -3.23197637044359e-05 + - -1.4539425373077393 + - -0.11438962072134018 + - -1.3077747821807861 + - 0.0 + - 0.0 + - 0.0 + q99: + - 2.0988931282772683e-05 + - 2.410673141479492 + - 2.2643818855285645 + - 1.7526803276268765e-05 + - 0.09094329178333282 + - 1.4207451343536377 + - 0.0 + - 0.9992868304252625 + - 2.5619561672210693 + - 2.3963286876678467 + - 2.7463351216283627e-05 + - 0.12538659572601318 + - 1.50703763961792 + - 0.0 + - 1.0 + - 1.0 + std: + - 0.2924744486808777 + - 0.9123055338859558 + - 0.6921584606170654 + - 0.5167360901832581 + - 0.024822130799293518 + - 0.7730752229690552 + - 0.0 + - 0.3022353947162628 + - 0.9363216161727905 + - 0.7197534441947937 + - 0.5159713625907898 + - 0.042254358530044556 + - 0.4932859241962433 + - 0.0 + - 0.4852803945541382 + - 0.43509986996650696 + count: + - 277350.0 + names: + - '0' + - '1' + - '2' + - '3' + - '4' + - '5' + - '6' + - '7' + - '8' + - '9' + - '10' + - '11' + - '12' + - '13' + - '14' + - '15' + mask: + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + state_stats: + mean: + - -0.28619107604026794 + - 1.1170974969863892 + - 0.8187518119812012 + - -0.6292704939842224 + - 0.002742639509961009 + - 0.38321176171302795 + - 0.0 + - 0.287190318107605 + - 1.1249585151672363 + - 0.8228866457939148 + - -0.6072733998298645 + - 0.003468270180746913 + - 0.18374495208263397 + - 0.0 + - 0.47009503841400146 + - 0.700372040271759 + std: + - 0.28622451424598694 + - 0.9085981845855713 + - 0.6900886297225952 + - 0.5145737528800964 + - 0.024824900552630424 + - 0.7737277746200562 + - 0.0 + - 0.3019956946372986 + - 0.9365604519844055 + - 0.720547616481781 + - 0.5160077214241028 + - 0.04248249530792236 + - 0.4887017011642456 + - 0.0 + - 0.4873206615447998 + - 0.4321308135986328 + q01: + - -0.991405725479126 + - -1.813583730836399e-05 + - -3.476843994576484e-05 + - -1.4938021898269653 + - -0.04476987197995186 + - -2.274427652359009 + - 0.0 + - -2.3844564566388726e-05 + - -1.9354150936123915e-05 + - -3.23197637044359e-05 + - -1.455243706703186 + - -0.11395290493965149 + - -1.3077747821807861 + - 0.0 + - 0.0 + - 0.0 + q99: + - 2.0988931282772683e-05 + - 2.410663366317749 + - 2.2642898559570312 + - 1.6971260265563615e-05 + - 0.09090737998485565 + - 1.4203816652297974 + - 0.0 + - 1.0001307725906372 + - 2.56237530708313 + - 2.397764205932617 + - 2.7463351216283627e-05 + - 0.12876388430595398 + - 1.4705802202224731 + - 0.0 + - 1.0 + - 1.0 + count: + - 277350.0 + names: + - '0' + - '1' + - '2' + - '3' + - '4' + - '5' + - '6' + - '7' + - '8' + - '9' + - '10' + - '11' + - '12' + - '13' + - '14' + - '15' + mask: + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + - true + repo_ids: + - rmbench_stage2 + norm_mode: q01_q99 +parallelism: + data_parallel_replicate_degree: 1 + enable_compiled_autograd: false + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: default + context_parallel_config: + degree: 1 + attention_type: ulysses + load_balancer: ulysses + head_stride: 1 + tensor_parallel_config: + degree: 1 + enable_async: false + data_parallel_config: + name: fsdp + param_dtype: null + reduce_dtype: float32 + num_replicas: null + shard_degree: null + wrapping_strategy: full + prefetch_factor: 0 + context_parallel_rotate_method: allgather +seed: 6198 +epoch: null +dry_run: false +ft_llm: true +ft_vit: false +ft_connector: true +ft_action_expert: true +ft_embedding: lm_head +optimizer: + name: adamw + learning_rate: 0.0001 + weight_decay: 0.01 + betas: + - 0.9 + - 0.95 + eps: 1.0e-05 + connector_learning_rate: 2.5e-05 + vit_learning_rate: 5.0e-06 + llm_learning_rate: 2.5e-05 + frame_selector_learning_rate: 0.0001 + temporal_token_scorer_learning_rate: 0.0001 + action_expert_learning_rate: 2.5e-05 + connector_weight_decay: 0.0 + vit_weight_decay: 0.0 + llm_weight_decay: 0.0 + frame_selector_weight_decay: 0.01 + temporal_token_scorer_weight_decay: 0.01 + action_expert_weight_decay: 0.0 + connector_betas: + - 0.9 + - 0.95 + vit_betas: + - 0.9 + - 0.95 + llm_betas: + - 0.9 + - 0.95 + frame_selector_betas: + - 0.9 + - 0.95 + temporal_token_scorer_betas: + - 0.9 + - 0.95 + action_expert_betas: + - 0.9 + - 0.95 + connector_eps: 1.0e-06 + vit_eps: 1.0e-06 + llm_eps: 1.0e-06 + frame_selector_eps: 1.0e-06 + temporal_token_scorer_eps: 1.0e-06 + action_expert_eps: 1.0e-06 + metrics_log_interval: -1 +scheduler: + name: multimodal + units: steps + t_warmup: 100 + t_max: null + alpha_f: 0.1 + connector_t_warmup: 200 + vit_t_warmup: 200 + llm_t_warmup: 200 + frame_selector_t_warmup: 200 + temporal_token_scorer_t_warmup: 200 + action_expert_t_warmup: 200 + grad_clip_warmup_steps: null + grad_clip_warmup_factor: null + warmup_min_lr: 0.0 +data: + dataset: null + mixture: null + root_size_mixture: null + kwargs_mixture: + - rate: 1.0 + datasets: + - dataset_name: lerobot:rmbench_stage2 + sampling_rate: null + root_size_factor: null + message_weight: null + override_p_high_res: null + name: lerobot:rmbench_stage2 + split: train + seed: 50189 + pad: null + sequence_length: 16384 + max_text_seq_len: null + crop_mode: null + shuffle: true + start_index: 0 + packing: null + enable_variable_sized_token_pooling: true + num_workers: 0 + drop_last: true + pin_memory: false + prefetch_factor: 4 + persistent_workers: false + timeout: 0 + skip_overlong_examples: true + skip_missing_vlm_examples: false +vlm_data: null +vlm_loader_rate: null +action_data: null +action_loader_rate: null +action_batch_interval: 1 +restore_dataloader: true +fast_forward_batches: null +evaluators: [] +eval_interval: 2000 +inf_evaluators: [] +inf_eval_interval: 2000 +save_inloop_predictions: true +eval_on_last_step: true +eval_on_load: false +eval_on: [] +save_folder: outputs/molmoact2-stage2 +reset_parameters: false +checkpointer_config: + save_thread_count: null + load_thread_count: null + pre_download: false + work_dir: null + throttle_uploads: false +canceled_check_interval: 50 +save_interval: 2000 +save_at: null +save_final_optim: true +save_num_checkpoints_to_keep: 1 +save_final_unsharded_checkpoint: false +save_merged_lora_checkpoint: true +save_interval_ephemeral: null +save_overwrite: false +load_path: null +checkpoint_load_strategy: auto +reset_optimizer_state: false +reset_trainer_state: false +initial_model_checkpoint: allenai/MolmoAct2 +allow_resume: true +max_duration: 100000 +global_train_batch_size: 4 +device_train_microbatch_size: 1 +max_grad_norm: 1.0 +multi_component_grad_norm: true +batch_divisor: global_batch +max_grad_norm_ratio: null +precision: amp_bf16 +wandb: + project: Robocasa + entity: ylinwu + group: null + name: molmoact2-stage2-dinov3-memory-reference-fixed-film-100k-v3 + tags: + - watching + log_artifacts: false + rank_zero_only: true + log_interval: 20 + allow_resume: true + finish_on_sigterm: true +beaker_log_interval: 50 +speed_monitor: + window_size: 20 + gpu_flops_available: null +console_log_interval: 20 +enable_timing_logs: false +gen1_gc_interval: 1 +compile: null +activation_checkpointing: true +fsdp: + fsdp2: true + precision: pure + use_orig_params: true + wrapping_strategy: null + sharding_strategy: FULL_SHARD + hybrid_sharding_num_model_replicas: null +softmax_auxiliary_loss: true +softmax_auxiliary_loss_scale: 0.0001 +response_logits_only: true +saliency_score_loss_wt: null +frame_score_loss_wt: null +frame_score_loss_type: mse +frame_score_loss_target: 0.7 +time_limit: null +extra_steps_after_cancel: 0 +python_profiling: false +torch_profiling: false +stop_at: 100000 +stop_after: null +fused_loss: false +compile_loss: true +distributed_eval: true diff --git a/simulation/rmbench/molmoact2/checkpoint/model.pt b/simulation/rmbench/molmoact2/checkpoint/model.pt new file mode 100644 index 0000000000000000000000000000000000000000..b4f0fbf2c2d8cb2d70cc955ee1ad3920dc886f17 --- /dev/null +++ b/simulation/rmbench/molmoact2/checkpoint/model.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ae0dacfcb2a45e82d18266ff5ed49edc26c35c41e038a43065c9a5671773360c +size 22747298964 diff --git a/simulation/rmbench/molmoact2/manifest.json b/simulation/rmbench/molmoact2/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..0c5169f4f2d80b6382c354dd2302708b83bcde02 --- /dev/null +++ b/simulation/rmbench/molmoact2/manifest.json @@ -0,0 +1,24 @@ +{ + "base_model": { + "repo": "allenai/MolmoAct2", + "revision": "e432d85f6e039edca44afb93c262f3084ab72a9c" + }, + "checkpoint": { + "path": "checkpoint/model.pt", + "sha256": "ae0dacfcb2a45e82d18266ff5ed49edc26c35c41e038a43065c9a5671773360c" + }, + "checkpoint_config": { + "path": "checkpoint/config.yaml", + "sha256": "ed4573be786a951ccb3319a090b8443db03acb1a204c827327eb6164ecfb905d" + }, + "checkpoint_step": 100000, + "eval_assets_path_prefix": "simulation/rmbench/molmoact2_eval_assets", + "eval_assets_repo": "Amazon-FAR/imprint_datasets", + "hf_repo": "Amazon-FAR/imprint_models", + "memory_enabled": true, + "offline_visual_features": false, + "path_prefix": "simulation/rmbench/molmoact2", + "schema_version": 1, + "source_run": "mit0jkx7", + "visual_history": "online_wan_and_dinov3" +} diff --git a/simulation/rmbench/pi0/README.md b/simulation/rmbench/pi0/README.md new file mode 100644 index 0000000000000000000000000000000000000000..f49d47e4014fe8f6594f6573b29c7b0c8d696302 --- /dev/null +++ b/simulation/rmbench/pi0/README.md @@ -0,0 +1,34 @@ +# Native Pi0 Memory Policy + +This directory contains the native Pi0 IMPRINT checkpoint used for the +RMBench backbone comparison. + +## Configuration + +- Policy backbone: native Pi0 +- Memory encoder: full-attention IMPRINT encoder +- Visual history: pooled DINOv3 features +- Stage-1 source: `j8t4ebpo`, step 160,000 +- Stage-2 source: `s746sl7k`, step 100,000 +- Memory bottleneck: 256 tokens of width 1024 +- Action horizon: 50 +- Image encoder: frozen in Stage 2 +- Memory encoder: finetuned in Stage 2 + +The `stage2/step-100000/` directory is a parameter-only Orbax checkpoint. +Optimizer and training state are intentionally omitted. + +Download this prefix with: + +```bash +hf download Amazon-FAR/imprint_models \ + --include "simulation/rmbench/pi0/**" \ + --local-dir artifacts/downloads/imprint_models +``` + +The matching training store is +`Amazon-FAR/imprint_datasets/simulation/rmbench/450/`. The released code and +evaluation harness are on the `pi0` branch of `Amazon-FAR/imprint`. + +The checkpoint contains Gemma/PaliGemma-derived weights. The Gemma Terms of +Use included with the IMPRINT code release apply. diff --git a/simulation/rmbench/pi0/SHA256SUMS b/simulation/rmbench/pi0/SHA256SUMS new file mode 100644 index 0000000000000000000000000000000000000000..3d00f23bdf456dc2f8b7f290fd95dcb56e466c4e --- /dev/null +++ b/simulation/rmbench/pi0/SHA256SUMS @@ -0,0 +1,23 @@ +121e80bc7d65582e61c44bf97cbbc7fae2e5245ec97fd62d75e00e90c8658010 README.md +a1f32c34a9c451cded8682b8e2ac484fdde2df3af806b4f262af8703429bf66e config/history_config.yaml +38a89ec49ceb4bd397b0792fe5704441588efa8aaa024bc6c66669c556500091 config/input_provenance.json +372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994 config/norm_stats.json +b0746fd94e92ebb8ffdc73a102e934202238c4d2d558f0b20ff878230acc409b manifest.json +271ec90649ef2f91404db80c799e01e60601a05510f4bb7d9ffcc166e48edce9 stage2/step-100000/_CHECKPOINT_METADATA +372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994 stage2/step-100000/assets/robomme/norm_stats.json +c52a0aaf801528076bb3c273e81100f81908e8b0b50a1e9b083e6f4a33eafd9d stage2/step-100000/params/_METADATA +307790b108b2f99519fda483bfedad812961006d7d0a21c6aa70149beb683326 stage2/step-100000/params/_sharding +e1ff7698d11457408eb93f14a5bd654c759102c00cceb45648070133409bcd2f stage2/step-100000/params/array_metadatas/process_0 +85ba1afab930ed68585aab9431d611d73a1c3c09e6703557b018320f80eace4c stage2/step-100000/params/d/3bd5e7a25312cf2bc1713db9a1d29fa5 +7b71a0eed34c83b3c408232fb342305a3e61df49b44cd84a8c59e71be40cf73e stage2/step-100000/params/manifest.ocdbt +5572405112e8e9a60f6755ee127bfb3dc54581e38c9328c1fddb731d89b9b614 stage2/step-100000/params/ocdbt.process_0/d/181a1cb7822e6e4ae74e5d52fd4c3e2c +6e8e872ea3a0a155be20628b8bb17893de384a3d06de2f924bfec7b4326b9f17 stage2/step-100000/params/ocdbt.process_0/d/24d3ba3b910a02ad89c9e9c5ee773b7d +b0e00f5770da3b87b76895a636beb69e089d6e96ecb77b3ffcde58de4412cd16 stage2/step-100000/params/ocdbt.process_0/d/3c822efa587f8375db55af5ad6ff4923 +e07b87ad1f9abfd99d86ec5f9b26ca7089caf89b8f859ee3902233018afa736e stage2/step-100000/params/ocdbt.process_0/d/4371d349f1ebc28e5203786efdbb69cc +386bb1cf6de7f6a3011524c526bd2c60965f5a2df2e3bcc4522e8d34924a5e82 stage2/step-100000/params/ocdbt.process_0/d/5c003ca5eacb776438fd966f8fe51bdf +4130dbf0683ef1f0f92ffee27c06d695f8360b0895364228a33591c66fc5f298 stage2/step-100000/params/ocdbt.process_0/d/6f0b57f972221d7c08e1c6e0542b0f59 +034c510b6b4a7b7a01bff53cd076593b5f8e993f6497015f311517b1feced6d7 stage2/step-100000/params/ocdbt.process_0/d/7cf131af294406195d880b45d0b24cdb +0fcbfee171ad9ec14676b5ed65d48e5fa9f6c88c9eff4edbc260be66e120d754 stage2/step-100000/params/ocdbt.process_0/d/b434ecacfeead5956996de4c9f1c8f19 +3a1cabf9535f848b13bd8013648130093f9937f4c1a3fe95e23ec513bbcd53b9 stage2/step-100000/params/ocdbt.process_0/d/bc0fe1585dbbc9796ccd89ef1cb53d7e +03e1f1fc55a7493b82d17d583063292ec0e0f1cd8b7d796a13e7cbc8cbb470c4 stage2/step-100000/params/ocdbt.process_0/d/e364328538b1fe570bf0733fb7086ce5 +8bf80ac089f1ca835aedcafd735621ea3c371f10fc19b42efd0a05c371ddced1 stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt diff --git a/simulation/rmbench/pi0/config/history_config.yaml b/simulation/rmbench/pi0/config/history_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a15cf28c4c117331a4bb823b23af04999683ecbc --- /dev/null +++ b/simulation/rmbench/pi0/config/history_config.yaml @@ -0,0 +1,67 @@ +budget: 256 +num_views: 3 +streaming_obs_horizon: 16 +integration_type: modulation +memory_token_dim: 1024 +representation_type: cosmo_bottleneck +memory_feature: + img: + net: identity + input_dim: 768 + pos: + input_dim: 768 + hidden_dim: 768 + state: + input_dim: 16 + hidden_dim: 512 +cosmo_bottleneck_memory: + visual_feature_stream: dinov3 + feature_store_path: artifacts/datasets/rmbench/stage1 + latent_source: offline + n_variants: 8 + num_query_samples: 16 + variant: Wan2.2-TI2V-5B + temporal_factor: 4 + downsample: 2 + num_latents: 256 + num_cross_blocks: 2 + num_heads: 8 + cosmos_latent_dim: 48 + siglip_dim: 768 + proprio_in_dim: 16 + proprio_encoding: raw + action_in_dim: 16 + action_encoding: raw + lang_emb_dim: 2048 + action_tokenizer_hidden: 512 + proprio_tokenizer_hidden: 512 + pool_h: 2 + pool_w: 2 + wan_h: 256 + wan_w: 256 + latent_cap: 4096 + max_recent_tokens: 768 + max_step_tokens: 1244 +cosmo_decoder: + num_heads: 8 + num_blocks: 2 + align_dim: 768 + align_num_patches: 192 + ffn_inner: 8192 + action_dim: 16 + proprio_dim: 16 + max_query_steps: 2048 +frame_jitter: + enabled: false + start_jitter: 3 + stride_choices: + - 1 + - 2 + - 3 + interval_jitter: 1 + canonical_downsample: 2 +stream_aug: + enabled: false + gauss_sigma: 0.05 + stream_drop_p: 0.1 + feat_drop_p: 0.0 diff --git a/simulation/rmbench/pi0/config/input_provenance.json b/simulation/rmbench/pi0/config/input_provenance.json new file mode 100644 index 0000000000000000000000000000000000000000..9a66b58c9a19045427a9830013dc2f5559896e45 --- /dev/null +++ b/simulation/rmbench/pi0/config/input_provenance.json @@ -0,0 +1,28 @@ +{ + "action_causality": "strictly_causal", + "align_num_patches": 192, + "checkpoint_hf_path": "simulation/rmbench/pi0/stage2/step-100000", + "checkpoint_hf_repo": "Amazon-FAR/imprint_models", + "checkpoint_metadata_sha256": "271ec90649ef2f91404db80c799e01e60601a05510f4bb7d9ffcc166e48edce9", + "checkpoint_step": 100000, + "eval_consolidate_all": 1, + "exact_source_sha256": "ffacf4518beb323c45972975d4b368d953ead306d9389a4a745f3f4d9d0047f6", + "history_config_sha256": "e3fd0c5bfd4d7f8f7455e9105a4e2ec0703bdbda67c50c16678288816fa4ba67", + "max_recent_tokens": 768, + "max_step_tokens": 1244, + "memory_language_source": "live", + "model_config_name": "mme_wan_stage2_pi0_memft", + "model_history_config": "cosmo-wan-bottleneck-modul-rmbench-all9-dinov3-memft.yaml", + "model_type": "pi0", + "num_latents": 256, + "pi05": false, + "stage1_checkpoint": "j8t4ebpo_160k", + "stage1_checkpoint_hf_path": "simulation/rmbench/full_attention/stage1/stage1_step160000.npz", + "stage1_checkpoint_hf_repo": "Amazon-FAR/imprint_models", + "stage2_memory_encoder": "finetuned", + "training_run": "s746sl7k", + "training_store_hf_path": "simulation/rmbench/450", + "training_store_hf_repo": "Amazon-FAR/imprint_datasets", + "visual_feature_stream": "dinov3", + "visual_width": 768 +} diff --git a/simulation/rmbench/pi0/config/norm_stats.json b/simulation/rmbench/pi0/config/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..631b84bab53fc384b33fa5468192db069a535b6b --- /dev/null +++ b/simulation/rmbench/pi0/config/norm_stats.json @@ -0,0 +1,152 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + -0.2884466350078583, + 1.116854190826416, + 0.8176121115684509, + -0.6286301016807556, + 0.0027539930306375027, + 0.3858006000518799, + 0.0, + 0.29012763500213623, + 1.1434277296066284, + 0.834911048412323, + -0.6164064407348633, + 0.0033871293999254704, + 0.1835649162530899, + 0.0, + 0.44485998153686523, + 0.693847119808197 + ], + "std": [ + 0.2924744486808777, + 0.9123055338859558, + 0.6921584606170654, + 0.5167360901832581, + 0.024822130799293518, + 0.7730752229690552, + 0.0, + 0.3022353947162628, + 0.9363216161727905, + 0.7197534441947937, + 0.5159713625907898, + 0.042254358530044556, + 0.4932859241962433, + 0.0, + 0.4852803945541382, + 0.43509986996650696 + ], + "q01": [ + -1.1566282510757446, + -0.00001813583730836399, + -0.000035173223295714706, + -1.493809461593628, + -0.044773634523153305, + -2.2744300365448, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.4539425373077393, + -0.11438962072134018, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410673141479492, + 2.2643818855285645, + 0.000017526803276268765, + 0.09094329178333282, + 1.4207451343536377, + 0.0, + 0.9992868304252625, + 2.5619561672210693, + 2.3963286876678467, + 0.000027463351216283627, + 0.12538659572601318, + 1.50703763961792, + 0.0, + 1.0, + 1.0 + ] + }, + "state": { + "mean": [ + -0.28619107604026794, + 1.1170974969863892, + 0.8187518119812012, + -0.6292704939842224, + 0.002742639509961009, + 0.38321176171302795, + 0.0, + 0.287190318107605, + 1.1249585151672363, + 0.8228866457939148, + -0.6072733998298645, + 0.003468270180746913, + 0.18374495208263397, + 0.0, + 0.47009503841400146, + 0.700372040271759 + ], + "std": [ + 0.28622451424598694, + 0.9085981845855713, + 0.6900886297225952, + 0.5145737528800964, + 0.024824900552630424, + 0.7737277746200562, + 0.0, + 0.3019956946372986, + 0.9365604519844055, + 0.720547616481781, + 0.5160077214241028, + 0.04248249530792236, + 0.4887017011642456, + 0.0, + 0.4873206615447998, + 0.4321308135986328 + ], + "q01": [ + -0.991405725479126, + -0.00001813583730836399, + -0.00003476843994576484, + -1.4938021898269653, + -0.04476987197995186, + -2.274427652359009, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.455243706703186, + -0.11395290493965149, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410663366317749, + 2.2642898559570312, + 0.000016971260265563615, + 0.09090737998485565, + 1.4203816652297974, + 0.0, + 1.0001307725906372, + 2.56237530708313, + 2.397764205932617, + 0.000027463351216283627, + 0.12876388430595398, + 1.4705802202224731, + 0.0, + 1.0, + 1.0 + ] + } + } +} \ No newline at end of file diff --git a/simulation/rmbench/pi0/manifest.json b/simulation/rmbench/pi0/manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..12040a5382a295ad2d880486b72e726984c14db3 --- /dev/null +++ b/simulation/rmbench/pi0/manifest.json @@ -0,0 +1,26 @@ +{ + "action_horizon": 50, + "architecture": "perceiver_v1", + "checkpoint_format": "orbax_params_only", + "checkpoint_step": 100000, + "hf_repo": "Amazon-FAR/imprint_models", + "image_encoder": "frozen", + "memory_encoder": "finetuned", + "memory_width": 1024, + "model_type": "pi0", + "num_latents": 256, + "path_prefix": "simulation/rmbench/pi0", + "payload_bytes": 11534439221, + "payload_file_count": 22, + "schema_version": 1, + "source_eval_bundle": "s746sl7k_100k_pi0_dinov3_memft_causal_c1", + "source_run": "s746sl7k", + "stage1_checkpoint_step": 160000, + "stage1_source_run": "j8t4ebpo", + "train_state_included": false, + "training_store": { + "path_prefix": "simulation/rmbench/450", + "repo": "Amazon-FAR/imprint_datasets" + }, + "visual_history": "dinov3" +} diff --git a/simulation/rmbench/pi0/stage2/step-100000/_CHECKPOINT_METADATA b/simulation/rmbench/pi0/stage2/step-100000/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..cf6914bee1c20f1a3828a67025608a1084699acb --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1787699801100333867, "commit_timestamp_nsecs": 1787699962354272918, "custom_metadata": {}} \ No newline at end of file diff --git a/simulation/rmbench/pi0/stage2/step-100000/assets/robomme/norm_stats.json b/simulation/rmbench/pi0/stage2/step-100000/assets/robomme/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..631b84bab53fc384b33fa5468192db069a535b6b --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/assets/robomme/norm_stats.json @@ -0,0 +1,152 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + -0.2884466350078583, + 1.116854190826416, + 0.8176121115684509, + -0.6286301016807556, + 0.0027539930306375027, + 0.3858006000518799, + 0.0, + 0.29012763500213623, + 1.1434277296066284, + 0.834911048412323, + -0.6164064407348633, + 0.0033871293999254704, + 0.1835649162530899, + 0.0, + 0.44485998153686523, + 0.693847119808197 + ], + "std": [ + 0.2924744486808777, + 0.9123055338859558, + 0.6921584606170654, + 0.5167360901832581, + 0.024822130799293518, + 0.7730752229690552, + 0.0, + 0.3022353947162628, + 0.9363216161727905, + 0.7197534441947937, + 0.5159713625907898, + 0.042254358530044556, + 0.4932859241962433, + 0.0, + 0.4852803945541382, + 0.43509986996650696 + ], + "q01": [ + -1.1566282510757446, + -0.00001813583730836399, + -0.000035173223295714706, + -1.493809461593628, + -0.044773634523153305, + -2.2744300365448, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.4539425373077393, + -0.11438962072134018, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410673141479492, + 2.2643818855285645, + 0.000017526803276268765, + 0.09094329178333282, + 1.4207451343536377, + 0.0, + 0.9992868304252625, + 2.5619561672210693, + 2.3963286876678467, + 0.000027463351216283627, + 0.12538659572601318, + 1.50703763961792, + 0.0, + 1.0, + 1.0 + ] + }, + "state": { + "mean": [ + -0.28619107604026794, + 1.1170974969863892, + 0.8187518119812012, + -0.6292704939842224, + 0.002742639509961009, + 0.38321176171302795, + 0.0, + 0.287190318107605, + 1.1249585151672363, + 0.8228866457939148, + -0.6072733998298645, + 0.003468270180746913, + 0.18374495208263397, + 0.0, + 0.47009503841400146, + 0.700372040271759 + ], + "std": [ + 0.28622451424598694, + 0.9085981845855713, + 0.6900886297225952, + 0.5145737528800964, + 0.024824900552630424, + 0.7737277746200562, + 0.0, + 0.3019956946372986, + 0.9365604519844055, + 0.720547616481781, + 0.5160077214241028, + 0.04248249530792236, + 0.4887017011642456, + 0.0, + 0.4873206615447998, + 0.4321308135986328 + ], + "q01": [ + -0.991405725479126, + -0.00001813583730836399, + -0.00003476843994576484, + -1.4938021898269653, + -0.04476987197995186, + -2.274427652359009, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.455243706703186, + -0.11395290493965149, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410663366317749, + 2.2642898559570312, + 0.000016971260265563615, + 0.09090737998485565, + 1.4203816652297974, + 0.0, + 1.0001307725906372, + 2.56237530708313, + 2.397764205932617, + 0.000027463351216283627, + 0.12876388430595398, + 1.4705802202224731, + 0.0, + 1.0, + 1.0 + ] + } + } +} \ No newline at end of file diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/_METADATA b/simulation/rmbench/pi0/stage2/step-100000/params/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..2632c3a62bcd33903496a7fb49b3ef30a3bf8548 --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_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": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1152, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 16, 72, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [14, 14, 3, 288]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "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": [1152, 512]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64288, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_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": [512]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'kv_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "kv_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'mem_rms_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "mem_rms_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": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'out_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "out_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'q_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "q_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 2048, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4096, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_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": [256, 32]}}, "('params', 'action_time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'action_time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'lang_film', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'lang_film', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('params', 'mem_encoder', 'latents', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "latents", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1024]}}, "('params', 'mem_encoder', 'proj_cosmo', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_cosmo', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "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, 1024]}}, "('params', 'mem_encoder', 'proj_recent', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_recent', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 1024]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'state_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "state_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'state_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "state_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/_sharding b/simulation/rmbench/pi0/stage2/step-100000/params/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..17cef6914bbaa8f44ee48bad18fe6f5e3279fa2a --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/_sharding @@ -0,0 +1 @@ +{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm1lbV9ybXNfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm91dF9laW5zdW1fbWVtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLmt2X2VpbnN1bV9tZW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLnFfZWluc3VtX21lbS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhdGVudHMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnN0YXRlX3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnN0YXRlX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}"} \ No newline at end of file diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/array_metadatas/process_0 b/simulation/rmbench/pi0/stage2/step-100000/params/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..f711d91a9bde79640d54be0928469c3b5b532a46 --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 1152, 1076], "chunk_shape": [27, 1152, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 16, 72, 288], "chunk_shape": [27, 16, 72, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [1152, 512], "chunk_shape": [1152, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 8, 256, 512], "chunk_shape": [18, 8, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 8, 512, 256], "chunk_shape": [18, 8, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.kv_einsum_mem.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.mem_rms_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.out_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.q_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.bias.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.kernel.value", "write_shape": [18, 1024, 512], "chunk_shape": [18, 1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 2048, 4096], "chunk_shape": [18, 2, 2048, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 1024, 1024], "chunk_shape": [18, 2, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_time_mlp_in.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_time_mlp_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.latents.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.kernel.value", "write_shape": [12, 1024], "chunk_shape": [12, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.kernel.value", "write_shape": [192, 1024], "chunk_shape": [192, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.state_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.state_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}]} \ No newline at end of file diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/d/3bd5e7a25312cf2bc1713db9a1d29fa5 b/simulation/rmbench/pi0/stage2/step-100000/params/d/3bd5e7a25312cf2bc1713db9a1d29fa5 new file mode 100644 index 0000000000000000000000000000000000000000..8f2c7764328b8ab5a20e5da5d1b602b99eb9a03a --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/d/3bd5e7a25312cf2bc1713db9a1d29fa5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85ba1afab930ed68585aab9431d611d73a1c3c09e6703557b018320f80eace4c +size 157877 diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/manifest.ocdbt b/simulation/rmbench/pi0/stage2/step-100000/params/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..5b16321dac5c6348ea49a31e08139f9e620a8e82 Binary files /dev/null and b/simulation/rmbench/pi0/stage2/step-100000/params/manifest.ocdbt differ diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/181a1cb7822e6e4ae74e5d52fd4c3e2c b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/181a1cb7822e6e4ae74e5d52fd4c3e2c new file mode 100644 index 0000000000000000000000000000000000000000..f2ffb2d725a935e177488beecc9048246729d8fd --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/181a1cb7822e6e4ae74e5d52fd4c3e2c @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5572405112e8e9a60f6755ee127bfb3dc54581e38c9328c1fddb731d89b9b614 +size 2240440598 diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/24d3ba3b910a02ad89c9e9c5ee773b7d b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/24d3ba3b910a02ad89c9e9c5ee773b7d new file mode 100644 index 0000000000000000000000000000000000000000..c55c1151b8f03d8d3d4afb69dadc31618fd9ab61 --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/24d3ba3b910a02ad89c9e9c5ee773b7d @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6e8e872ea3a0a155be20628b8bb17893de384a3d06de2f924bfec7b4326b9f17 +size 995921598 diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/3c822efa587f8375db55af5ad6ff4923 b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/3c822efa587f8375db55af5ad6ff4923 new file mode 100644 index 0000000000000000000000000000000000000000..d618b4ba57ec44065e9552110e52bfb66228797d --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/3c822efa587f8375db55af5ad6ff4923 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b0e00f5770da3b87b76895a636beb69e089d6e96ecb77b3ffcde58de4412cd16 +size 673740 diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/4371d349f1ebc28e5203786efdbb69cc b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/4371d349f1ebc28e5203786efdbb69cc new file mode 100644 index 0000000000000000000000000000000000000000..e4ad2fd53434229fdb414290f9c93c99aa31938e Binary files /dev/null and b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/4371d349f1ebc28e5203786efdbb69cc differ diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/5c003ca5eacb776438fd966f8fe51bdf b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/5c003ca5eacb776438fd966f8fe51bdf new file mode 100644 index 0000000000000000000000000000000000000000..eb309d75042521bb090470201c4984755cd098d7 --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/5c003ca5eacb776438fd966f8fe51bdf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:386bb1cf6de7f6a3011524c526bd2c60965f5a2df2e3bcc4522e8d34924a5e82 +size 1120292288 diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/6f0b57f972221d7c08e1c6e0542b0f59 b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/6f0b57f972221d7c08e1c6e0542b0f59 new file mode 100644 index 0000000000000000000000000000000000000000..58766aab8cdd6cc2447a8f4e785f790080bc0c62 --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/6f0b57f972221d7c08e1c6e0542b0f59 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4130dbf0683ef1f0f92ffee27c06d695f8360b0895364228a33591c66fc5f298 +size 2240243222 diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/7cf131af294406195d880b45d0b24cdb b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/7cf131af294406195d880b45d0b24cdb new file mode 100644 index 0000000000000000000000000000000000000000..2a0b58f2c74a2f8b9539cde943dce7b6fad34243 Binary files /dev/null and b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/7cf131af294406195d880b45d0b24cdb differ diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/b434ecacfeead5956996de4c9f1c8f19 b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/b434ecacfeead5956996de4c9f1c8f19 new file mode 100644 index 0000000000000000000000000000000000000000..82e70003f4924735684c6c69d8fca8ad5e67fbc3 Binary files /dev/null and b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/b434ecacfeead5956996de4c9f1c8f19 differ diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/bc0fe1585dbbc9796ccd89ef1cb53d7e b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/bc0fe1585dbbc9796ccd89ef1cb53d7e new file mode 100644 index 0000000000000000000000000000000000000000..1efba3168a2f9c0dba642ea1d8d099fd8e60d904 --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/bc0fe1585dbbc9796ccd89ef1cb53d7e @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3a1cabf9535f848b13bd8013648130093f9937f4c1a3fe95e23ec513bbcd53b9 +size 2348121321 diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/e364328538b1fe570bf0733fb7086ce5 b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/e364328538b1fe570bf0733fb7086ce5 new file mode 100644 index 0000000000000000000000000000000000000000..e78ac98589a6873f59408f67758d3161bed6a124 --- /dev/null +++ b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/e364328538b1fe570bf0733fb7086ce5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:03e1f1fc55a7493b82d17d583063292ec0e0f1cd8b7d796a13e7cbc8cbb470c4 +size 2588466495 diff --git a/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..ef823e7f17b055403e0dce78d128c0d887242db9 Binary files /dev/null and b/simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt differ diff --git a/simulation/rmbench/recurrent/.gitattributes b/simulation/rmbench/recurrent/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..575804fdf56c218a3da1b6d6a68613a6be31d7ae --- /dev/null +++ b/simulation/rmbench/recurrent/.gitattributes @@ -0,0 +1,44 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz 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 +stage2/step-100000/params/d/cbf2fef1100ba6c460caea5d0023e183 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/3a109d42459c2150f8b988ce2442e2b5 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/fb91f39ce5a047cdcaebc0bad36caad6 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/7dc0a88b6d199ccbb016750c98012cf1 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/88cd20aa3cd4faa14d8123fbe93a44b9 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/b7327531c46c5cc244cc7981e9f4da90 filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/fa7cea5bceef1307d13d9a61ea91a35e filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/fa71f45e6411d633e6827c04a989a92f filter=lfs diff=lfs merge=lfs -text +stage2/step-100000/params/ocdbt.process_0/d/054944d2cae72ec6dcde8844c294284c filter=lfs diff=lfs merge=lfs -text diff --git a/simulation/rmbench/recurrent/LICENSE_GEMMA.txt b/simulation/rmbench/recurrent/LICENSE_GEMMA.txt new file mode 100644 index 0000000000000000000000000000000000000000..e0f8930d83fe9eef924c0eed50ee0e572eb471fa --- /dev/null +++ b/simulation/rmbench/recurrent/LICENSE_GEMMA.txt @@ -0,0 +1,77 @@ +Gemma Terms of Use + +Last modified: February 21, 2024 + +By using, reproducing, modifying, distributing, performing or displaying any portion or element of Gemma, Model Derivatives including via any Hosted Service, (each as defined below) (collectively, the "Gemma Services") or otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement. + +Section 1: DEFINITIONS +1.1 Definitions +(a) "Agreement" or "Gemma Terms of Use" means these terms and conditions that govern the use, reproduction, Distribution or modification of the Gemma Services and any terms and conditions incorporated by reference. + +(b) "Distribution" or "Distribute" means any transmission, publication, or other sharing of Gemma or Model Derivatives to a third party, including by providing or making Gemma or its functionality available as a hosted service via API, web access, or any other electronic or remote means ("Hosted Service"). + +(c) "Gemma" means the set of machine learning language models, trained model weights and parameters identified at ai.google.dev/gemma, regardless of the source that you obtained it from. + +(d) "Google" means Google LLC. + +(e) "Model Derivatives" means all (i) modifications to Gemma, (ii) works based on Gemma, or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Gemma, to that model in order to cause that model to perform similarly to Gemma, including distillation methods that use intermediate data representations or methods based on the generation of synthetic data Outputs by Gemma for training that model. For clarity, Outputs are not deemed Model Derivatives. + +(f) "Output" means the information content output of Gemma or a Model Derivative that results from operating or otherwise using Gemma or the Model Derivative, including via a Hosted Service. + +1.2 +As used in this Agreement, "including" means "including without limitation". + +Section 2: ELIGIBILITY AND USAGE +2.1 Eligibility +You represent and warrant that you have the legal capacity to enter into this Agreement (including being of sufficient age of consent). If you are accessing or using any of the Gemma Services for or on behalf of a legal entity, (a) you are entering into this Agreement on behalf of yourself and that legal entity, (b) you represent and warrant that you have the authority to act on behalf of and bind that entity to this Agreement and (c) references to "you" or "your" in the remainder of this Agreement refers to both you (as an individual) and that entity. + +2.2 Use +You may use, reproduce, modify, Distribute, perform or display any of the Gemma Services only in accordance with the terms of this Agreement, and must not violate (or encourage or permit anyone else to violate) any term of this Agreement. + +Section 3: DISTRIBUTION AND RESTRICTIONS +3.1 Distribution and Redistribution +You may reproduce or Distribute copies of Gemma or Model Derivatives if you meet all of the following conditions: + +You must include the use restrictions referenced in Section 3.2 as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Gemma or Model Derivatives and you must provide notice to subsequent users you Distribute to that Gemma or Model Derivatives are subject to the use restrictions in Section 3.2. +You must provide all third party recipients of Gemma or Model Derivatives a copy of this Agreement. +You must cause any modified files to carry prominent notices stating that you modified the files. +All Distributions (other than through a Hosted Service) must be accompanied by a "Notice" text file that contains the following notice: "Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms". +You may add your own intellectual property statement to your modifications and, except as set forth in this Section, may provide additional or different terms and conditions for use, reproduction, or Distribution of your modifications, or for any such Model Derivatives as a whole, provided your use, reproduction, modification, Distribution, performance, and display of Gemma otherwise complies with the terms and conditions of this Agreement. Any additional or different terms and conditions you impose must not conflict with the terms of this Agreement. + +3.2 Use Restrictions +You must not use any of the Gemma Services: + +for the restricted uses set forth in the Gemma Prohibited Use Policy at ai.google.dev/gemma/prohibited_use_policy ("Prohibited Use Policy"), which is hereby incorporated by reference into this Agreement; or +in violation of applicable laws and regulations. +To the maximum extent permitted by law, Google reserves the right to restrict (remotely or otherwise) usage of any of the Gemma Services that Google reasonably believes are in violation of this Agreement. + +3.3 Generated Output +Google claims no rights in Outputs you generate using Gemma. You and your users are solely responsible for Outputs and their subsequent uses. + +Section 4: ADDITIONAL PROVISIONS +4.1 Updates +Google may update Gemma from time to time, and you must make reasonable efforts to use the latest version of Gemma. + +4.2 Trademarks +Nothing in this Agreement grants you any rights to use Google's trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between you and Google. Google reserves any rights not expressly granted herein. + +4.3 DISCLAIMER OF WARRANTY +UNLESS REQUIRED BY APPLICABLE LAW, THE GEMMA SERVICES, AND OUTPUTS, ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY WARRANTIES OR CONDITIONS OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR OR DISTRIBUTING ANY OF THE GEMMA SERVICES OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR USE OR DISTRIBUTION OF ANY OF THE GEMMA SERVICES OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT. + +4.4 LIMITATION OF LIABILITY +TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY, CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW, SHALL GOOGLE OR ITS AFFILIATES BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL, OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO, ANY OF THE GEMMA SERVICES OR OUTPUTS EVEN IF GOOGLE OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. + +4.5 Term, Termination, and Survival +The term of this Agreement will commence upon your acceptance of this Agreement (including acceptance by your use, modification, or Distribution, reproduction, performance or display of any portion or element of the Gemma Services) and will continue in full force and effect until terminated in accordance with the terms of this Agreement. Google may terminate this Agreement if you are in breach of any term of this Agreement. Upon termination of this Agreement, you must delete and cease use and Distribution of all copies of Gemma and Model Derivatives in your possession or control. Sections 1, 2.1, 3.3, 4.2 to 4.9 shall survive the termination of this Agreement. + +4.6 Governing Law and Jurisdiction +This Agreement will be governed by the laws of the State of California without regard to choice of law principles. The UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The state and federal courts of Santa Clara County, California shall have exclusive jurisdiction of any dispute arising out of this Agreement. + +4.7 Severability +If any provision of this Agreement is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein. + +4.8 Entire Agreement +This Agreement states all the terms agreed between the parties and supersedes all other agreements between the parties as of the date of acceptance relating to its subject matter. + +4.9 No Waiver +Google will not be treated as having waived any rights by not exercising (or delaying the exercise of) any rights under this Agreement. \ No newline at end of file diff --git a/simulation/rmbench/recurrent/NOTICE b/simulation/rmbench/recurrent/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..1b9cb6541d7135c4c78a542ac4a3e04a971c96db --- /dev/null +++ b/simulation/rmbench/recurrent/NOTICE @@ -0,0 +1,11 @@ +IMPRINT +Copyright 2026 + +This distribution contains modifications to the OpenPI project: +https://github.com/Physical-Intelligence/openpi + +OpenPI is distributed under the Apache License 2.0. Gemma components and +derivatives are subject to the Gemma Terms of Use in LICENSE_GEMMA.txt. + +Gemma is provided under and subject to the Gemma Terms of Use found at +ai.google.dev/gemma/terms. diff --git a/simulation/rmbench/recurrent/README.md b/simulation/rmbench/recurrent/README.md new file mode 100644 index 0000000000000000000000000000000000000000..75a7f16cc673f099e950b8916c79a971042667e2 --- /dev/null +++ b/simulation/rmbench/recurrent/README.md @@ -0,0 +1,38 @@ +--- +license: gemma +library_name: jax +tags: + - robotics + - vision-language-action + - memory + - rmbench + - dinov3 + - recurrent +--- + +# IMPRINT RMBench Recurrent Checkpoint + +Canonical Pi0.5 IMPRINT checkpoint for RMBench using the recurrent gated-delta +memory encoder. + +- Stage 1: 160,000 updates, W&B `liqhuiha/Robocasa/9bm5uvc7` +- Stage 2: 100,000 updates, W&B `liqhuiha/Robocasa/0f4k18rw` +- memory tokens: 256 x 1024 +- policy horizon: 50 +- training scan: chunkwise WY +- inference update: incremental recurrent +- Stage 2: memory encoder finetuned; image tower frozen + +The release contains the Stage-1 encoder, the complete Stage-2 Orbax parameter +checkpoint (including the PaliGemma/Gemma backbone), normalization statistics, +portable history configuration, and provenance. The checkpoint is a Gemma +model derivative and is distributed subject to the included Gemma Terms of +Use. + +Use the matching `yilin-wu/imprint-rmbench-450` dataset and the evaluation +contract in the IMPRINT GitHub repository. Evaluation requires strict causal +actions, online recent DINOv3 features, `EVAL_CONSOLIDATE_ALL=1`, BGR policy +inputs, and 50-step prediction/execution chunks. + +The matching training store and reproducibility code are publicly available +under the IMPRINT release. diff --git a/simulation/rmbench/recurrent/SHA256SUMS b/simulation/rmbench/recurrent/SHA256SUMS new file mode 100644 index 0000000000000000000000000000000000000000..af742e649405d040e92a77a3b73e0b9cbade3275 --- /dev/null +++ b/simulation/rmbench/recurrent/SHA256SUMS @@ -0,0 +1,29 @@ +3e2c24001f9ef57bf7ec959a3658fbb49cdad113cdf394c264da9d16f9bdd132 ./LICENSE_GEMMA.txt +5b21d7116ac42a0ab6141488cb654c2411557ac977e2276accb88ef774f28c19 ./NOTICE +ad0a0597a9e8433ca73ddafe3e9ae5869845525f1e98a6b048f6956252c8ea51 ./README.md +aa75ca05051d41efc9d146c1b022e15759f153a078f60bb6ae73c3eb59f80fdb ./config/history_config.yaml +1b63e66aa24c04f109fe0764b25d0dbf1d9e4425330eca4957fd8218d7a06b9a ./config/input_provenance.json +372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994 ./config/norm_stats.json +76f7ecea50d65c374c79bba071074fc35aded1dea22846d86cfe90b772f4bea3 ./stage1/stage1_step160000.npz +dd12980babae314a5097df31c47da055e3ced591c4fc44e1a394442365cb4323 ./stage2/step-100000/_CHECKPOINT_METADATA +372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994 ./stage2/step-100000/assets/robomme/norm_stats.json +9f10bcbbb6a6fe31bee656d8abbeb81270c56bece5a5cd1e9fd2b002070cd272 ./stage2/step-100000/params/_METADATA +306825e3332fd927d96181a73228f3480a3ceb705c7db27fe49a2563204ebb8f ./stage2/step-100000/params/_sharding +5b2c403538eea21cabdd0456dcce9684c2e0323104a5d3e69181d9b59cb57114 ./stage2/step-100000/params/array_metadatas/process_0 +e80e4713e4be13d069d87c35df5a3ce55913e9a59e363810a2f135829e4b29d2 ./stage2/step-100000/params/d/cbf2fef1100ba6c460caea5d0023e183 +9e6440f2971369604ff6089d15b5bd706be91dc016a20bb834d5ed99615c9b54 ./stage2/step-100000/params/manifest.ocdbt +d9e1b37586d2fcc49d688a62bbbea569be2891978a8e874277f4e2ee5364fdb3 ./stage2/step-100000/params/ocdbt.process_0/d/02ce27709e083f4ed8844fe4ab47131a +e2785efb474684260ba2731c22f74f0b19dd21c0b1ca25839ab7a168560228cd ./stage2/step-100000/params/ocdbt.process_0/d/054944d2cae72ec6dcde8844c294284c +3e2a41b88e4ffa38f3d5365512f411998579c0c83e4d3289536d56c02354f29a ./stage2/step-100000/params/ocdbt.process_0/d/1407cedaaf9ac43d641e631a808f0451 +65cb9182ab3e67386b7764d421a58d0a914ac86e242fafccce0ca65a7f6c0802 ./stage2/step-100000/params/ocdbt.process_0/d/3a109d42459c2150f8b988ce2442e2b5 +6461b15269d8d88d6b2c5c491ea53b52fcca71aa1857b5be9a5049574ccbc121 ./stage2/step-100000/params/ocdbt.process_0/d/7dc0a88b6d199ccbb016750c98012cf1 +4c1b1551057185e883abb2aa55c9acc068704464f1f7c4eceff23516b3acba4a ./stage2/step-100000/params/ocdbt.process_0/d/88cd20aa3cd4faa14d8123fbe93a44b9 +bab2c4f38a6328f378704e03c47ba826b3913fe69dfde7abe50e9ef6573af0b8 ./stage2/step-100000/params/ocdbt.process_0/d/910ee254009b8baf9f584972f0ee07f3 +6874c71af0dd426161b75b61198d1070aa80385eafb7e9494a7f4a4890727eff ./stage2/step-100000/params/ocdbt.process_0/d/a7b5b4d199bf5bf9a3712c2ecb35ce3c +ea9e3cf82bb7d52258ab6a694a7071b17539db9bb76498d43655096aa708915c ./stage2/step-100000/params/ocdbt.process_0/d/b7327531c46c5cc244cc7981e9f4da90 +12d6fece5e98bcf3c423319e7202d61e1085758fe312458473a3d9b37dbf7318 ./stage2/step-100000/params/ocdbt.process_0/d/dc503237bbcd88e2662f0cbd88978e03 +82b4331fec1fa27f3b9f675d2b7e3c7624afc8da83e7a66d4b8cd94384e43dfa ./stage2/step-100000/params/ocdbt.process_0/d/fa71f45e6411d633e6827c04a989a92f +81be2ef1c30c7599f1efbbe6b850b305480d4824945129adca337cdeabf25b71 ./stage2/step-100000/params/ocdbt.process_0/d/fa7cea5bceef1307d13d9a61ea91a35e +65bc792db617ab7945f5a7f5fd9f4fe207cf0f1185b693e91dc426820a564a29 ./stage2/step-100000/params/ocdbt.process_0/d/fb91f39ce5a047cdcaebc0bad36caad6 +034c510b6b4a7b7a01bff53cd076593b5f8e993f6497015f311517b1feced6d7 ./stage2/step-100000/params/ocdbt.process_0/d/fd4807cac3a561adf47064faff5b21c5 +2e21288d10e4e43b5c6228cca04c1bbcda6740006ba577f1bd60d7e1b4e80e8b ./stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt diff --git a/simulation/rmbench/recurrent/config/history_config.yaml b/simulation/rmbench/recurrent/config/history_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8ee16383754533b1112997733a85ef5f6e2fb234 --- /dev/null +++ b/simulation/rmbench/recurrent/config/history_config.yaml @@ -0,0 +1,82 @@ +budget: 256 +num_views: 3 +streaming_obs_horizon: 16 +integration_type: modulation +memory_token_dim: 1024 +representation_type: cosmo_bottleneck +memory_feature: + img: + net: identity + input_dim: 768 + pos: + input_dim: 768 + hidden_dim: 768 + state: + input_dim: 16 + hidden_dim: 512 +cosmo_bottleneck_memory: + memory_encoder_type: gated_delta_hybrid + memory_architecture_version: gated_delta_hybrid_v1 + memory_compute_dtype: bfloat16 + visual_feature_stream: dinov3 + feature_store_path: data/rmbench-450/stage1 + latent_source: offline + n_variants: 8 + num_query_samples: 16 + variant: Wan2.2-TI2V-5B + temporal_factor: 4 + downsample: 2 + num_latents: 256 + num_heads: 8 + cosmos_latent_dim: 48 + siglip_dim: 768 + proprio_in_dim: 16 + proprio_encoding: raw + action_in_dim: 16 + action_encoding: raw + lang_emb_dim: 2048 + action_tokenizer_hidden: 512 + proprio_tokenizer_hidden: 512 + pool_h: 2 + pool_w: 2 + wan_h: 256 + wan_w: 256 + latent_cap: 4096 + max_recent_tokens: 768 + max_step_tokens: 1244 + delta_num_layers: 2 + delta_num_heads: 8 + delta_head_dim: 128 + delta_conv_kernel: 4 + delta_ffn_inner: 2048 + delta_decay_bias: -4.0 + delta_beta_bias: -1.0 + delta_scan_checkpoint_size: 20 + delta_full_scan_mode: chunkwise_wy + delta_wy_matmul_dtype: bfloat16 + delta_recent_fusion_blocks: 1 + delta_max_chunks: 384 + delta_event_order: action_proprio_then_finalized_visual_v1 +cosmo_decoder: + num_heads: 8 + num_blocks: 2 + align_dim: 768 + align_num_patches: 192 + ffn_inner: 8192 + action_dim: 16 + proprio_dim: 16 + max_query_steps: 2048 +frame_jitter: + enabled: false + start_jitter: 3 + stride_choices: + - 1 + - 2 + - 3 + interval_jitter: 1 + canonical_downsample: 2 +stream_aug: + enabled: false + gauss_sigma: 0.05 + stream_drop_p: 0.1 + feat_drop_p: 0.0 diff --git a/simulation/rmbench/recurrent/config/input_provenance.json b/simulation/rmbench/recurrent/config/input_provenance.json new file mode 100644 index 0000000000000000000000000000000000000000..6aa230029751a684ff7de30192d6e5b47704453b --- /dev/null +++ b/simulation/rmbench/recurrent/config/input_provenance.json @@ -0,0 +1,27 @@ +{ + "action_causality": "strictly_causal", + "architecture": "gated_delta_hybrid_v1", + "base_policy": "pi0.5", + "checkpoint_metadata_sha256": "dd12980babae314a5097df31c47da055e3ced591c4fc44e1a394442365cb4323", + "checkpoint_path": "stage2/step-100000", + "checkpoint_step": 100000, + "dataset": "yilin-wu/imprint-rmbench-450", + "eval_consolidate_all": 1, + "execution_horizon": 50, + "freeze_variant": "image_only", + "latent_cap": 4096, + "max_recent_tokens": 768, + "max_step_tokens": 1244, + "memory_inference_mode": "incremental_recurrent", + "memory_language_source": "live", + "memory_training_scan": "chunkwise_wy", + "num_latents": 256, + "prediction_horizon": 50, + "stage1_checkpoint": "stage1/stage1_step160000.npz", + "stage1_sha256": "76f7ecea50d65c374c79bba071074fc35aded1dea22846d86cfe90b772f4bea3", + "stage1_wandb_run": "liqhuiha/Robocasa/9bm5uvc7", + "stage2_auxiliary_loss_weight": 0.0, + "stage2_memory_encoder": "finetuned", + "stage2_wandb_run": "liqhuiha/Robocasa/0f4k18rw", + "visual_feature_stream": "dinov3" +} diff --git a/simulation/rmbench/recurrent/config/norm_stats.json b/simulation/rmbench/recurrent/config/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..631b84bab53fc384b33fa5468192db069a535b6b --- /dev/null +++ b/simulation/rmbench/recurrent/config/norm_stats.json @@ -0,0 +1,152 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + -0.2884466350078583, + 1.116854190826416, + 0.8176121115684509, + -0.6286301016807556, + 0.0027539930306375027, + 0.3858006000518799, + 0.0, + 0.29012763500213623, + 1.1434277296066284, + 0.834911048412323, + -0.6164064407348633, + 0.0033871293999254704, + 0.1835649162530899, + 0.0, + 0.44485998153686523, + 0.693847119808197 + ], + "std": [ + 0.2924744486808777, + 0.9123055338859558, + 0.6921584606170654, + 0.5167360901832581, + 0.024822130799293518, + 0.7730752229690552, + 0.0, + 0.3022353947162628, + 0.9363216161727905, + 0.7197534441947937, + 0.5159713625907898, + 0.042254358530044556, + 0.4932859241962433, + 0.0, + 0.4852803945541382, + 0.43509986996650696 + ], + "q01": [ + -1.1566282510757446, + -0.00001813583730836399, + -0.000035173223295714706, + -1.493809461593628, + -0.044773634523153305, + -2.2744300365448, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.4539425373077393, + -0.11438962072134018, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410673141479492, + 2.2643818855285645, + 0.000017526803276268765, + 0.09094329178333282, + 1.4207451343536377, + 0.0, + 0.9992868304252625, + 2.5619561672210693, + 2.3963286876678467, + 0.000027463351216283627, + 0.12538659572601318, + 1.50703763961792, + 0.0, + 1.0, + 1.0 + ] + }, + "state": { + "mean": [ + -0.28619107604026794, + 1.1170974969863892, + 0.8187518119812012, + -0.6292704939842224, + 0.002742639509961009, + 0.38321176171302795, + 0.0, + 0.287190318107605, + 1.1249585151672363, + 0.8228866457939148, + -0.6072733998298645, + 0.003468270180746913, + 0.18374495208263397, + 0.0, + 0.47009503841400146, + 0.700372040271759 + ], + "std": [ + 0.28622451424598694, + 0.9085981845855713, + 0.6900886297225952, + 0.5145737528800964, + 0.024824900552630424, + 0.7737277746200562, + 0.0, + 0.3019956946372986, + 0.9365604519844055, + 0.720547616481781, + 0.5160077214241028, + 0.04248249530792236, + 0.4887017011642456, + 0.0, + 0.4873206615447998, + 0.4321308135986328 + ], + "q01": [ + -0.991405725479126, + -0.00001813583730836399, + -0.00003476843994576484, + -1.4938021898269653, + -0.04476987197995186, + -2.274427652359009, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.455243706703186, + -0.11395290493965149, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410663366317749, + 2.2642898559570312, + 0.000016971260265563615, + 0.09090737998485565, + 1.4203816652297974, + 0.0, + 1.0001307725906372, + 2.56237530708313, + 2.397764205932617, + 0.000027463351216283627, + 0.12876388430595398, + 1.4705802202224731, + 0.0, + 1.0, + 1.0 + ] + } + } +} \ No newline at end of file diff --git a/simulation/rmbench/recurrent/stage1/stage1_step160000.npz b/simulation/rmbench/recurrent/stage1/stage1_step160000.npz new file mode 100644 index 0000000000000000000000000000000000000000..52394d3d836a982b0610d1af66ecb79a0b6f627f --- /dev/null +++ b/simulation/rmbench/recurrent/stage1/stage1_step160000.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:76f7ecea50d65c374c79bba071074fc35aded1dea22846d86cfe90b772f4bea3 +size 391036002 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/_CHECKPOINT_METADATA b/simulation/rmbench/recurrent/stage2/step-100000/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..83f3f94411032806cbc47174615d27218d2ff7ec --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1788751407534277577, "commit_timestamp_nsecs": 1788751421836431897, "custom_metadata": {}} \ No newline at end of file diff --git a/simulation/rmbench/recurrent/stage2/step-100000/assets/robomme/norm_stats.json b/simulation/rmbench/recurrent/stage2/step-100000/assets/robomme/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..631b84bab53fc384b33fa5468192db069a535b6b --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/assets/robomme/norm_stats.json @@ -0,0 +1,152 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + -0.2884466350078583, + 1.116854190826416, + 0.8176121115684509, + -0.6286301016807556, + 0.0027539930306375027, + 0.3858006000518799, + 0.0, + 0.29012763500213623, + 1.1434277296066284, + 0.834911048412323, + -0.6164064407348633, + 0.0033871293999254704, + 0.1835649162530899, + 0.0, + 0.44485998153686523, + 0.693847119808197 + ], + "std": [ + 0.2924744486808777, + 0.9123055338859558, + 0.6921584606170654, + 0.5167360901832581, + 0.024822130799293518, + 0.7730752229690552, + 0.0, + 0.3022353947162628, + 0.9363216161727905, + 0.7197534441947937, + 0.5159713625907898, + 0.042254358530044556, + 0.4932859241962433, + 0.0, + 0.4852803945541382, + 0.43509986996650696 + ], + "q01": [ + -1.1566282510757446, + -0.00001813583730836399, + -0.000035173223295714706, + -1.493809461593628, + -0.044773634523153305, + -2.2744300365448, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.4539425373077393, + -0.11438962072134018, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410673141479492, + 2.2643818855285645, + 0.000017526803276268765, + 0.09094329178333282, + 1.4207451343536377, + 0.0, + 0.9992868304252625, + 2.5619561672210693, + 2.3963286876678467, + 0.000027463351216283627, + 0.12538659572601318, + 1.50703763961792, + 0.0, + 1.0, + 1.0 + ] + }, + "state": { + "mean": [ + -0.28619107604026794, + 1.1170974969863892, + 0.8187518119812012, + -0.6292704939842224, + 0.002742639509961009, + 0.38321176171302795, + 0.0, + 0.287190318107605, + 1.1249585151672363, + 0.8228866457939148, + -0.6072733998298645, + 0.003468270180746913, + 0.18374495208263397, + 0.0, + 0.47009503841400146, + 0.700372040271759 + ], + "std": [ + 0.28622451424598694, + 0.9085981845855713, + 0.6900886297225952, + 0.5145737528800964, + 0.024824900552630424, + 0.7737277746200562, + 0.0, + 0.3019956946372986, + 0.9365604519844055, + 0.720547616481781, + 0.5160077214241028, + 0.04248249530792236, + 0.4887017011642456, + 0.0, + 0.4873206615447998, + 0.4321308135986328 + ], + "q01": [ + -0.991405725479126, + -0.00001813583730836399, + -0.00003476843994576484, + -1.4938021898269653, + -0.04476987197995186, + -2.274427652359009, + 0.0, + -0.000023844564566388726, + -0.000019354150936123915, + -0.0000323197637044359, + -1.455243706703186, + -0.11395290493965149, + -1.3077747821807861, + 0.0, + 0.0, + 0.0 + ], + "q99": [ + 0.000020988931282772683, + 2.410663366317749, + 2.2642898559570312, + 0.000016971260265563615, + 0.09090737998485565, + 1.4203816652297974, + 0.0, + 1.0001307725906372, + 2.56237530708313, + 2.397764205932617, + 0.000027463351216283627, + 0.12876388430595398, + 1.4705802202224731, + 0.0, + 1.0, + 1.0 + ] + } + } +} \ No newline at end of file diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/_METADATA b/simulation/rmbench/recurrent/stage2/step-100000/params/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..840952cd0fbce3705e7f845c3c25426533fe9a31 --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_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": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1152, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 16, 72, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [14, 14, 3, 288]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "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": [1152, 512]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64288, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_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": [512]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'kv_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "kv_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'mem_rms_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "mem_rms_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": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'out_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "out_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'q_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "q_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 2048, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4096, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_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": [256, 32]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'beta_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "beta_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'beta_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "beta_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": [256, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'decay_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "decay_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'decay_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "decay_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": [256, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 512]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'norm_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'norm_in', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_gate', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_gate', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'beta_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "beta_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'beta_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "beta_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": [256, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'decay_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "decay_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'decay_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "decay_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": [256, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 512]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'norm_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'norm_in', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_gate', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_gate', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1024, 256]}}, "('params', 'mem_encoder', 'event_slot', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "event_slot", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('params', 'mem_encoder', 'event_type', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "event_type", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('params', 'mem_encoder', 'lang_film', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'lang_film', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('params', 'mem_encoder', 'latents', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "latents", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1024]}}, "('params', 'mem_encoder', 'proj_cosmo', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_cosmo', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "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, 1024]}}, "('params', 'mem_encoder', 'proj_recent', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_recent', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 1024]}}, "('params', 'mem_encoder', 'read_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'read_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_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": [256]}}, "('params', 'mem_encoder', 'read_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'read_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'read_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'read_q', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_q", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_kv_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_kv_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_kv_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_kv_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": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_q_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_q_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_q_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_q_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": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 512]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_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": [256]}}, "('params', 'mem_encoder', 'recent_null', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_null", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 256]}}, "('params', 'mem_encoder', 'time_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "time_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'time_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "time_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 256]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/_sharding b/simulation/rmbench/recurrent/stage2/step-100000/params/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..4a6e96dbc26510168ad386a01cd30f863dfac7c5 --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/_sharding @@ -0,0 +1 @@ +{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm1lbV9ybXNfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm91dF9laW5zdW1fbWVtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLmt2X2VpbnN1bV9tZW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLnFfZWluc3VtX21lbS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm5vcm1faW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm5vcm1faW4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9nYXRlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9nYXRlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmJldGFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmJldGFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmNvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmNvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmRlY2F5X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmRlY2F5X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmtfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnZfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm5vcm1faW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm5vcm1faW4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9nYXRlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9nYXRlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmJldGFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmJldGFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmNvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmNvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmRlY2F5X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmRlY2F5X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmtfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnZfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmV2ZW50X3Nsb3QudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmV2ZW50X3R5cGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhdGVudHMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19rdl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19rdl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19xX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19xX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5faW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5faW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9udWxsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfcS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRpbWVfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRpbWVfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}"} \ No newline at end of file diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/array_metadatas/process_0 b/simulation/rmbench/recurrent/stage2/step-100000/params/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..22bc4ef5038e08f3de4e6cc5cf99f1237dd22b1e --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 1152, 1076], "chunk_shape": [27, 1152, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 16, 72, 288], "chunk_shape": [27, 16, 72, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [1152, 512], "chunk_shape": [1152, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [1024, 768], "chunk_shape": [1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 8, 256, 512], "chunk_shape": [18, 8, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 8, 512, 256], "chunk_shape": [18, 8, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.kv_einsum_mem.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.mem_rms_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.out_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.q_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.bias.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.kernel.value", "write_shape": [18, 1024, 512], "chunk_shape": [18, 1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 2048, 4096], "chunk_shape": [18, 2, 2048, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 1024, 1024], "chunk_shape": [18, 2, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.beta_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.beta_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.conv.kernel.value", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.decay_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.decay_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_in.kernel.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.k_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.norm_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.norm_in.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_gate.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_gate.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.q_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.v_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.beta_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.beta_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.conv.kernel.value", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.decay_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.decay_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_in.kernel.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.k_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.norm_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.norm_in.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_gate.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_gate.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.q_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.v_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.event_slot.value", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.event_type.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.latents.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.kernel.value", "write_shape": [12, 1024], "chunk_shape": [12, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.kernel.value", "write_shape": [192, 1024], "chunk_shape": [192, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_q.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_in.kernel.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_null.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.time_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.time_proj.kernel.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}]} \ No newline at end of file diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/d/cbf2fef1100ba6c460caea5d0023e183 b/simulation/rmbench/recurrent/stage2/step-100000/params/d/cbf2fef1100ba6c460caea5d0023e183 new file mode 100644 index 0000000000000000000000000000000000000000..ff2f7734cdc816ec0945ee3e6632a6baeaaa6b7c --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/d/cbf2fef1100ba6c460caea5d0023e183 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e80e4713e4be13d069d87c35df5a3ce55913e9a59e363810a2f135829e4b29d2 +size 190592 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/manifest.ocdbt b/simulation/rmbench/recurrent/stage2/step-100000/params/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..ebf7584b81f28f88b138596bf32fef9d982989f5 Binary files /dev/null and b/simulation/rmbench/recurrent/stage2/step-100000/params/manifest.ocdbt differ diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/02ce27709e083f4ed8844fe4ab47131a b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/02ce27709e083f4ed8844fe4ab47131a new file mode 100644 index 0000000000000000000000000000000000000000..305b42819ecdffa7ad69f717cb95e83c0393b3d0 Binary files /dev/null and b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/02ce27709e083f4ed8844fe4ab47131a differ diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/054944d2cae72ec6dcde8844c294284c b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/054944d2cae72ec6dcde8844c294284c new file mode 100644 index 0000000000000000000000000000000000000000..0066f0072a044cc37eb10a6b0abe09f159d58d61 --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/054944d2cae72ec6dcde8844c294284c @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e2785efb474684260ba2731c22f74f0b19dd21c0b1ca25839ab7a168560228cd +size 2236773228 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/1407cedaaf9ac43d641e631a808f0451 b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/1407cedaaf9ac43d641e631a808f0451 new file mode 100644 index 0000000000000000000000000000000000000000..3b88a45ead114d9476fc3b5fd74fe41291f2a5ed Binary files /dev/null and b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/1407cedaaf9ac43d641e631a808f0451 differ diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/3a109d42459c2150f8b988ce2442e2b5 b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/3a109d42459c2150f8b988ce2442e2b5 new file mode 100644 index 0000000000000000000000000000000000000000..15b96ed9dbb54ec38ed078ca5ec0cc7e6a39f819 --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/3a109d42459c2150f8b988ce2442e2b5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:65cb9182ab3e67386b7764d421a58d0a914ac86e242fafccce0ca65a7f6c0802 +size 769587458 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/7dc0a88b6d199ccbb016750c98012cf1 b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/7dc0a88b6d199ccbb016750c98012cf1 new file mode 100644 index 0000000000000000000000000000000000000000..e152b16728c3d9dc16ab85841c495a06d98a2043 --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/7dc0a88b6d199ccbb016750c98012cf1 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6461b15269d8d88d6b2c5c491ea53b52fcca71aa1857b5be9a5049574ccbc121 +size 24030053 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/88cd20aa3cd4faa14d8123fbe93a44b9 b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/88cd20aa3cd4faa14d8123fbe93a44b9 new file mode 100644 index 0000000000000000000000000000000000000000..98e91132e8dd1f7c60cbe99a137fb593e1c2485c --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/88cd20aa3cd4faa14d8123fbe93a44b9 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4c1b1551057185e883abb2aa55c9acc068704464f1f7c4eceff23516b3acba4a +size 2240204089 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/910ee254009b8baf9f584972f0ee07f3 b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/910ee254009b8baf9f584972f0ee07f3 new file mode 100644 index 0000000000000000000000000000000000000000..3753ad9a89861eb1ec7d507312d6db81bf61cfb9 Binary files /dev/null and b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/910ee254009b8baf9f584972f0ee07f3 differ diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/a7b5b4d199bf5bf9a3712c2ecb35ce3c b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/a7b5b4d199bf5bf9a3712c2ecb35ce3c new file mode 100644 index 0000000000000000000000000000000000000000..0c412b61105d1e53d2cfda5c582ccf02b48f1b61 Binary files /dev/null and b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/a7b5b4d199bf5bf9a3712c2ecb35ce3c differ diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/b7327531c46c5cc244cc7981e9f4da90 b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/b7327531c46c5cc244cc7981e9f4da90 new file mode 100644 index 0000000000000000000000000000000000000000..36caa3980d572ff171cf9d66f5f6b36980e311f5 --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/b7327531c46c5cc244cc7981e9f4da90 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ea9e3cf82bb7d52258ab6a694a7071b17539db9bb76498d43655096aa708915c +size 2240473834 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/dc503237bbcd88e2662f0cbd88978e03 b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/dc503237bbcd88e2662f0cbd88978e03 new file mode 100644 index 0000000000000000000000000000000000000000..2b898bd1367e67d2bdb6a43500a64dc0593715cd Binary files /dev/null and b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/dc503237bbcd88e2662f0cbd88978e03 differ diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa71f45e6411d633e6827c04a989a92f b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa71f45e6411d633e6827c04a989a92f new file mode 100644 index 0000000000000000000000000000000000000000..b7b514338fbed2d6ab8975cc94235c47a240df95 --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa71f45e6411d633e6827c04a989a92f @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:82b4331fec1fa27f3b9f675d2b7e3c7624afc8da83e7a66d4b8cd94384e43dfa +size 2263394319 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa7cea5bceef1307d13d9a61ea91a35e b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa7cea5bceef1307d13d9a61ea91a35e new file mode 100644 index 0000000000000000000000000000000000000000..da06644f5bf492141aff87db5136b10bfc8300cc --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa7cea5bceef1307d13d9a61ea91a35e @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:81be2ef1c30c7599f1efbbe6b850b305480d4824945129adca337cdeabf25b71 +size 190584 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fb91f39ce5a047cdcaebc0bad36caad6 b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fb91f39ce5a047cdcaebc0bad36caad6 new file mode 100644 index 0000000000000000000000000000000000000000..16c0d0ab5945c13d999984098fa9cd80437a744b --- /dev/null +++ b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fb91f39ce5a047cdcaebc0bad36caad6 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:65bc792db617ab7945f5a7f5fd9f4fe207cf0f1185b693e91dc426820a564a29 +size 2240206785 diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fd4807cac3a561adf47064faff5b21c5 b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fd4807cac3a561adf47064faff5b21c5 new file mode 100644 index 0000000000000000000000000000000000000000..2a0b58f2c74a2f8b9539cde943dce7b6fad34243 Binary files /dev/null and b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fd4807cac3a561adf47064faff5b21c5 differ diff --git a/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..f9af5277b14194f373779654935494f2755004f1 Binary files /dev/null and b/simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt differ diff --git a/simulation/robocasa/full_attention/.gitattributes b/simulation/robocasa/full_attention/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..44cd336bab8c87ebabca1e27f11e7ecf9a60a23c --- /dev/null +++ b/simulation/robocasa/full_attention/.gitattributes @@ -0,0 +1,45 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz 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 +stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1 filter=lfs diff=lfs merge=lfs -text diff --git a/simulation/robocasa/full_attention/LICENSE_GEMMA.txt b/simulation/robocasa/full_attention/LICENSE_GEMMA.txt new file mode 100644 index 0000000000000000000000000000000000000000..e0f8930d83fe9eef924c0eed50ee0e572eb471fa --- /dev/null +++ b/simulation/robocasa/full_attention/LICENSE_GEMMA.txt @@ -0,0 +1,77 @@ +Gemma Terms of Use + +Last modified: February 21, 2024 + +By using, reproducing, modifying, distributing, performing or displaying any portion or element of Gemma, Model Derivatives including via any Hosted Service, (each as defined below) (collectively, the "Gemma Services") or otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement. + +Section 1: DEFINITIONS +1.1 Definitions +(a) "Agreement" or "Gemma Terms of Use" means these terms and conditions that govern the use, reproduction, Distribution or modification of the Gemma Services and any terms and conditions incorporated by reference. + +(b) "Distribution" or "Distribute" means any transmission, publication, or other sharing of Gemma or Model Derivatives to a third party, including by providing or making Gemma or its functionality available as a hosted service via API, web access, or any other electronic or remote means ("Hosted Service"). + +(c) "Gemma" means the set of machine learning language models, trained model weights and parameters identified at ai.google.dev/gemma, regardless of the source that you obtained it from. + +(d) "Google" means Google LLC. + +(e) "Model Derivatives" means all (i) modifications to Gemma, (ii) works based on Gemma, or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Gemma, to that model in order to cause that model to perform similarly to Gemma, including distillation methods that use intermediate data representations or methods based on the generation of synthetic data Outputs by Gemma for training that model. For clarity, Outputs are not deemed Model Derivatives. + +(f) "Output" means the information content output of Gemma or a Model Derivative that results from operating or otherwise using Gemma or the Model Derivative, including via a Hosted Service. + +1.2 +As used in this Agreement, "including" means "including without limitation". + +Section 2: ELIGIBILITY AND USAGE +2.1 Eligibility +You represent and warrant that you have the legal capacity to enter into this Agreement (including being of sufficient age of consent). If you are accessing or using any of the Gemma Services for or on behalf of a legal entity, (a) you are entering into this Agreement on behalf of yourself and that legal entity, (b) you represent and warrant that you have the authority to act on behalf of and bind that entity to this Agreement and (c) references to "you" or "your" in the remainder of this Agreement refers to both you (as an individual) and that entity. + +2.2 Use +You may use, reproduce, modify, Distribute, perform or display any of the Gemma Services only in accordance with the terms of this Agreement, and must not violate (or encourage or permit anyone else to violate) any term of this Agreement. + +Section 3: DISTRIBUTION AND RESTRICTIONS +3.1 Distribution and Redistribution +You may reproduce or Distribute copies of Gemma or Model Derivatives if you meet all of the following conditions: + +You must include the use restrictions referenced in Section 3.2 as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Gemma or Model Derivatives and you must provide notice to subsequent users you Distribute to that Gemma or Model Derivatives are subject to the use restrictions in Section 3.2. +You must provide all third party recipients of Gemma or Model Derivatives a copy of this Agreement. +You must cause any modified files to carry prominent notices stating that you modified the files. +All Distributions (other than through a Hosted Service) must be accompanied by a "Notice" text file that contains the following notice: "Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms". +You may add your own intellectual property statement to your modifications and, except as set forth in this Section, may provide additional or different terms and conditions for use, reproduction, or Distribution of your modifications, or for any such Model Derivatives as a whole, provided your use, reproduction, modification, Distribution, performance, and display of Gemma otherwise complies with the terms and conditions of this Agreement. Any additional or different terms and conditions you impose must not conflict with the terms of this Agreement. + +3.2 Use Restrictions +You must not use any of the Gemma Services: + +for the restricted uses set forth in the Gemma Prohibited Use Policy at ai.google.dev/gemma/prohibited_use_policy ("Prohibited Use Policy"), which is hereby incorporated by reference into this Agreement; or +in violation of applicable laws and regulations. +To the maximum extent permitted by law, Google reserves the right to restrict (remotely or otherwise) usage of any of the Gemma Services that Google reasonably believes are in violation of this Agreement. + +3.3 Generated Output +Google claims no rights in Outputs you generate using Gemma. You and your users are solely responsible for Outputs and their subsequent uses. + +Section 4: ADDITIONAL PROVISIONS +4.1 Updates +Google may update Gemma from time to time, and you must make reasonable efforts to use the latest version of Gemma. + +4.2 Trademarks +Nothing in this Agreement grants you any rights to use Google's trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between you and Google. Google reserves any rights not expressly granted herein. + +4.3 DISCLAIMER OF WARRANTY +UNLESS REQUIRED BY APPLICABLE LAW, THE GEMMA SERVICES, AND OUTPUTS, ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY WARRANTIES OR CONDITIONS OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR OR DISTRIBUTING ANY OF THE GEMMA SERVICES OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR USE OR DISTRIBUTION OF ANY OF THE GEMMA SERVICES OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT. + +4.4 LIMITATION OF LIABILITY +TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY, CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW, SHALL GOOGLE OR ITS AFFILIATES BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL, OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO, ANY OF THE GEMMA SERVICES OR OUTPUTS EVEN IF GOOGLE OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. + +4.5 Term, Termination, and Survival +The term of this Agreement will commence upon your acceptance of this Agreement (including acceptance by your use, modification, or Distribution, reproduction, performance or display of any portion or element of the Gemma Services) and will continue in full force and effect until terminated in accordance with the terms of this Agreement. Google may terminate this Agreement if you are in breach of any term of this Agreement. Upon termination of this Agreement, you must delete and cease use and Distribution of all copies of Gemma and Model Derivatives in your possession or control. Sections 1, 2.1, 3.3, 4.2 to 4.9 shall survive the termination of this Agreement. + +4.6 Governing Law and Jurisdiction +This Agreement will be governed by the laws of the State of California without regard to choice of law principles. The UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The state and federal courts of Santa Clara County, California shall have exclusive jurisdiction of any dispute arising out of this Agreement. + +4.7 Severability +If any provision of this Agreement is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein. + +4.8 Entire Agreement +This Agreement states all the terms agreed between the parties and supersedes all other agreements between the parties as of the date of acceptance relating to its subject matter. + +4.9 No Waiver +Google will not be treated as having waived any rights by not exercising (or delaying the exercise of) any rights under this Agreement. \ No newline at end of file diff --git a/simulation/robocasa/full_attention/NOTICE b/simulation/robocasa/full_attention/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..1b9cb6541d7135c4c78a542ac4a3e04a971c96db --- /dev/null +++ b/simulation/robocasa/full_attention/NOTICE @@ -0,0 +1,11 @@ +IMPRINT +Copyright 2026 + +This distribution contains modifications to the OpenPI project: +https://github.com/Physical-Intelligence/openpi + +OpenPI is distributed under the Apache License 2.0. Gemma components and +derivatives are subject to the Gemma Terms of Use in LICENSE_GEMMA.txt. + +Gemma is provided under and subject to the Gemma Terms of Use found at +ai.google.dev/gemma/terms. diff --git a/simulation/robocasa/full_attention/README.md b/simulation/robocasa/full_attention/README.md new file mode 100644 index 0000000000000000000000000000000000000000..181e6ae94807372a7819474c142325ef95b6f4a1 --- /dev/null +++ b/simulation/robocasa/full_attention/README.md @@ -0,0 +1,31 @@ +--- +license: gemma +library_name: jax +tags: + - robotics + - vision-language-action + - memory + - robocasa + - dinov3 +--- + +# IMPRINT RoboCasa Full-Attention Checkpoint + +Canonical Pi0.5 IMPRINT checkpoint for the customized 12-task RoboCasa +benchmark using the full-attention memory encoder. + +- Stage 1: trained 180,000 updates, selected at 160,000, W&B `far-wandb/Robocasa/yw0l1a20` +- Stage 2: 150,000 updates, W&B `far-wandb/Robocasa/c0t2f1cg` +- memory tokens: 256 x 1024 +- policy horizon: 20 +- Stage 2: memory encoder finetuned; image tower frozen + +Use `yilin-wu/imprint-robocasa-600` for training and +`yilin-wu/imprint-robocasa-environment` for the exact patched RoboCasa 1.0.1 +simulator source. This release contains the complete Stage-2 Orbax parameter +checkpoint, including the PaliGemma/Gemma backbone. It is a Gemma model +derivative and is distributed subject to the included Gemma Terms of Use. +Large simulator assets are not duplicated. + +The matching training store, patched simulator source, and reproducibility +code are publicly available under the IMPRINT release. diff --git a/simulation/robocasa/full_attention/SHA256SUMS b/simulation/robocasa/full_attention/SHA256SUMS new file mode 100644 index 0000000000000000000000000000000000000000..7fdc1d2775ffaf64e6500a0fe51d7e43e241cca6 --- /dev/null +++ b/simulation/robocasa/full_attention/SHA256SUMS @@ -0,0 +1,30 @@ +3e2c24001f9ef57bf7ec959a3658fbb49cdad113cdf394c264da9d16f9bdd132 ./LICENSE_GEMMA.txt +5b21d7116ac42a0ab6141488cb654c2411557ac977e2276accb88ef774f28c19 ./NOTICE +688f34b0256eb63dff227c29811b7e1b916e503a2eb56888e5b52430707ee2bc ./README.md +6ed9405ca99c33d57f8101c1c98130e74a33602dadbc2aeb1e0cac72e03a9788 ./config/history_config.yaml +a7dda1ef87035e50a39390c549d7e4de8fb6359d3f972f28f140f02dc2029138 ./config/input_provenance.json +39a8f9a475d18c9167a3872c7679a348b81a165776fe3e16c391f354c7ad8514 ./config/norm_stats.json +34387cbbd3891cc7bb082516dfb4993ee9d2319e32344a364ea32bd611cd924e ./stage1/stage1_step160000.npz +d962b86d3894c8afdbadc56c986b8c6f9d96a453989e2fe77bc9794ef4c125e6 ./stage2/step-150000/_CHECKPOINT_METADATA +39a8f9a475d18c9167a3872c7679a348b81a165776fe3e16c391f354c7ad8514 ./stage2/step-150000/assets/robomme/norm_stats.json +168e29fbef9d082040ff482e5d1015975d2bcde4a72f6e7bb226c3706443c28d ./stage2/step-150000/params/_METADATA +19d4e4b2d799c0b0300b17f32821411136fa0ff4c638b616bae2250d42c82056 ./stage2/step-150000/params/_sharding +8e51278a946b172415780a940115a3dc79ce6f711541150b9ff8639ec4ef008d ./stage2/step-150000/params/array_metadatas/process_0 +476ff74dc38c756c9c195faec2152f34b052f4263d9e19899bc40360de746aae ./stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9 +29b0ae4aa14607930fc553bd6de727fa2065be5f83da82a65f46b9280673cb2b ./stage2/step-150000/params/manifest.ocdbt +034c510b6b4a7b7a01bff53cd076593b5f8e993f6497015f311517b1feced6d7 ./stage2/step-150000/params/ocdbt.process_0/d/06aabeab1a0e602ca8f8e217fcb96709 +824eaea25c59bac7143776b1c37425d47a9426986d0246d389cc6a5eca912ed8 ./stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9 +2290868f5d3f5fb15c3e9342cc3fa0adeb4e0517ea2ba667c00e1f9555b32be0 ./stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60 +7088b323898c94fe320df837059a255f9ac31db7a06b98cb350a373b4e7e855f ./stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8 +3855446dea4243997fe5ad2e90d0e7a9f76d000c414f3e6998606716384078db ./stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5 +ed6ad005edb24d33eb04226b8bee988a0ef5c44e3ef94ec84b56765001326be4 ./stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1 +401e509a3f9037aeb23fb337db4564689bf1f439c21e2c5adac1519ba21f42f1 ./stage2/step-150000/params/ocdbt.process_0/d/822b3aef48f3ec46a514c1a9b663d129 +d60b237dba5c9d4a74f511cb55156e8a59b605339a72cbfb12a6c3fab1b1a6f8 ./stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891 +d7c51b12565e3a360a071dd3385e6932c528325a779fce441268c7cccb030685 ./stage2/step-150000/params/ocdbt.process_0/d/874b4177fd57749dda8e5c8a08b34524 +789e58de6a76f486337c20368b11cda5b3ff58c8c55dfd753c646afb14c629f1 ./stage2/step-150000/params/ocdbt.process_0/d/96836e12295a196861bcac1091e408c6 +a8f1de421cb76f94b3d0d0c14a93d8c204759063b5585b3fce25d2798aa7619d ./stage2/step-150000/params/ocdbt.process_0/d/b1b9c620755ae82fa0958a1aba0a6853 +6a4a5b9583174ad28223727df900bf05edbd0f03b50fd85948750876e13fd66a ./stage2/step-150000/params/ocdbt.process_0/d/b726d8d65685cae9811fa5a7918320d3 +d58075ff02f957a5bfc1f9892c8956b19df5dfd13968cea4cb51530ded6b1137 ./stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065 +b42d860c3b83040a1bf3da8ae1e73497e29fa83d4ab6d7c7bf2ee1f9a9bac22c ./stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761 +7f3e1d8d11d6d0f2d58a47be5d1c47d533b9c07d892e41057f1acc6cb2be23a9 ./stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25 +2a4d84154371e69384a7713f2f1bdd5d136db4482bf341b347345d30d8878bc9 ./stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt diff --git a/simulation/robocasa/full_attention/config/history_config.yaml b/simulation/robocasa/full_attention/config/history_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..025c64a0ba07defb0f5024ac9f715287d723ec0a --- /dev/null +++ b/simulation/robocasa/full_attention/config/history_config.yaml @@ -0,0 +1,74 @@ +budget: 256 +num_views: 3 +streaming_obs_horizon: 16 +integration_type: modulation +memory_token_dim: 1024 +representation_type: cosmo_bottleneck +memory_feature: + img: + net: identity + input_dim: 768 + pos: + input_dim: 768 + hidden_dim: 768 + state: + input_dim: 16 + hidden_dim: 512 +cosmo_bottleneck_memory: + visual_feature_stream: dinov3 + latent_source: offline + episode_cache_size: 8 + pool_h: 2 + pool_w: 2 + n_variants: 8 + max_recent_tokens: 768 + max_step_tokens: 1776 + max_stream_steps: 1776 + latent_cap: 4096 + latent_buckets: + - 256 + - 512 + - 1024 + - 2048 + - 4096 + num_query_samples: 16 + variant: Wan2.2-TI2V-5B + temporal_factor: 4 + downsample: 2 + wan_h: 256 + wan_w: 256 + num_latents: 256 + num_cross_blocks: 2 + num_heads: 8 + cosmos_latent_dim: 48 + siglip_dim: 768 + proprio_in_dim: 16 + proprio_encoding: raw + action_in_dim: 12 + action_encoding: raw + lang_emb_dim: 2048 + action_tokenizer_hidden: 512 + proprio_tokenizer_hidden: 512 +cosmo_decoder: + num_heads: 8 + num_blocks: 2 + align_dim: 768 + align_num_patches: 192 + ffn_inner: 8192 + action_dim: 12 + proprio_dim: 16 + max_query_steps: 2048 +frame_jitter: + enabled: false + start_jitter: 3 + stride_choices: + - 1 + - 2 + - 3 + interval_jitter: 1 + canonical_downsample: 2 +stream_aug: + enabled: false + gauss_sigma: 0.05 + stream_drop_p: 0.1 + feat_drop_p: 0.0 diff --git a/simulation/robocasa/full_attention/config/input_provenance.json b/simulation/robocasa/full_attention/config/input_provenance.json new file mode 100644 index 0000000000000000000000000000000000000000..60554a6b5eb890f7134be09692b0b825f5f24ae3 --- /dev/null +++ b/simulation/robocasa/full_attention/config/input_provenance.json @@ -0,0 +1,25 @@ +{ + "action_causality": "strictly_causal", + "architecture": "perceiver_v1", + "base_policy": "pi0.5", + "checkpoint_metadata_sha256": "d962b86d3894c8afdbadc56c986b8c6f9d96a453989e2fe77bc9794ef4c125e6", + "checkpoint_path": "stage2/step-150000", + "checkpoint_step": 150000, + "dataset": "yilin-wu/imprint-robocasa-600", + "eval_consolidate_all": 1, + "execution_horizon": 20, + "freeze_variant": "image_only", + "latent_cap": 4096, + "max_recent_tokens": 768, + "max_step_tokens": 1776, + "memory_language_source": "live", + "num_latents": 256, + "prediction_horizon": 20, + "stage1_checkpoint": "stage1/stage1_step160000.npz", + "stage1_sha256": "34387cbbd3891cc7bb082516dfb4993ee9d2319e32344a364ea32bd611cd924e", + "stage1_wandb_run": "far-wandb/Robocasa/yw0l1a20", + "stage2_auxiliary_loss_weight": 0.0, + "stage2_memory_encoder": "finetuned", + "stage2_wandb_run": "far-wandb/Robocasa/c0t2f1cg", + "visual_feature_stream": "dinov3" +} diff --git a/simulation/robocasa/full_attention/config/norm_stats.json b/simulation/robocasa/full_attention/config/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..84350088fc65d4bab950ca32bb284a71339e6d46 --- /dev/null +++ b/simulation/robocasa/full_attention/config/norm_stats.json @@ -0,0 +1,136 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + 0.005957702174782753, + -0.009433974511921406, + -0.027537286281585693, + -0.0002329433336853981, + -0.0019665092695504427, + 0.0010746258776634932, + -0.17213615775108337, + 0.0022201647516340017, + -0.0015274320030584931, + 0.0009210301213897765, + -0.000019904651708202437, + -0.7345612645149231 + ], + "std": [ + 0.2822888493537903, + 0.31111952662467957, + 0.30976882576942444, + 0.060622796416282654, + 0.08188065886497498, + 0.08514650166034698, + 0.9420629143714905, + 0.06350194662809372, + 0.1455080211162567, + 0.05478019639849663, + 0.0006542060291394591, + 0.6832558512687683 + ], + "q01": [ + -1.0, + -1.0, + -0.6785714030265808, + -0.1971428543329239, + -0.27142858505249023, + -0.30571427941322327, + -1.0, + -0.07321428507566452, + -0.5982142686843872, + -0.24081368744373322, + 0.0, + -1.0 + ], + "q99": [ + 0.6482142806053162, + 1.0, + 1.0, + 0.1971428543329239, + 0.2514285743236542, + 0.3028571307659149, + 1.0, + 0.19107143580913544, + 0.612500011920929, + 0.24883341789245605, + 0.0, + 1.0 + ] + }, + "state": { + "mean": [ + 0.5936015248298645, + -0.20600375533103943, + 0.4474135637283325, + -0.30925947427749634, + -0.007771184667944908, + -0.04887167364358902, + 0.1312701404094696, + 2.4705684185028076, + -1.8421560525894165, + 0.6735693216323853, + -0.018410686403512955, + -0.002943282015621662, + 0.05312279239296913, + 0.595700740814209, + 0.03035394847393036, + -0.030480381101369858 + ], + "std": [ + 0.8089187145233154, + 0.6286742687225342, + 0.18121899664402008, + 0.8043732047080994, + 0.2596682906150818, + 0.349939227104187, + 0.21797646582126617, + 1.5330910682678223, + 1.3586890697479248, + 0.09442085027694702, + 0.2916492521762848, + 0.08581095188856125, + 0.6641682982444763, + 0.32726120948791504, + 0.011618087068200111, + 0.011510285548865795 + ], + "q01": [ + 0.042294517159461975, + -3.4426209926605225, + 0.0796390026807785, + -0.9991932511329651, + -0.7087220549583435, + -1.0, + 8.079589264298193e-8, + -3.263056755065918, + -5.687641620635986, + 0.23626811802387238, + -0.9920241236686707, + -0.3793386220932007, + -1.0, + 9.924601727107074e-7, + 0.0007141401874832809, + -0.04061261937022209 + ], + "q99": [ + 4.271139144897461, + 0.4542047083377838, + 0.8608905673027039, + 0.9988747239112854, + 0.813831627368927, + 1.0, + 1.0, + 6.125694274902344, + 0.5118050575256348, + 0.7076432704925537, + 0.9900418519973755, + 0.32602807879447937, + 1.0, + 1.0, + 0.04024730622768402, + -0.0001981566019821912 + ] + } + } +} \ No newline at end of file diff --git a/simulation/robocasa/full_attention/stage1/stage1_step160000.npz b/simulation/robocasa/full_attention/stage1/stage1_step160000.npz new file mode 100644 index 0000000000000000000000000000000000000000..82d11c4f2f77cfdd5b6fbe335b6ed890b80a2100 --- /dev/null +++ b/simulation/robocasa/full_attention/stage1/stage1_step160000.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:34387cbbd3891cc7bb082516dfb4993ee9d2319e32344a364ea32bd611cd924e +size 331912420 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/_CHECKPOINT_METADATA b/simulation/robocasa/full_attention/stage2/step-150000/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..75b9ff8b8eb8cba1e3c00483f7bdd39c10026a49 --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1788499436824938211, "commit_timestamp_nsecs": 1788499445933718228, "custom_metadata": {}} \ No newline at end of file diff --git a/simulation/robocasa/full_attention/stage2/step-150000/assets/robomme/norm_stats.json b/simulation/robocasa/full_attention/stage2/step-150000/assets/robomme/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..84350088fc65d4bab950ca32bb284a71339e6d46 --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/assets/robomme/norm_stats.json @@ -0,0 +1,136 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + 0.005957702174782753, + -0.009433974511921406, + -0.027537286281585693, + -0.0002329433336853981, + -0.0019665092695504427, + 0.0010746258776634932, + -0.17213615775108337, + 0.0022201647516340017, + -0.0015274320030584931, + 0.0009210301213897765, + -0.000019904651708202437, + -0.7345612645149231 + ], + "std": [ + 0.2822888493537903, + 0.31111952662467957, + 0.30976882576942444, + 0.060622796416282654, + 0.08188065886497498, + 0.08514650166034698, + 0.9420629143714905, + 0.06350194662809372, + 0.1455080211162567, + 0.05478019639849663, + 0.0006542060291394591, + 0.6832558512687683 + ], + "q01": [ + -1.0, + -1.0, + -0.6785714030265808, + -0.1971428543329239, + -0.27142858505249023, + -0.30571427941322327, + -1.0, + -0.07321428507566452, + -0.5982142686843872, + -0.24081368744373322, + 0.0, + -1.0 + ], + "q99": [ + 0.6482142806053162, + 1.0, + 1.0, + 0.1971428543329239, + 0.2514285743236542, + 0.3028571307659149, + 1.0, + 0.19107143580913544, + 0.612500011920929, + 0.24883341789245605, + 0.0, + 1.0 + ] + }, + "state": { + "mean": [ + 0.5936015248298645, + -0.20600375533103943, + 0.4474135637283325, + -0.30925947427749634, + -0.007771184667944908, + -0.04887167364358902, + 0.1312701404094696, + 2.4705684185028076, + -1.8421560525894165, + 0.6735693216323853, + -0.018410686403512955, + -0.002943282015621662, + 0.05312279239296913, + 0.595700740814209, + 0.03035394847393036, + -0.030480381101369858 + ], + "std": [ + 0.8089187145233154, + 0.6286742687225342, + 0.18121899664402008, + 0.8043732047080994, + 0.2596682906150818, + 0.349939227104187, + 0.21797646582126617, + 1.5330910682678223, + 1.3586890697479248, + 0.09442085027694702, + 0.2916492521762848, + 0.08581095188856125, + 0.6641682982444763, + 0.32726120948791504, + 0.011618087068200111, + 0.011510285548865795 + ], + "q01": [ + 0.042294517159461975, + -3.4426209926605225, + 0.0796390026807785, + -0.9991932511329651, + -0.7087220549583435, + -1.0, + 8.079589264298193e-8, + -3.263056755065918, + -5.687641620635986, + 0.23626811802387238, + -0.9920241236686707, + -0.3793386220932007, + -1.0, + 9.924601727107074e-7, + 0.0007141401874832809, + -0.04061261937022209 + ], + "q99": [ + 4.271139144897461, + 0.4542047083377838, + 0.8608905673027039, + 0.9988747239112854, + 0.813831627368927, + 1.0, + 1.0, + 6.125694274902344, + 0.5118050575256348, + 0.7076432704925537, + 0.9900418519973755, + 0.32602807879447937, + 1.0, + 1.0, + 0.04024730622768402, + -0.0001981566019821912 + ] + } + } +} \ No newline at end of file diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/_METADATA b/simulation/robocasa/full_attention/stage2/step-150000/params/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..3174ecb3aadf31e14f1f7c8b230a9cbe3b80c0dc --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_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": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 4304]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [14, 14, 3, 288]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "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": [288, 2048]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64288, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_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": [512]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 256, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 256, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 2048, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1024, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'kv_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "kv_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'mem_rms_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "mem_rms_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": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'out_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "out_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1, 256, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'q_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "q_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1, 1024, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 512, 16384]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4096, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 256, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256, 3072]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256, 3072]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_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": [256, 32]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128, 1024]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1024]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128, 1024]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128, 1024]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1024]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128, 1024]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'lang_film', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'lang_film', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 2048]}}, "('params', 'mem_encoder', 'latents', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "latents", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1024]}}, "('params', 'mem_encoder', 'proj_cosmo', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_cosmo', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "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, 1024]}}, "('params', 'mem_encoder', 'proj_recent', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_recent', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 1024]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 512]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1024]}}, "('params', 'time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/_sharding b/simulation/robocasa/full_attention/stage2/step-150000/params/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..8e8a2723950e24f11d6c7d7ef67a79f8ccf59c4d --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/_sharding @@ -0,0 +1 @@ +{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm1lbV9ybXNfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm91dF9laW5zdW1fbWVtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLmt2X2VpbnN1bV9tZW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLnFfZWluc3VtX21lbS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhdGVudHMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [4, 1], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}], [{\"id\": 1}], [{\"id\": 2}], [{\"id\": 3}]]}}"} \ No newline at end of file diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/array_metadatas/process_0 b/simulation/robocasa/full_attention/stage2/step-150000/params/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..5eacdbf35307909fee5f510e0873b2c738cafdbf --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 288, 4304], "chunk_shape": [27, 288, 4304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 4, 72, 1152], "chunk_shape": [27, 4, 72, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [288, 2048], "chunk_shape": [288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [256, 3072], "chunk_shape": [256, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 2, 256, 2048], "chunk_shape": [18, 2, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 2, 256, 1024], "chunk_shape": [18, 2, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 2, 2048, 256], "chunk_shape": [18, 2, 2048, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 2, 1024, 256], "chunk_shape": [18, 2, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.kv_einsum_mem.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.mem_rms_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.out_einsum_mem.w.value", "write_shape": [18, 1, 256, 1024], "chunk_shape": [18, 1, 256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.q_einsum_mem.w.value", "write_shape": [18, 1, 1024, 256], "chunk_shape": [18, 1, 1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.bias.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.kernel.value", "write_shape": [18, 256, 2048], "chunk_shape": [18, 256, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 512, 16384], "chunk_shape": [18, 2, 512, 16384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 256, 4096], "chunk_shape": [18, 2, 256, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [18, 256, 3072], "chunk_shape": [18, 256, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [18, 256, 3072], "chunk_shape": [18, 256, 3072], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.kernel.value", "write_shape": [2, 128, 1024], "chunk_shape": [2, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.kernel.value", "write_shape": [2, 128, 1024], "chunk_shape": [2, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.kernel.value", "write_shape": [2, 128, 1024], "chunk_shape": [2, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.kernel.value", "write_shape": [2, 128, 1024], "chunk_shape": [2, 128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.kernel.value", "write_shape": [512, 2048], "chunk_shape": [512, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.latents.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.kernel.value", "write_shape": [12, 1024], "chunk_shape": [12, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.kernel.value", "write_shape": [192, 1024], "chunk_shape": [192, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.kernel.value", "write_shape": [3, 512], "chunk_shape": [3, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [256, 1024], "chunk_shape": [256, 1024], "ext_metadata": null}}]} \ No newline at end of file diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9 b/simulation/robocasa/full_attention/stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9 new file mode 100644 index 0000000000000000000000000000000000000000..c2829ba644b3c90f69e2a24dbb19c2579eee130e --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:476ff74dc38c756c9c195faec2152f34b052f4263d9e19899bc40360de746aae +size 158450 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/manifest.ocdbt b/simulation/robocasa/full_attention/stage2/step-150000/params/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..f9a67dc627e4ec7ffe90cae0ae33f8049c68fea5 Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/manifest.ocdbt differ diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/06aabeab1a0e602ca8f8e217fcb96709 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/06aabeab1a0e602ca8f8e217fcb96709 new file mode 100644 index 0000000000000000000000000000000000000000..2a0b58f2c74a2f8b9539cde943dce7b6fad34243 Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/06aabeab1a0e602ca8f8e217fcb96709 differ diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9 new file mode 100644 index 0000000000000000000000000000000000000000..dba5303318fbf97dda54ca5bc7d3aa75ac8b3f99 --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:824eaea25c59bac7143776b1c37425d47a9426986d0246d389cc6a5eca912ed8 +size 606012846 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60 new file mode 100644 index 0000000000000000000000000000000000000000..236c8b96073f50c3551cf32f34d6f98c5882e494 --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2290868f5d3f5fb15c3e9342cc3fa0adeb4e0517ea2ba667c00e1f9555b32be0 +size 258762 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8 new file mode 100644 index 0000000000000000000000000000000000000000..3375329894cd2666f0875e243ab0b76f6482beea --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7088b323898c94fe320df837059a255f9ac31db7a06b98cb350a373b4e7e855f +size 1048614538 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5 new file mode 100644 index 0000000000000000000000000000000000000000..e9b55d519ab04cfff1769c24d2c9b5154df4bd03 --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3855446dea4243997fe5ad2e90d0e7a9f76d000c414f3e6998606716384078db +size 2239695988 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1 new file mode 100644 index 0000000000000000000000000000000000000000..80fb627928c156305708bd239c9cb1e017f26eec --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ed6ad005edb24d33eb04226b8bee988a0ef5c44e3ef94ec84b56765001326be4 +size 2239686238 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/822b3aef48f3ec46a514c1a9b663d129 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/822b3aef48f3ec46a514c1a9b663d129 new file mode 100644 index 0000000000000000000000000000000000000000..ad055f7a7f18fe1bccc605d0ae354067803db860 Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/822b3aef48f3ec46a514c1a9b663d129 differ diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891 new file mode 100644 index 0000000000000000000000000000000000000000..db849424a23ca20f90e975d1112ba1634ff4cbb3 --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d60b237dba5c9d4a74f511cb55156e8a59b605339a72cbfb12a6c3fab1b1a6f8 +size 1120293104 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/874b4177fd57749dda8e5c8a08b34524 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/874b4177fd57749dda8e5c8a08b34524 new file mode 100644 index 0000000000000000000000000000000000000000..0d52b5cf70d83b1810119f630b224e340320c7a6 Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/874b4177fd57749dda8e5c8a08b34524 differ diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/96836e12295a196861bcac1091e408c6 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/96836e12295a196861bcac1091e408c6 new file mode 100644 index 0000000000000000000000000000000000000000..eb22654e4958d0d465397b732958faf33fe3dac0 Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/96836e12295a196861bcac1091e408c6 differ diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b1b9c620755ae82fa0958a1aba0a6853 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b1b9c620755ae82fa0958a1aba0a6853 new file mode 100644 index 0000000000000000000000000000000000000000..0277dc69c3e76357fded92f1415e96dd146bf0c7 Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b1b9c620755ae82fa0958a1aba0a6853 differ diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b726d8d65685cae9811fa5a7918320d3 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b726d8d65685cae9811fa5a7918320d3 new file mode 100644 index 0000000000000000000000000000000000000000..1a2ea0ca6e8386bbc6eb3e79fdfea23edef271b3 Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/b726d8d65685cae9811fa5a7918320d3 differ diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065 new file mode 100644 index 0000000000000000000000000000000000000000..004d36c4434232db9f660e7a9f77e3cef88c328e --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d58075ff02f957a5bfc1f9892c8956b19df5dfd13968cea4cb51530ded6b1137 +size 158430 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761 new file mode 100644 index 0000000000000000000000000000000000000000..d484b7f2caa187a89b614304a4f65e483bc496e5 --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b42d860c3b83040a1bf3da8ae1e73497e29fa83d4ab6d7c7bf2ee1f9a9bac22c +size 2587810029 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25 b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25 new file mode 100644 index 0000000000000000000000000000000000000000..af09e1794f91285055749415128ca909d6fe5040 --- /dev/null +++ b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7f3e1d8d11d6d0f2d58a47be5d1c47d533b9c07d892e41057f1acc6cb2be23a9 +size 2117122051 diff --git a/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..b3c0725729cac774cfbd49806482419fa6f9405b Binary files /dev/null and b/simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt differ diff --git a/simulation/robocasa/recurrent/.gitattributes b/simulation/robocasa/recurrent/.gitattributes new file mode 100644 index 0000000000000000000000000000000000000000..8df4ba5a44a559b25144859e1653e7f8f309c3af --- /dev/null +++ b/simulation/robocasa/recurrent/.gitattributes @@ -0,0 +1,44 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.xz 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 +stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 filter=lfs diff=lfs merge=lfs -text +stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 filter=lfs diff=lfs merge=lfs -text diff --git a/simulation/robocasa/recurrent/LICENSE_GEMMA.txt b/simulation/robocasa/recurrent/LICENSE_GEMMA.txt new file mode 100644 index 0000000000000000000000000000000000000000..e0f8930d83fe9eef924c0eed50ee0e572eb471fa --- /dev/null +++ b/simulation/robocasa/recurrent/LICENSE_GEMMA.txt @@ -0,0 +1,77 @@ +Gemma Terms of Use + +Last modified: February 21, 2024 + +By using, reproducing, modifying, distributing, performing or displaying any portion or element of Gemma, Model Derivatives including via any Hosted Service, (each as defined below) (collectively, the "Gemma Services") or otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement. + +Section 1: DEFINITIONS +1.1 Definitions +(a) "Agreement" or "Gemma Terms of Use" means these terms and conditions that govern the use, reproduction, Distribution or modification of the Gemma Services and any terms and conditions incorporated by reference. + +(b) "Distribution" or "Distribute" means any transmission, publication, or other sharing of Gemma or Model Derivatives to a third party, including by providing or making Gemma or its functionality available as a hosted service via API, web access, or any other electronic or remote means ("Hosted Service"). + +(c) "Gemma" means the set of machine learning language models, trained model weights and parameters identified at ai.google.dev/gemma, regardless of the source that you obtained it from. + +(d) "Google" means Google LLC. + +(e) "Model Derivatives" means all (i) modifications to Gemma, (ii) works based on Gemma, or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Gemma, to that model in order to cause that model to perform similarly to Gemma, including distillation methods that use intermediate data representations or methods based on the generation of synthetic data Outputs by Gemma for training that model. For clarity, Outputs are not deemed Model Derivatives. + +(f) "Output" means the information content output of Gemma or a Model Derivative that results from operating or otherwise using Gemma or the Model Derivative, including via a Hosted Service. + +1.2 +As used in this Agreement, "including" means "including without limitation". + +Section 2: ELIGIBILITY AND USAGE +2.1 Eligibility +You represent and warrant that you have the legal capacity to enter into this Agreement (including being of sufficient age of consent). If you are accessing or using any of the Gemma Services for or on behalf of a legal entity, (a) you are entering into this Agreement on behalf of yourself and that legal entity, (b) you represent and warrant that you have the authority to act on behalf of and bind that entity to this Agreement and (c) references to "you" or "your" in the remainder of this Agreement refers to both you (as an individual) and that entity. + +2.2 Use +You may use, reproduce, modify, Distribute, perform or display any of the Gemma Services only in accordance with the terms of this Agreement, and must not violate (or encourage or permit anyone else to violate) any term of this Agreement. + +Section 3: DISTRIBUTION AND RESTRICTIONS +3.1 Distribution and Redistribution +You may reproduce or Distribute copies of Gemma or Model Derivatives if you meet all of the following conditions: + +You must include the use restrictions referenced in Section 3.2 as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Gemma or Model Derivatives and you must provide notice to subsequent users you Distribute to that Gemma or Model Derivatives are subject to the use restrictions in Section 3.2. +You must provide all third party recipients of Gemma or Model Derivatives a copy of this Agreement. +You must cause any modified files to carry prominent notices stating that you modified the files. +All Distributions (other than through a Hosted Service) must be accompanied by a "Notice" text file that contains the following notice: "Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms". +You may add your own intellectual property statement to your modifications and, except as set forth in this Section, may provide additional or different terms and conditions for use, reproduction, or Distribution of your modifications, or for any such Model Derivatives as a whole, provided your use, reproduction, modification, Distribution, performance, and display of Gemma otherwise complies with the terms and conditions of this Agreement. Any additional or different terms and conditions you impose must not conflict with the terms of this Agreement. + +3.2 Use Restrictions +You must not use any of the Gemma Services: + +for the restricted uses set forth in the Gemma Prohibited Use Policy at ai.google.dev/gemma/prohibited_use_policy ("Prohibited Use Policy"), which is hereby incorporated by reference into this Agreement; or +in violation of applicable laws and regulations. +To the maximum extent permitted by law, Google reserves the right to restrict (remotely or otherwise) usage of any of the Gemma Services that Google reasonably believes are in violation of this Agreement. + +3.3 Generated Output +Google claims no rights in Outputs you generate using Gemma. You and your users are solely responsible for Outputs and their subsequent uses. + +Section 4: ADDITIONAL PROVISIONS +4.1 Updates +Google may update Gemma from time to time, and you must make reasonable efforts to use the latest version of Gemma. + +4.2 Trademarks +Nothing in this Agreement grants you any rights to use Google's trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between you and Google. Google reserves any rights not expressly granted herein. + +4.3 DISCLAIMER OF WARRANTY +UNLESS REQUIRED BY APPLICABLE LAW, THE GEMMA SERVICES, AND OUTPUTS, ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY WARRANTIES OR CONDITIONS OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR OR DISTRIBUTING ANY OF THE GEMMA SERVICES OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR USE OR DISTRIBUTION OF ANY OF THE GEMMA SERVICES OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT. + +4.4 LIMITATION OF LIABILITY +TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY, CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW, SHALL GOOGLE OR ITS AFFILIATES BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL, OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO, ANY OF THE GEMMA SERVICES OR OUTPUTS EVEN IF GOOGLE OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. + +4.5 Term, Termination, and Survival +The term of this Agreement will commence upon your acceptance of this Agreement (including acceptance by your use, modification, or Distribution, reproduction, performance or display of any portion or element of the Gemma Services) and will continue in full force and effect until terminated in accordance with the terms of this Agreement. Google may terminate this Agreement if you are in breach of any term of this Agreement. Upon termination of this Agreement, you must delete and cease use and Distribution of all copies of Gemma and Model Derivatives in your possession or control. Sections 1, 2.1, 3.3, 4.2 to 4.9 shall survive the termination of this Agreement. + +4.6 Governing Law and Jurisdiction +This Agreement will be governed by the laws of the State of California without regard to choice of law principles. The UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The state and federal courts of Santa Clara County, California shall have exclusive jurisdiction of any dispute arising out of this Agreement. + +4.7 Severability +If any provision of this Agreement is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein. + +4.8 Entire Agreement +This Agreement states all the terms agreed between the parties and supersedes all other agreements between the parties as of the date of acceptance relating to its subject matter. + +4.9 No Waiver +Google will not be treated as having waived any rights by not exercising (or delaying the exercise of) any rights under this Agreement. \ No newline at end of file diff --git a/simulation/robocasa/recurrent/NOTICE b/simulation/robocasa/recurrent/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..1b9cb6541d7135c4c78a542ac4a3e04a971c96db --- /dev/null +++ b/simulation/robocasa/recurrent/NOTICE @@ -0,0 +1,11 @@ +IMPRINT +Copyright 2026 + +This distribution contains modifications to the OpenPI project: +https://github.com/Physical-Intelligence/openpi + +OpenPI is distributed under the Apache License 2.0. Gemma components and +derivatives are subject to the Gemma Terms of Use in LICENSE_GEMMA.txt. + +Gemma is provided under and subject to the Gemma Terms of Use found at +ai.google.dev/gemma/terms. diff --git a/simulation/robocasa/recurrent/README.md b/simulation/robocasa/recurrent/README.md new file mode 100644 index 0000000000000000000000000000000000000000..7d64b6dee24b77f06ef7f2a5f55a3ba9504289f2 --- /dev/null +++ b/simulation/robocasa/recurrent/README.md @@ -0,0 +1,38 @@ +--- +license: gemma +library_name: jax +tags: + - robotics + - vision-language-action + - memory + - robocasa + - dinov3 + - recurrent +--- + +# IMPRINT RoboCasa Recurrent Checkpoint + +Canonical Pi0.5 IMPRINT checkpoint for the customized 12-task RoboCasa +benchmark using the recurrent gated-delta memory encoder. + +- Stage 1: trained 180,000 updates, selected at 160,000, W&B `far-wandb/Robocasa/ag6ytzrm` +- Stage 2: 150,000 updates, W&B `far-wandb/Robocasa/1ozs9qy8` +- memory tokens: 256 x 1024 +- policy horizon: 20 +- training scan: chunkwise WY +- inference update: incremental recurrent +- Stage 2: memory encoder finetuned; image tower frozen + +The Stage-2 model-level auxiliary weight was zero. The archived history YAML +retains a legacy `aux_weight: 0.25` preload flag, which caused auxiliary +targets to be read but did not contribute an auxiliary loss. + +Use `yilin-wu/imprint-robocasa-600` for training and +`yilin-wu/imprint-robocasa-environment` for the exact patched RoboCasa 1.0.1 +simulator source. This release contains the complete Stage-2 Orbax parameter +checkpoint, including the PaliGemma/Gemma backbone. It is a Gemma model +derivative and is distributed subject to the included Gemma Terms of Use. +Large simulator assets are not duplicated. + +The matching training store, patched simulator source, and reproducibility +code are publicly available under the IMPRINT release. diff --git a/simulation/robocasa/recurrent/SHA256SUMS b/simulation/robocasa/recurrent/SHA256SUMS new file mode 100644 index 0000000000000000000000000000000000000000..181d534d8e76656e3b938d7257ebceb4a7a1bf91 --- /dev/null +++ b/simulation/robocasa/recurrent/SHA256SUMS @@ -0,0 +1,27 @@ +3e2c24001f9ef57bf7ec959a3658fbb49cdad113cdf394c264da9d16f9bdd132 ./LICENSE_GEMMA.txt +5b21d7116ac42a0ab6141488cb654c2411557ac977e2276accb88ef774f28c19 ./NOTICE +1c4d7719baef48359e81dc02e729ec9361c51acf5ba3bd788f21dcd3268b5277 ./README.md +3f5ee9416db5224d6effaea03e799540023817301a0ba2677922a31709216b1d ./config/history_config.yaml +4143882de7d4e41230b03bd95a7001745139f6a757e64b4db3afa2857ec5b9e7 ./config/input_provenance.json +39a8f9a475d18c9167a3872c7679a348b81a165776fe3e16c391f354c7ad8514 ./config/norm_stats.json +7567aca14012be88045a10203432c7117c128b4b1455718143c429cea8155fc1 ./stage1/stage1_step160000.npz +fe7ad82eeb9e1096c90f433a50fb4b2fabe000779a6d14bfe8ee317f6aea9c72 ./stage2/step-150000/_CHECKPOINT_METADATA +39a8f9a475d18c9167a3872c7679a348b81a165776fe3e16c391f354c7ad8514 ./stage2/step-150000/assets/robomme/norm_stats.json +68305898798e4aa166cbeff6f5375e02d51d18c84920784c0a7b34159292f481 ./stage2/step-150000/params/_METADATA +306825e3332fd927d96181a73228f3480a3ceb705c7db27fe49a2563204ebb8f ./stage2/step-150000/params/_sharding +d86f5bf955a3a93391b23797e2c7e1cd146341f989074e96b0678b8351067c88 ./stage2/step-150000/params/array_metadatas/process_0 +1d455b9c7cd4010bac48990f5b286f3d4764e677b577361e584d5de51702d72d ./stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 +c56bbba464e09062f252686a9b61c77d21f4f6c9f5176c0ce0e07f5ccab96860 ./stage2/step-150000/params/manifest.ocdbt +2f44abe9cc4c9b57417a84eccbc4cfaacbf78138b754f58deb3663e8f0b9cd0c ./stage2/step-150000/params/ocdbt.process_0/d/2a2b816d8bc9e5ca108e10d8c0ec697b +78c760f9903cc094536408330040c07125539871571bcd8f66addfeff3489211 ./stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 +ede3e60d71ddb1726e98428f57b9be80a113cc3101a18c8c1a0362224951497a ./stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb +57a8cccbbd6268bbe3cd7892014c27da7f4bc9e734cea8dd64881c144d15a9f5 ./stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 +d5c5bbdb19cc2144308820ef54b0382b2eb07092c2ed0ea67a251914752735ae ./stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 +81594df81b1b9b7ae76963f695bff5e4ed1da3a0d7e29dea42475d84e432e6b9 ./stage2/step-150000/params/ocdbt.process_0/d/5a62244cc2b71e7460df9b3d05faf966 +8b3d3995e3aef785f5244dae3e6588ddf1f0386876778f95303b68f3b41d71ef ./stage2/step-150000/params/ocdbt.process_0/d/62a2c855e8bf17ca63d6be7a948a943f +d86dfa44bb73ca13d774185159e069a4d7aa6585af4a569f9f7569b9460b3392 ./stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 +2144a693e034a9dab2595f94cfaf7f29f9a92384fdb8827e7556ecd3b8d42f53 ./stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 +034c510b6b4a7b7a01bff53cd076593b5f8e993f6497015f311517b1feced6d7 ./stage2/step-150000/params/ocdbt.process_0/d/d891aac544c185edcc57e57a7f69561c +af709070d5c1a4b648d7799923f503c658d2403dfffed9a54fc301a89d18a1f9 ./stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 +3614c36b08b05532184f2b34cd514d9d2dcc3f0be7038e846356371d62e52e54 ./stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 +7c124c19a20c79bc57bde2839dee54753c2290f0fe166070dcf59258e3222f1b ./stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt diff --git a/simulation/robocasa/recurrent/config/history_config.yaml b/simulation/robocasa/recurrent/config/history_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e3d972ab3e3d527bbefd8115c06108707d3eda0b --- /dev/null +++ b/simulation/robocasa/recurrent/config/history_config.yaml @@ -0,0 +1,89 @@ +budget: 256 +num_views: 3 +streaming_obs_horizon: 16 +integration_type: modulation +memory_token_dim: 1024 +representation_type: cosmo_bottleneck +memory_feature: + img: + net: identity + input_dim: 768 + pos: + input_dim: 768 + hidden_dim: 768 + state: + input_dim: 16 + hidden_dim: 512 +cosmo_bottleneck_memory: + memory_encoder_type: gated_delta_hybrid + memory_architecture_version: gated_delta_hybrid_v1 + memory_compute_dtype: bfloat16 + visual_feature_stream: dinov3 + latent_source: offline + n_variants: 8 + num_query_samples: 16 + variant: Wan2.2-TI2V-5B + temporal_factor: 4 + downsample: 2 + num_latents: 256 + num_heads: 8 + cosmos_latent_dim: 48 + siglip_dim: 768 + proprio_in_dim: 16 + proprio_encoding: raw + action_in_dim: 12 + action_encoding: raw + lang_emb_dim: 2048 + action_tokenizer_hidden: 512 + proprio_tokenizer_hidden: 512 + pool_hw: 2 + wan_h: 256 + wan_w: 256 + latent_cap: 4096 + latent_buckets: + - 256 + - 512 + - 1024 + - 2048 + - 4096 + max_recent_tokens: 768 + max_step_tokens: 1776 + max_stream_steps: 1776 + delta_num_layers: 2 + delta_num_heads: 8 + delta_head_dim: 128 + delta_conv_kernel: 4 + delta_ffn_inner: 2048 + delta_decay_bias: -4.0 + delta_beta_bias: -1.0 + delta_scan_checkpoint_size: 20 + delta_full_scan_mode: chunkwise_wy + delta_wy_matmul_dtype: bfloat16 + delta_recent_fusion_blocks: 1 + delta_max_chunks: 384 + delta_event_order: action_proprio_then_finalized_visual_v1 +cosmo_decoder: + num_heads: 8 + num_blocks: 2 + align_dim: 768 + align_num_patches: 192 + ffn_inner: 8192 + action_dim: 12 + proprio_dim: 16 + max_query_steps: 2048 +frame_jitter: + enabled: false + start_jitter: 3 + stride_choices: + - 1 + - 2 + - 3 + interval_jitter: 1 + canonical_downsample: 2 +stream_aug: + enabled: false + gauss_sigma: 0.05 + stream_drop_p: 0.1 + feat_drop_p: 0.0 +aux_weight: 0.25 +action_horizon: 20 diff --git a/simulation/robocasa/recurrent/config/input_provenance.json b/simulation/robocasa/recurrent/config/input_provenance.json new file mode 100644 index 0000000000000000000000000000000000000000..59e7878ede6291481d9539c527bd4e7be4b106d7 --- /dev/null +++ b/simulation/robocasa/recurrent/config/input_provenance.json @@ -0,0 +1,28 @@ +{ + "action_causality": "strictly_causal", + "architecture": "gated_delta_hybrid_v1", + "base_policy": "pi0.5", + "checkpoint_metadata_sha256": "fe7ad82eeb9e1096c90f433a50fb4b2fabe000779a6d14bfe8ee317f6aea9c72", + "checkpoint_path": "stage2/step-150000", + "checkpoint_step": 150000, + "dataset": "yilin-wu/imprint-robocasa-600", + "eval_consolidate_all": 1, + "execution_horizon": 20, + "freeze_variant": "image_only", + "latent_cap": 4096, + "legacy_history_aux_preload_flag": 0.25, + "max_recent_tokens": 768, + "max_step_tokens": 1776, + "memory_inference_mode": "incremental_recurrent", + "memory_language_source": "live", + "memory_training_scan": "chunkwise_wy", + "num_latents": 256, + "prediction_horizon": 20, + "stage1_checkpoint": "stage1/stage1_step160000.npz", + "stage1_sha256": "7567aca14012be88045a10203432c7117c128b4b1455718143c429cea8155fc1", + "stage1_wandb_run": "far-wandb/Robocasa/ag6ytzrm", + "stage2_auxiliary_loss_weight": 0.0, + "stage2_memory_encoder": "finetuned", + "stage2_wandb_run": "far-wandb/Robocasa/1ozs9qy8", + "visual_feature_stream": "dinov3" +} diff --git a/simulation/robocasa/recurrent/config/norm_stats.json b/simulation/robocasa/recurrent/config/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..84350088fc65d4bab950ca32bb284a71339e6d46 --- /dev/null +++ b/simulation/robocasa/recurrent/config/norm_stats.json @@ -0,0 +1,136 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + 0.005957702174782753, + -0.009433974511921406, + -0.027537286281585693, + -0.0002329433336853981, + -0.0019665092695504427, + 0.0010746258776634932, + -0.17213615775108337, + 0.0022201647516340017, + -0.0015274320030584931, + 0.0009210301213897765, + -0.000019904651708202437, + -0.7345612645149231 + ], + "std": [ + 0.2822888493537903, + 0.31111952662467957, + 0.30976882576942444, + 0.060622796416282654, + 0.08188065886497498, + 0.08514650166034698, + 0.9420629143714905, + 0.06350194662809372, + 0.1455080211162567, + 0.05478019639849663, + 0.0006542060291394591, + 0.6832558512687683 + ], + "q01": [ + -1.0, + -1.0, + -0.6785714030265808, + -0.1971428543329239, + -0.27142858505249023, + -0.30571427941322327, + -1.0, + -0.07321428507566452, + -0.5982142686843872, + -0.24081368744373322, + 0.0, + -1.0 + ], + "q99": [ + 0.6482142806053162, + 1.0, + 1.0, + 0.1971428543329239, + 0.2514285743236542, + 0.3028571307659149, + 1.0, + 0.19107143580913544, + 0.612500011920929, + 0.24883341789245605, + 0.0, + 1.0 + ] + }, + "state": { + "mean": [ + 0.5936015248298645, + -0.20600375533103943, + 0.4474135637283325, + -0.30925947427749634, + -0.007771184667944908, + -0.04887167364358902, + 0.1312701404094696, + 2.4705684185028076, + -1.8421560525894165, + 0.6735693216323853, + -0.018410686403512955, + -0.002943282015621662, + 0.05312279239296913, + 0.595700740814209, + 0.03035394847393036, + -0.030480381101369858 + ], + "std": [ + 0.8089187145233154, + 0.6286742687225342, + 0.18121899664402008, + 0.8043732047080994, + 0.2596682906150818, + 0.349939227104187, + 0.21797646582126617, + 1.5330910682678223, + 1.3586890697479248, + 0.09442085027694702, + 0.2916492521762848, + 0.08581095188856125, + 0.6641682982444763, + 0.32726120948791504, + 0.011618087068200111, + 0.011510285548865795 + ], + "q01": [ + 0.042294517159461975, + -3.4426209926605225, + 0.0796390026807785, + -0.9991932511329651, + -0.7087220549583435, + -1.0, + 8.079589264298193e-8, + -3.263056755065918, + -5.687641620635986, + 0.23626811802387238, + -0.9920241236686707, + -0.3793386220932007, + -1.0, + 9.924601727107074e-7, + 0.0007141401874832809, + -0.04061261937022209 + ], + "q99": [ + 4.271139144897461, + 0.4542047083377838, + 0.8608905673027039, + 0.9988747239112854, + 0.813831627368927, + 1.0, + 1.0, + 6.125694274902344, + 0.5118050575256348, + 0.7076432704925537, + 0.9900418519973755, + 0.32602807879447937, + 1.0, + 1.0, + 0.04024730622768402, + -0.0001981566019821912 + ] + } + } +} \ No newline at end of file diff --git a/simulation/robocasa/recurrent/stage1/stage1_step160000.npz b/simulation/robocasa/recurrent/stage1/stage1_step160000.npz new file mode 100644 index 0000000000000000000000000000000000000000..3a3e7c32bac553285e3444b4950470e53a2efb12 --- /dev/null +++ b/simulation/robocasa/recurrent/stage1/stage1_step160000.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7567aca14012be88045a10203432c7117c128b4b1455718143c429cea8155fc1 +size 391011410 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/_CHECKPOINT_METADATA b/simulation/robocasa/recurrent/stage2/step-150000/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..6a4b614b88361e569f91e3b944e3c966ee7b2974 --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/_CHECKPOINT_METADATA @@ -0,0 +1 @@ +{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1789085274077480773, "commit_timestamp_nsecs": 1789085308428284253, "custom_metadata": {}} \ No newline at end of file diff --git a/simulation/robocasa/recurrent/stage2/step-150000/assets/robomme/norm_stats.json b/simulation/robocasa/recurrent/stage2/step-150000/assets/robomme/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..84350088fc65d4bab950ca32bb284a71339e6d46 --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/assets/robomme/norm_stats.json @@ -0,0 +1,136 @@ +{ + "norm_stats": { + "actions": { + "mean": [ + 0.005957702174782753, + -0.009433974511921406, + -0.027537286281585693, + -0.0002329433336853981, + -0.0019665092695504427, + 0.0010746258776634932, + -0.17213615775108337, + 0.0022201647516340017, + -0.0015274320030584931, + 0.0009210301213897765, + -0.000019904651708202437, + -0.7345612645149231 + ], + "std": [ + 0.2822888493537903, + 0.31111952662467957, + 0.30976882576942444, + 0.060622796416282654, + 0.08188065886497498, + 0.08514650166034698, + 0.9420629143714905, + 0.06350194662809372, + 0.1455080211162567, + 0.05478019639849663, + 0.0006542060291394591, + 0.6832558512687683 + ], + "q01": [ + -1.0, + -1.0, + -0.6785714030265808, + -0.1971428543329239, + -0.27142858505249023, + -0.30571427941322327, + -1.0, + -0.07321428507566452, + -0.5982142686843872, + -0.24081368744373322, + 0.0, + -1.0 + ], + "q99": [ + 0.6482142806053162, + 1.0, + 1.0, + 0.1971428543329239, + 0.2514285743236542, + 0.3028571307659149, + 1.0, + 0.19107143580913544, + 0.612500011920929, + 0.24883341789245605, + 0.0, + 1.0 + ] + }, + "state": { + "mean": [ + 0.5936015248298645, + -0.20600375533103943, + 0.4474135637283325, + -0.30925947427749634, + -0.007771184667944908, + -0.04887167364358902, + 0.1312701404094696, + 2.4705684185028076, + -1.8421560525894165, + 0.6735693216323853, + -0.018410686403512955, + -0.002943282015621662, + 0.05312279239296913, + 0.595700740814209, + 0.03035394847393036, + -0.030480381101369858 + ], + "std": [ + 0.8089187145233154, + 0.6286742687225342, + 0.18121899664402008, + 0.8043732047080994, + 0.2596682906150818, + 0.349939227104187, + 0.21797646582126617, + 1.5330910682678223, + 1.3586890697479248, + 0.09442085027694702, + 0.2916492521762848, + 0.08581095188856125, + 0.6641682982444763, + 0.32726120948791504, + 0.011618087068200111, + 0.011510285548865795 + ], + "q01": [ + 0.042294517159461975, + -3.4426209926605225, + 0.0796390026807785, + -0.9991932511329651, + -0.7087220549583435, + -1.0, + 8.079589264298193e-8, + -3.263056755065918, + -5.687641620635986, + 0.23626811802387238, + -0.9920241236686707, + -0.3793386220932007, + -1.0, + 9.924601727107074e-7, + 0.0007141401874832809, + -0.04061261937022209 + ], + "q99": [ + 4.271139144897461, + 0.4542047083377838, + 0.8608905673027039, + 0.9988747239112854, + 0.813831627368927, + 1.0, + 1.0, + 6.125694274902344, + 0.5118050575256348, + 0.7076432704925537, + 0.9900418519973755, + 0.32602807879447937, + 1.0, + 1.0, + 0.04024730622768402, + -0.0001981566019821912 + ] + } + } +} \ No newline at end of file diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/_METADATA b/simulation/robocasa/recurrent/stage2/step-150000/params/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..636ae950c6862f3ae9d4c1a0b4e40f216af0afd0 --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/_METADATA @@ -0,0 +1 @@ +{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_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": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1152, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 16, 72, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [14, 14, 3, 288]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "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": [1152, 512]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64288, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_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": [512]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'kv_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "kv_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'mem_rms_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "mem_rms_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": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'out_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "out_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'q_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "q_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 2048, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4096, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_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": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_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": [256, 32]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'beta_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "beta_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'beta_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "beta_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": [256, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'decay_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "decay_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'decay_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "decay_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": [256, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 512]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'ffn_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'norm_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'norm_in', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_gate', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_gate', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "out_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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '0', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "0", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'beta_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "beta_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'beta_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "beta_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": [256, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'conv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'conv', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'decay_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "decay_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'decay_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "decay_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": [256, 8]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 512]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'ffn_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'k_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'k_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'norm_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'norm_in', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_in", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_gate', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_gate', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_gate", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_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": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "out_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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'q_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'q_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1024, 256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'v_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'delta_blocks', '1', 'v_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "delta_blocks", "key_type": 2}, {"key": "1", "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": [1024, 256]}}, "('params', 'mem_encoder', 'event_slot', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "event_slot", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [5, 1024]}}, "('params', 'mem_encoder', 'event_type', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "event_type", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 256]}}, "('params', 'mem_encoder', 'lang_film', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'lang_film', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('params', 'mem_encoder', 'latents', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "latents", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1024]}}, "('params', 'mem_encoder', 'proj_cosmo', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_cosmo', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "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, 1024]}}, "('params', 'mem_encoder', 'proj_recent', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_recent', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 1024]}}, "('params', 'mem_encoder', 'read_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'read_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_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": [256]}}, "('params', 'mem_encoder', 'read_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'read_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'read_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'read_q', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "read_q", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_kv_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_kv_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_kv_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_kv_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": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_q_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_q_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'cross_q_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "cross_q_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": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 512]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_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": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'ffn_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'recent_blocks', '0', 'self_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_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": [256]}}, "('params', 'mem_encoder', 'recent_null', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "recent_null", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 256]}}, "('params', 'mem_encoder', 'time_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "time_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'time_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "time_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 256]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [3, 512]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null} \ No newline at end of file diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/_sharding b/simulation/robocasa/recurrent/stage2/step-150000/params/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..4a6e96dbc26510168ad386a01cd30f863dfac7c5 --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/_sharding @@ -0,0 +1 @@ +{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm1lbV9ybXNfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm91dF9laW5zdW1fbWVtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLmt2X2VpbnN1bV9tZW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLnFfZWluc3VtX21lbS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm5vcm1faW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm5vcm1faW4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9nYXRlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9nYXRlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmJldGFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmJldGFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmNvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmNvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmRlY2F5X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmRlY2F5X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmZmbl9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmtfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnZfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4wLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm5vcm1faW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm5vcm1faW4uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9nYXRlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9nYXRlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLm91dF9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmJldGFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmJldGFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmNvbnYuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmNvbnYua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmRlY2F5X3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmRlY2F5X3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmZmbl9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmtfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLmtfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnFfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnFfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnZfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmRlbHRhX2Jsb2Nrcy4xLnZfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmV2ZW50X3Nsb3QudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmV2ZW50X3R5cGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhdGVudHMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLm91dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLm91dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLmtleS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLmtleS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnF1ZXJ5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnF1ZXJ5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnZhbHVlLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19hdHRuLnZhbHVlLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19rdl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19rdl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19xX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5jcm9zc19xX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5faW4uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5faW4ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5mZm5fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX25vcm0uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX25vcm0uc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9ibG9ja3MuMC5zZWxmX2F0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlY2VudF9udWxsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfb3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfb3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfbm9ybS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnJlYWRfcS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRpbWVfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRpbWVfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}"} \ No newline at end of file diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/array_metadatas/process_0 b/simulation/robocasa/recurrent/stage2/step-150000/params/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..1674d2dff173c2f90742beebf091c36816411bbe --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/array_metadatas/process_0 @@ -0,0 +1 @@ +{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 1152, 1076], "chunk_shape": [27, 1152, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 16, 72, 288], "chunk_shape": [27, 16, 72, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [1152, 512], "chunk_shape": [1152, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [1024, 768], "chunk_shape": [1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 8, 256, 512], "chunk_shape": [18, 8, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 8, 512, 256], "chunk_shape": [18, 8, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.kv_einsum_mem.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.mem_rms_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.out_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.q_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.bias.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.kernel.value", "write_shape": [18, 1024, 512], "chunk_shape": [18, 1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 2048, 4096], "chunk_shape": [18, 2, 2048, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 1024, 1024], "chunk_shape": [18, 2, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.beta_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.beta_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.conv.kernel.value", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.decay_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.decay_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_in.kernel.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.ffn_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.k_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.norm_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.norm_in.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_gate.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_gate.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.out_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.q_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.0.v_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.beta_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.beta_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.conv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.conv.kernel.value", "write_shape": [1, 1024], "chunk_shape": [1, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.decay_proj.bias.value", "write_shape": [2], "chunk_shape": [2], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.decay_proj.kernel.value", "write_shape": [256, 8], "chunk_shape": [256, 8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_in.kernel.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.ffn_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.k_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.k_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.norm_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.norm_in.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_gate.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_gate.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.out_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.q_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.q_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.v_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.delta_blocks.1.v_proj.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.event_slot.value", "write_shape": [5, 1024], "chunk_shape": [5, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.event_type.value", "write_shape": [3, 256], "chunk_shape": [3, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.latents.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.kernel.value", "write_shape": [12, 1024], "chunk_shape": [12, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.kernel.value", "write_shape": [192, 1024], "chunk_shape": [192, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.read_q.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_kv_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_kv_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_q_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.cross_q_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_in.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_in.kernel.value", "write_shape": [1024, 512], "chunk_shape": [1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.ffn_out.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_norm.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_blocks.0.self_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.recent_null.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.time_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.time_proj.kernel.value", "write_shape": [1, 256], "chunk_shape": [1, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.kernel.value", "write_shape": [3, 512], "chunk_shape": [3, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}]} \ No newline at end of file diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 b/simulation/robocasa/recurrent/stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 new file mode 100644 index 0000000000000000000000000000000000000000..e27154cde53ee258ce7f1c5abe7a4473a84c77d3 --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1d455b9c7cd4010bac48990f5b286f3d4764e677b577361e584d5de51702d72d +size 196861 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/manifest.ocdbt b/simulation/robocasa/recurrent/stage2/step-150000/params/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..a1a37be397ad0662fac95ccf7f00b4b7936da6ec Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/manifest.ocdbt differ diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2a2b816d8bc9e5ca108e10d8c0ec697b b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2a2b816d8bc9e5ca108e10d8c0ec697b new file mode 100644 index 0000000000000000000000000000000000000000..cd287190512c3f96a737bc04bdb8b504e4e9ef59 Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2a2b816d8bc9e5ca108e10d8c0ec697b differ diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 new file mode 100644 index 0000000000000000000000000000000000000000..2a43b49d5ce574d513aa1404b5b0c95e29748d27 --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:78c760f9903cc094536408330040c07125539871571bcd8f66addfeff3489211 +size 2030582856 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb new file mode 100644 index 0000000000000000000000000000000000000000..258f02dc60b38c6e34b3e0def1c9855a70f1ff07 --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ede3e60d71ddb1726e98428f57b9be80a113cc3101a18c8c1a0362224951497a +size 196857 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 new file mode 100644 index 0000000000000000000000000000000000000000..6887e66c412833401819a5005c6060f358e1302c --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:57a8cccbbd6268bbe3cd7892014c27da7f4bc9e734cea8dd64881c144d15a9f5 +size 2240054650 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 new file mode 100644 index 0000000000000000000000000000000000000000..252e72af5487357dd2d528fa1fd2d7bdb006a9ba --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d5c5bbdb19cc2144308820ef54b0382b2eb07092c2ed0ea67a251914752735ae +size 79040671 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/5a62244cc2b71e7460df9b3d05faf966 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/5a62244cc2b71e7460df9b3d05faf966 new file mode 100644 index 0000000000000000000000000000000000000000..49ff9c74e5e88dcd4713c142d6a0599d79afc148 Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/5a62244cc2b71e7460df9b3d05faf966 differ diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/62a2c855e8bf17ca63d6be7a948a943f b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/62a2c855e8bf17ca63d6be7a948a943f new file mode 100644 index 0000000000000000000000000000000000000000..55ca384744eb81310252f6fce617225cb731dae0 Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/62a2c855e8bf17ca63d6be7a948a943f differ diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 new file mode 100644 index 0000000000000000000000000000000000000000..4e1fc196beb146f21ef9a34c84078a1564a9d77b --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d86dfa44bb73ca13d774185159e069a4d7aa6585af4a569f9f7569b9460b3392 +size 1227340570 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 new file mode 100644 index 0000000000000000000000000000000000000000..57d53d046386f84a55550d8bb651e1e608b39d4c --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2144a693e034a9dab2595f94cfaf7f29f9a92384fdb8827e7556ecd3b8d42f53 +size 1680405616 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/d891aac544c185edcc57e57a7f69561c b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/d891aac544c185edcc57e57a7f69561c new file mode 100644 index 0000000000000000000000000000000000000000..2a0b58f2c74a2f8b9539cde943dce7b6fad34243 Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/d891aac544c185edcc57e57a7f69561c differ diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 new file mode 100644 index 0000000000000000000000000000000000000000..126dffad8b23049ab59381236b4fe8259cf3ef1a --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:af709070d5c1a4b648d7799923f503c658d2403dfffed9a54fc301a89d18a1f9 +size 2240068387 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 new file mode 100644 index 0000000000000000000000000000000000000000..a92395dff6dcdfb3eb3be9c27515f4a00f4107e3 --- /dev/null +++ b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3614c36b08b05532184f2b34cd514d9d2dcc3f0be7038e846356371d62e52e54 +size 2517086701 diff --git a/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..2d77683f5468fb8cdd93e9db50f808e3b24a7ae1 Binary files /dev/null and b/simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/manifest.ocdbt differ