Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,29 @@ 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
+ckpt-25000/params/ocdbt.process_0/d/b83f5e590fde16d2ffcdcd6e831c1afa filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/b3aa0e59038363c810430bcd75926994 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/f14c5cb70c664815cf05e4b8902bf828 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/fef0dc9c901358f5ea9051ea10470d6a filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/844a9f6ed46ed50d43d01eda1e1b79e3 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/a504e8af8a397251f252ea602774d558 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/814ea7572de2bf1d46b47e1f7e766f3b filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/642cfa8dc1baeeb752b25846dd7adccc filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/b7b32cf0341d6cd18f1cbedb38e302fe filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/cba6f2686dfdc394f5bf01f73d73b3d9 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/84608024d12e3db50a429e221c114f0f filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/e33291732b948f87f2326174174a0a18 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/383f471a0dc0bbb778b32cede406bdf4 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/2f113304c7a6aa52b440b7789543990f filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/3fb585923941d5ebbeda9609f2414820 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/b686b6f5922b31086c318788cc36c207 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/9a5e633956844d5ae85b3b8b137e0eef filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/1d9cd396aea3a0cf32deb4ba7415adc6 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/77cc68f1580fb6fe84d8afc9f53fc590 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/3756e21310bc007794d3260e4315ce01 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/9a4b698a6474ff2170b0041bd0278bde filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/c404c61702c03623a07974424801c937 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/cee0e52b99a22c65aaacaf3d2e7ce8b6 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/6e5037f59b74d62fdc7499e356b644e8 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/f26316144892aa02c7dc5e289d6870b1 filter=lfs diff=lfs merge=lfs -text
+ckpt-25000/params/ocdbt.process_0/d/8d499e6b9adad4a17c6168d3003a9248 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,75 @@
+# pi05-robodyna-baseline-bs256
+
+π0.5 fine-tune on the **RoboDyna** benchmark (dual UR5 + WSG grippers, 35 tasks).
+
+## Checkpoint
+
+| Field | Value |
+|---|---|
+| Base | `openpi-assets/checkpoints/pi05_base/params` (π0.5) |
+| Training config | `pi05_robodyna` (openpi TrainConfig, see `training_config/train_config.py`) |
+| Exp name | `baseline_bs256_30k` |
+| Ckpt uploaded | `ckpt-25000` (step 25000 of a planned 30000) |
+| Loss at step 25000 | ~0.0046 (grad_norm ~0.05, param_norm ~1806) |
+| Why not 30k | Slurm 12h time limit hit at step ~25.8k; last saved ckpt is 25000. |
+
+## Training setup (as-run)
+
+Sbatch launched training with these overrides on top of the openpi TrainConfig:
+
+```
+python scripts/train.py pi05_robodyna \
+    --exp-name baseline_bs256_30k \
+    --overwrite \
+    --fsdp-devices 8 \
+    --batch-size 256 \
+    --num-train-steps 30000 \
+    --log-interval 100 \
+    --save-interval 5000 \
+    --keep-period 10000
+```
+
+Config-file defaults (see `training_config/train_config.py`):
+- Model: `Pi0Config(pi05=True, action_horizon=50)`
+- Data: `LeRobotAlohaDataConfig(repo_id="robodyna", adapt_to_pi=False, use_delta_joint_actions=True)`
+- Repack: images = head → cam_high, left_wrist → cam_left_wrist, right_wrist → cam_right_wrist
+- LR: CosineDecay, warmup=1000, peak=2.5e-5, decay=2.5e-6 over 30000 steps
+- Optimizer: AdamW, clip_gradient_norm=1.0, ema_decay=0.999
+- **BS 256** at launch (overrides the file's default BS 32) on 8× H200, FSDP
+
+Hardware / env: 1× gpu-h200-106, 8× H200, XLA mem 0.95, NCCL_NVLS_ENABLE=0.
+
+## Dataset (RoboDyna)
+
+| Field | Value |
+|---|---|
+| Robot | dual UR5 + WSG grippers (14-D state/action) |
+| Total episodes | 4050 |
+| Total frames | 1,374,883 |
+| Total tasks | 35 |
+| FPS | ~16.67 |
+| Cameras (3) | head, left_wrist, right_wrist |
+| LeRobot version | v2.1 |
+| Local path | `/work/markhsp/datasets/robodyna` (7.5 GB) |
+| Build script | `training_config/build_robodyna_lerobot.py` |
+| norm_stats | Baked into `ckpt-25000/assets/robodyna/norm_stats.json` |
+
+## Layout
+
+```
+ckpt-25000/
+├── _CHECKPOINT_METADATA
+├── assets/
+│   └── robodyna/
+│       └── norm_stats.json      # normalization stats (baked in)
+└── params/                       # π0.5 weights (Orbax checkpoint)
+training_config/
+├── train_config.py              # excerpt of openpi TrainConfig for pi05_robodyna
+├── train_30k.sbatch             # slurm launch script (as-run)
+├── pi05_robodyna_config.patch   # adds the TrainConfig to openpi
+├── pi05_robodyna_workers16.patch  # num_workers=16 adjustment
+├── build_robodyna_lerobot.py    # dataset builder (source→LeRobot v2.1)
+└── compute_norm_stats_fast.py   # norm_stats generator
+```
+
+`train_state/` (~31 GB Orbax optimizer/rng state) is intentionally NOT included. To resume, re-train from the base and load `params/`.

ckpt-25000/_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": 1787193487883803906, "commit_timestamp_nsecs": 1787193572757794238, "custom_metadata": {}}

ckpt-25000/assets/robodyna/norm_stats.json

@@ -0,0 +1,136 @@
+{
+  "norm_stats": {
+    "state": {
+      "mean": [
+        -1.4256925582885742,
+        -1.6636379957199097,
+        -1.507103681564331,
+        -1.3564404249191284,
+        1.4291787147521973,
+        0.025553449988365173,
+        0.6001958847045898,
+        -1.2182224988937378,
+        -1.618809700012207,
+        -1.508061170578003,
+        -1.3469772338867188,
+        1.3814505338668823,
+        0.15097208321094513,
+        0.632867157459259
+      ],
+      "std": [
+        0.5894335508346558,
+        0.3616091012954712,
+        0.4077964127063751,
+        0.36711007356643677,
+        0.659265398979187,
+        0.35710880160331726,
+        0.4553719460964203,
+        0.7272299528121948,
+        0.34507909417152405,
+        0.4401741027832031,
+        0.40912431478500366,
+        0.7529579997062683,
+        0.44152185320854187,
+        0.45681440830230713
+      ],
+      "q01": [
+        -2.087451553249359,
+        -2.3722099697589876,
+        -2.1884448766708373,
+        -1.678600004863739,
+        -1.5752634632110596,
+        -0.6940703403472899,
+        0.0,
+        -1.551709079504013,
+        -2.2875228883266447,
+        -2.197360730791092,
+        -1.6594589138031006,
+        -1.5861233316421508,
+        -1.1734905878067017,
+        0.0
+      ],
+      "q99": [
+        1.3530198209285738,
+        -0.18097919769287074,
+        0.8057692181110383,
+        0.055999876976013274,
+        1.7439356647491455,
+        1.6303220832824712,
+        0.9998,
+        1.7716080257892606,
+        -0.4247041611194611,
+        0.9245013095140457,
+        0.44391815614700336,
+        1.6014086872100828,
+        2.0506450589179988,
+        0.9998
+      ]
+    },
+    "actions": {
+      "mean": [
+        0.010506438091397285,
+        0.0008242920157499611,
+        0.002111168345436454,
+        0.010532110929489136,
+        -0.011367308907210827,
+        0.0010756880510598421,
+        0.5846506357192993,
+        0.015274097211658955,
+        -0.00027031003264710307,
+        0.00371064362116158,
+        0.010530831292271614,
+        -0.010845364071428776,
+        0.001888549653813243,
+        0.616425096988678
+      ],
+      "std": [
+        0.2341788113117218,
+        0.17944715917110443,
+        0.18930457532405853,
+        0.20791380107402802,
+        0.21762460470199585,
+        0.16937041282653809,
+        0.4587770700454712,
+        0.22700554132461548,
+        0.15701399743556976,
+        0.1718132048845291,
+        0.1901271790266037,
+        0.23562085628509521,
+        0.18265898525714874,
+        0.4619460701942444
+      ],
+      "q01": [
+        -0.451742915725708,
+        -0.5705871927261352,
+        -0.5596735925674441,
+        -0.564864376306534,
+        -0.1737762044906619,
+        -0.42154834899902305,
+        0.0,
+        -0.43785261926651,
+        -0.5650164951324463,
+        -0.5262246330261231,
+        -0.5235988674163818,
+        -0.17716022343635585,
+        -0.46765870952606203,
+        0.0
+      ],
+      "q99": [
+        0.44777951803207383,
+        0.5190583933830264,
+        0.6351090844154355,
+        0.6399086740016937,
+        0.09972469253539984,
+        0.3921274749755863,
+        0.9998,
+        0.5755057651519775,
+        0.5371956399917601,
+        0.5852298011779782,
+        0.6093252430915834,
+        0.07817415833473218,
+        0.4746520328521733,
+        0.9998
+      ]
+    }
+  }
+}

ckpt-25000/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": [144]}}, "('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": [144]}}, "('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, 144]}}, "('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, 144]}}, "('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, 144]}}, "('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, 144]}}, "('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, 538]}}, "('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, 538]}}, "('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, 144]}}, "('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, 538, 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, 2, 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, 144, 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, 144]}}, "('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, 144]}}, "('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, 2, 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, 144, 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, 2, 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, 144, 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": [144]}}, "('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, 144]}}, "('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": [256]}}, "('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, 256]}}, "('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, 32, 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": [32144, 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": [256]}}, "('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": [384]}}, "('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, 384]}}, "('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, 256]}}, "('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, 128]}}, "('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, 256, 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, 128, 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, 256, 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, 128, 256]}}, "('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, 2048]}}, "('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, 2048, 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, 512]}}, "('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, 512, 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, 256]}}, "('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, 384]}}, "('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, 384]}}, "('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, 256]}}, "('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, 384]}}, "('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, 384]}}, "('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": [128]}}, "('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": [4, 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": [4]}}, "('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": [128, 32]}}, "('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": [128]}}, "('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, 128]}}, "('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": [128]}}, "('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, 128]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

ckpt-25000/params/_sharding

@@ -0,0 +1 @@
+{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 8], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}, {\"id\": 4}, {\"id\": 5}, {\"id\": 6}, {\"id\": 7}]]}}"}

ckpt-25000/params/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [144], "chunk_shape": [144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [144], "chunk_shape": [144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 144], "chunk_shape": [27, 144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 144], "chunk_shape": [27, 144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 144], "chunk_shape": [27, 144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 144], "chunk_shape": [27, 144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 538], "chunk_shape": [27, 538], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 1152, 538], "chunk_shape": [27, 1152, 538], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 144], "chunk_shape": [27, 144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 538, 1152], "chunk_shape": [27, 538, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 2, 72], "chunk_shape": [27, 2, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 144, 16, 72], "chunk_shape": [27, 144, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 144], "chunk_shape": [27, 144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 16, 72, 144], "chunk_shape": [27, 16, 72, 144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 2, 72], "chunk_shape": [27, 2, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 144, 16, 72], "chunk_shape": [27, 144, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 2, 72], "chunk_shape": [27, 2, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 144, 16, 72], "chunk_shape": [27, 144, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [144], "chunk_shape": [144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 144], "chunk_shape": [14, 14, 3, 144], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [1152, 256], "chunk_shape": [1152, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 32, 1152], "chunk_shape": [1, 32, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [32144, 2048], "chunk_shape": [32144, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [384], "chunk_shape": [384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [1024, 384], "chunk_shape": [1024, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.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.attn_vec_einsum_1.w.value", "write_shape": [18, 8, 256, 128], "chunk_shape": [18, 8, 256, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.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.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 128, 256], "chunk_shape": [18, 2, 1, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.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.q_einsum_1.w.value", "write_shape": [18, 8, 128, 256], "chunk_shape": [18, 8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 2048, 2048], "chunk_shape": [18, 2, 2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 2048, 2048], "chunk_shape": [18, 2048, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 1024, 512], "chunk_shape": [18, 2, 1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 512, 1024], "chunk_shape": [18, 512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [18, 384], "chunk_shape": [18, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 384], "chunk_shape": [18, 1024, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [18, 384], "chunk_shape": [18, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 384], "chunk_shape": [18, 1024, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [4, 1024], "chunk_shape": [4, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [4], "chunk_shape": [4], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [128, 32], "chunk_shape": [128, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [1024, 128], "chunk_shape": [1024, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [1024, 128], "chunk_shape": [1024, 128], "ext_metadata": null}}]}

ckpt-25000/params/ocdbt.process_0/d/1d9cd396aea3a0cf32deb4ba7415adc6

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

ckpt-25000/params/ocdbt.process_0/d/2f113304c7a6aa52b440b7789543990f

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

ckpt-25000/params/ocdbt.process_0/d/3756e21310bc007794d3260e4315ce01

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

ckpt-25000/params/ocdbt.process_0/d/383f471a0dc0bbb778b32cede406bdf4

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

ckpt-25000/params/ocdbt.process_0/d/3fb585923941d5ebbeda9609f2414820

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

ckpt-25000/params/ocdbt.process_0/d/642cfa8dc1baeeb752b25846dd7adccc

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

ckpt-25000/params/ocdbt.process_0/d/6e5037f59b74d62fdc7499e356b644e8

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

ckpt-25000/params/ocdbt.process_0/d/77cc68f1580fb6fe84d8afc9f53fc590

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

ckpt-25000/params/ocdbt.process_0/d/814ea7572de2bf1d46b47e1f7e766f3b

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

ckpt-25000/params/ocdbt.process_0/d/844a9f6ed46ed50d43d01eda1e1b79e3

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

ckpt-25000/params/ocdbt.process_0/d/84608024d12e3db50a429e221c114f0f

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

ckpt-25000/params/ocdbt.process_0/d/8d499e6b9adad4a17c6168d3003a9248

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

ckpt-25000/params/ocdbt.process_0/d/9a4b698a6474ff2170b0041bd0278bde

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

ckpt-25000/params/ocdbt.process_0/d/9a5e633956844d5ae85b3b8b137e0eef

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

ckpt-25000/params/ocdbt.process_0/d/a504e8af8a397251f252ea602774d558

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

ckpt-25000/params/ocdbt.process_0/d/b3aa0e59038363c810430bcd75926994

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

ckpt-25000/params/ocdbt.process_0/d/b686b6f5922b31086c318788cc36c207

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

ckpt-25000/params/ocdbt.process_0/d/b7b32cf0341d6cd18f1cbedb38e302fe

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

ckpt-25000/params/ocdbt.process_0/d/b83f5e590fde16d2ffcdcd6e831c1afa

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

ckpt-25000/params/ocdbt.process_0/d/c404c61702c03623a07974424801c937

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

ckpt-25000/params/ocdbt.process_0/d/cba6f2686dfdc394f5bf01f73d73b3d9

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

ckpt-25000/params/ocdbt.process_0/d/cee0e52b99a22c65aaacaf3d2e7ce8b6

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

ckpt-25000/params/ocdbt.process_0/d/e33291732b948f87f2326174174a0a18

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

ckpt-25000/params/ocdbt.process_0/d/f14c5cb70c664815cf05e4b8902bf828

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

ckpt-25000/params/ocdbt.process_0/d/f26316144892aa02c7dc5e289d6870b1

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

ckpt-25000/params/ocdbt.process_0/d/fef0dc9c901358f5ea9051ea10470d6a

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

training_config/build_robodyna_lerobot.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+"""Merge RoboDyna's 81 LeRobot v2.1 datasets into one training dataset.
+
+The source revision reuses episode indices across conditions, so this builder
+assigns contiguous episode indices, rewrites the affected parquet columns and
+metadata, and hard-links videos. Only episodes referenced by episodes.jsonl are
+included; stale files in older snapshots are ignored.
+"""
+
+from __future__ import annotations
+
+import argparse
+import errno
+import json
+import math
+import os
+from pathlib import Path
+import shutil
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+
+DEFAULT_SOURCE = Path("/work/markhsp/datasets/.robodyna_staging_645afe3/lerobot")
+DEFAULT_OUTPUT = Path("/work/markhsp/datasets/robodyna")
+DEFAULT_REVISION = "0dd81b41aa4156731aa0ee09df8f50a7ff724f08"
+CAMERAS = (
+    "observation.images.head",
+    "observation.images.left_wrist",
+    "observation.images.right_wrist",
+)
+
+
+def read_jsonl(path: Path) -> list[dict]:
+    with path.open(encoding="utf-8") as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def write_jsonl(path: Path, rows: list[dict]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        for row in rows:
+            f.write(json.dumps(row, separators=(",", ":"), ensure_ascii=False) + "\n")
+
+
+def replace_int_column(table: pa.Table, name: str, values: range | list[int]) -> pa.Table:
+    column_index = table.schema.get_field_index(name)
+    if column_index < 0:
+        raise KeyError(f"Required parquet column is missing: {name}")
+    field = table.schema.field(column_index)
+    return table.set_column(column_index, field, pa.array(values, type=field.type))
+
+
+def constant_stats(value: int, count: int) -> dict:
+    return {
+        "min": [value],
+        "max": [value],
+        "mean": [float(value)],
+        "std": [0.0],
+        "count": [count],
+    }
+
+
+def consecutive_stats(start: int, count: int) -> dict:
+    end = start + count - 1
+    return {
+        "min": [start],
+        "max": [end],
+        "mean": [start + (count - 1) / 2.0],
+        "std": [math.sqrt((count * count - 1) / 12.0)] if count > 1 else [0.0],
+        "count": [count],
+    }
+
+
+def link_or_copy(source: Path, destination: Path) -> None:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        os.link(source, destination)
+    except OSError as exc:
+        if exc.errno != errno.EXDEV:
+            raise
+        shutil.copy2(source, destination)
+
+
+def build(source: Path, output: Path, revision: str) -> None:
+    if output.exists():
+        raise FileExistsError(f"Refusing to overwrite existing output: {output}")
+    building = output.with_name(f".{output.name}.building")
+    if building.exists():
+        raise FileExistsError(f"Remove or inspect previous partial build first: {building}")
+
+    metadata_files = sorted(source.glob("*/*/meta/episodes.jsonl"))
+    if len(metadata_files) != 81:
+        raise RuntimeError(f"Expected 81 task-condition datasets, found {len(metadata_files)}")
+
+    building.mkdir(parents=True)
+    episode_rows: list[dict] = []
+    stats_rows: list[dict] = []
+    manifest_rows: list[dict] = []
+    prompt_to_task_index: dict[str, int] = {}
+    total_frames = 0
+
+    reference_info = json.loads((metadata_files[0].parent / "info.json").read_text())
+    expected_features = reference_info["features"]
+
+    for metadata_path in metadata_files:
+        dataset_root = metadata_path.parent.parent
+        combo = "/".join(dataset_root.parts[-2:])
+        info = json.loads((dataset_root / "meta/info.json").read_text())
+        if info["features"] != expected_features:
+            raise RuntimeError(f"Feature schema differs in {combo}")
+
+        source_episodes = read_jsonl(metadata_path)
+        source_stats = {
+            row["episode_index"]: row for row in read_jsonl(dataset_root / "meta/episodes_stats.jsonl")
+        }
+        if len(source_episodes) != 50:
+            raise RuntimeError(f"Expected 50 episodes in {combo}, found {len(source_episodes)}")
+
+        for episode in sorted(source_episodes, key=lambda row: row["episode_index"]):
+            old_index = int(episode["episode_index"])
+            new_index = len(episode_rows)
+            length = int(episode["length"])
+            prompt = episode["tasks"][0]
+            task_index = prompt_to_task_index.setdefault(prompt, len(prompt_to_task_index))
+            old_chunk = old_index // int(info["chunks_size"])
+            new_chunk = new_index // 1000
+
+            source_parquet = dataset_root / info["data_path"].format(
+                episode_chunk=old_chunk, episode_index=old_index
+            )
+            destination_parquet = building / (
+                f"data/task-{new_chunk:04d}/episode_{new_index:08d}.parquet"
+            )
+            table = pq.read_table(source_parquet)
+            if table.num_rows != length:
+                raise RuntimeError(
+                    f"Length mismatch in {combo}/{old_index}: metadata={length}, parquet={table.num_rows}"
+                )
+            table = replace_int_column(table, "episode_index", [new_index] * length)
+            table = replace_int_column(table, "index", range(total_frames, total_frames + length))
+            table = replace_int_column(table, "task_index", [task_index] * length)
+            destination_parquet.parent.mkdir(parents=True, exist_ok=True)
+            pq.write_table(table, destination_parquet, compression="zstd")
+
+            for camera in CAMERAS:
+                source_video = dataset_root / info["video_path"].format(
+                    episode_chunk=old_chunk,
+                    episode_index=old_index,
+                    video_key=camera,
+                )
+                destination_video = building / (
+                    f"videos/task-{new_chunk:04d}/{camera}/episode_{new_index:08d}.mp4"
+                )
+                if not source_video.is_file():
+                    raise FileNotFoundError(source_video)
+                link_or_copy(source_video, destination_video)
+
+            annotation_template = info.get(
+                "annotation_path",
+                "annotations/task-{episode_chunk:04d}/episode_{episode_index:08d}.json",
+            )
+            source_annotation = dataset_root / annotation_template.format(
+                episode_chunk=old_chunk, episode_index=old_index
+            )
+            annotation = json.loads(source_annotation.read_text())
+            annotation["episode_index"] = new_index
+            if isinstance(annotation.get("meta_data"), dict):
+                annotation["meta_data"]["length"] = length
+            destination_annotation = building / (
+                f"annotations/task-{new_chunk:04d}/episode_{new_index:08d}.json"
+            )
+            destination_annotation.parent.mkdir(parents=True, exist_ok=True)
+            destination_annotation.write_text(
+                json.dumps(annotation, separators=(",", ":"), ensure_ascii=False) + "\n",
+                encoding="utf-8",
+            )
+
+            output_episode = dict(episode)
+            output_episode["episode_index"] = new_index
+            episode_rows.append(output_episode)
+
+            output_stats = json.loads(json.dumps(source_stats[old_index]))
+            output_stats["episode_index"] = new_index
+            stats = output_stats["stats"]
+            stats["episode_index"] = constant_stats(new_index, length)
+            stats["index"] = consecutive_stats(total_frames, length)
+            stats["task_index"] = constant_stats(task_index, length)
+            stats_rows.append(output_stats)
+
+            manifest_rows.append(
+                {
+                    "new_episode_index": new_index,
+                    "source_combo": combo,
+                    "source_episode_index": old_index,
+                    "task_index": task_index,
+                    "length": length,
+                }
+            )
+            total_frames += length
+
+    if len(episode_rows) != 4050:
+        raise RuntimeError(f"Expected 4050 episodes, built {len(episode_rows)}")
+
+    task_rows = [
+        {"task_index": task_index, "task": prompt}
+        for prompt, task_index in sorted(prompt_to_task_index.items(), key=lambda item: item[1])
+    ]
+    info = dict(reference_info)
+    info.update(
+        {
+            "total_episodes": len(episode_rows),
+            "total_frames": total_frames,
+            "total_tasks": len(task_rows),
+            "total_videos": len(episode_rows) * len(CAMERAS),
+            "total_chunks": math.ceil(len(episode_rows) / 1000),
+            "chunks_size": 1000,
+            "splits": {"train": f"0:{len(episode_rows)}"},
+            "data_path": "data/task-{episode_chunk:04d}/episode_{episode_index:08d}.parquet",
+            "video_path": "videos/task-{episode_chunk:04d}/{video_key}/episode_{episode_index:08d}.mp4",
+            "annotation_path": "annotations/task-{episode_chunk:04d}/episode_{episode_index:08d}.json",
+        }
+    )
+    meta = building / "meta"
+    meta.mkdir()
+    (meta / "info.json").write_text(json.dumps(info, indent=2) + "\n", encoding="utf-8")
+    write_jsonl(meta / "tasks.jsonl", task_rows)
+    write_jsonl(meta / "episodes.jsonl", episode_rows)
+    write_jsonl(meta / "episodes_stats.jsonl", stats_rows)
+    (meta / "source_manifest.json").write_text(
+        json.dumps(
+            {
+                "dataset": "RoboDyna/robodyna-benchmark-v2",
+                "revision": revision,
+                "source_root": str(source),
+                "episodes": manifest_rows,
+            },
+            indent=2,
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+
+    expected_counts = {
+        "parquet": len(list(building.glob("data/task-*/episode_*.parquet"))),
+        "annotations": len(list(building.glob("annotations/task-*/episode_*.json"))),
+        "videos": len(list(building.glob("videos/task-*/*/episode_*.mp4"))),
+    }
+    if expected_counts != {"parquet": 4050, "annotations": 4050, "videos": 12150}:
+        raise RuntimeError(f"Output count mismatch: {expected_counts}")
+
+    building.rename(output)
+    print(
+        json.dumps(
+            {
+                "output": str(output),
+                "episodes": len(episode_rows),
+                "frames": total_frames,
+                "tasks": len(task_rows),
+                **expected_counts,
+            },
+            indent=2,
+        )
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE)
+    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
+    parser.add_argument("--revision", default=DEFAULT_REVISION)
+    args = parser.parse_args()
+    build(args.source, args.output, args.revision)
+
+
+if __name__ == "__main__":
+    main()

training_config/compute_norm_stats_fast.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""Compute exact frame coverage norm stats without decoding RoboDyna videos."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as pq
+
+from openpi.shared import normalize
+
+
+DELTA_MASK = np.array(
+    [True, True, True, True, True, True, False, True, True, True, True, True, True, False]
+)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=Path, default=Path("/work/markhsp/datasets/robodyna"))
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=Path("/work/markhsp/openpi/assets/pi05_robodyna/robodyna"),
+    )
+    parser.add_argument("--action-horizon", type=int, default=50)
+    args = parser.parse_args()
+
+    episodes = [
+        json.loads(line)
+        for line in (args.dataset / "meta/episodes.jsonl").read_text().splitlines()
+        if line.strip()
+    ]
+    state_stats = normalize.RunningStats()
+    action_stats = normalize.RunningStats()
+    frame_count = 0
+
+    for position, episode in enumerate(episodes, start=1):
+        episode_index = int(episode["episode_index"])
+        episode_chunk = episode_index // 1000
+        parquet_path = (
+            args.dataset
+            / f"data/task-{episode_chunk:04d}/episode_{episode_index:08d}.parquet"
+        )
+        table = pq.read_table(parquet_path, columns=["observation.state", "action"])
+        state = np.asarray(table["observation.state"].to_pylist(), dtype=np.float32)
+        action = np.asarray(table["action"].to_pylist(), dtype=np.float32)
+        if state.shape != action.shape or state.shape[1] != 14:
+            raise RuntimeError(
+                f"Unexpected state/action shapes for episode {episode_index}: "
+                f"{state.shape}, {action.shape}"
+            )
+
+        length = state.shape[0]
+        offsets = np.arange(args.action_horizon, dtype=np.int64)
+        indices = np.minimum(np.arange(length, dtype=np.int64)[:, None] + offsets, length - 1)
+        action_chunks = action[indices].copy()
+        action_chunks[:, :, DELTA_MASK] -= state[:, None, DELTA_MASK]
+
+        state_stats.update(state)
+        action_stats.update(action_chunks)
+        frame_count += length
+        if position % 100 == 0 or position == len(episodes):
+            print(f"episodes={position}/{len(episodes)} frames={frame_count}", flush=True)
+
+    expected_frames = json.loads((args.dataset / "meta/info.json").read_text())["total_frames"]
+    if frame_count != expected_frames:
+        raise RuntimeError(f"Frame count mismatch: processed={frame_count}, expected={expected_frames}")
+
+    norm_stats = {
+        "state": state_stats.get_statistics(),
+        "actions": action_stats.get_statistics(),
+    }
+    normalize.save(args.output, norm_stats)
+    print(f"wrote {args.output / 'norm_stats.json'}")
+    for name, stats in norm_stats.items():
+        print(name, "mean", stats.mean.tolist())
+        print(name, "std", stats.std.tolist())
+        print(name, "q01", stats.q01.tolist())
+        print(name, "q99", stats.q99.tolist())
+
+
+if __name__ == "__main__":
+    main()

training_config/pi05_robodyna_config.patch

@@ -0,0 +1,52 @@
+diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py
+--- a/src/openpi/training/config.py
++++ b/src/openpi/training/config.py
+@@ -1853,5 +1853,47 @@ for _tid in range(0, 100):
+         num_train_steps=50_000,
+     ))
+ 
++# RoboDyna benchmark v2: dual UR5 + WSG, three RGB cameras, 14-D state/action.
++# The dataset is built locally by /work/markhsp/openpi_robodyna/
++# build_robodyna_lerobot.py. HF_LEROBOT_HOME must be /work/markhsp/datasets.
++_CONFIGS.append(TrainConfig(
++    name="pi05_robodyna",
++    wandb_enabled=False,
++    num_workers=48,
++    model=pi0_config.Pi0Config(pi05=True, action_horizon=50),
++    data=LeRobotAlohaDataConfig(
++        repo_id="robodyna",
++        adapt_to_pi=False,
++        use_delta_joint_actions=True,
++        repack_transforms=_transforms.Group(
++            inputs=[
++                _transforms.RepackTransform(
++                    {
++                        "images": {
++                            "cam_high": "observation.images.head",
++                            "cam_left_wrist": "observation.images.left_wrist",
++                            "cam_right_wrist": "observation.images.right_wrist",
++                        },
++                        "state": "observation.state",
++                        "actions": "action",
++                        "prompt": "prompt",
++                    }
++                )
++            ]
++        ),
++        base_config=DataConfig(prompt_from_task=True),
++    ),
++    batch_size=32,
++    lr_schedule=_optimizer.CosineDecaySchedule(
++        warmup_steps=1_000, peak_lr=2.5e-5, decay_steps=30_000, decay_lr=2.5e-6),
++    optimizer=_optimizer.AdamW(clip_gradient_norm=1.0),
++    ema_decay=0.999,
++    weight_loader=weight_loaders.CheckpointWeightLoader(
++        "gs://openpi-assets/checkpoints/pi05_base/params"),
++    num_train_steps=30_000,
++    save_interval=5_000,
++    keep_period=10_000,
++))
++
+ if len({config.name for config in _CONFIGS}) != len(_CONFIGS):
+     raise ValueError("Config names must be unique.")
+ _CONFIGS_DICT = {config.name: config for config in _CONFIGS}

training_config/pi05_robodyna_workers16.patch

@@ -0,0 +1,12 @@
+diff --git a/src/openpi/training/config.py b/src/openpi/training/config.py
+--- a/src/openpi/training/config.py
++++ b/src/openpi/training/config.py
+@@ -1860,6 +1860,6 @@ _CONFIGS.append(TrainConfig(
+ _CONFIGS.append(TrainConfig(
+     name="pi05_robodyna",
+     wandb_enabled=False,
+-    num_workers=48,
++    num_workers=16,
+     model=pi0_config.Pi0Config(pi05=True, action_horizon=50),
+     data=LeRobotAlohaDataConfig(
+         repo_id="robodyna",

training_config/train_30k.sbatch

@@ -0,0 +1,31 @@
+#!/bin/bash
+#SBATCH --job-name=pi05-robodyna
+#SBATCH --partition=gpu
+#SBATCH --nodelist=gpu-h200-106
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --gres=gpu:h200:8
+#SBATCH --time=12:00:00
+#SBATCH --output=/work/markhsp/openpi_robodyna/logs/train_30k_%j.out
+#SBATCH --error=/work/markhsp/openpi_robodyna/logs/train_30k_%j.err
+
+set -euo pipefail
+cd /work/markhsp/openpi
+
+export HF_LEROBOT_HOME=/work/markhsp/datasets
+export LD_LIBRARY_PATH=/work/markhsp/miniforge3/envs/ffmpeg7/lib:${LD_LIBRARY_PATH:-}
+export XLA_PYTHON_CLIENT_MEM_FRACTION=0.95
+export JAX_PLATFORMS=cuda
+export NCCL_P2P_DISABLE=0
+export NCCL_IB_DISABLE=0
+export NCCL_NVLS_ENABLE=0
+
+exec ./.venv/bin/python scripts/train.py pi05_robodyna \
+    --exp-name baseline_bs256_30k \
+    --overwrite \
+    --fsdp-devices 8 \
+    --batch-size 256 \
+    --num-train-steps 30000 \
+    --log-interval 100 \
+    --save-interval 5000 \
+    --keep-period 10000