diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..b32d1cd8174d7d7136a2dfaf32da9a9fdb03e617 100644 --- a/.gitattributes +++ b/.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 diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..44604f7a4877c063ad19159b056e67a80164b4e1 --- /dev/null +++ b/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/`. diff --git a/ckpt-25000/_CHECKPOINT_METADATA b/ckpt-25000/_CHECKPOINT_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..920be7225c49f6985c97ab43718aaed4eaff40e0 --- /dev/null +++ b/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": {}} \ No newline at end of file diff --git a/ckpt-25000/assets/robodyna/norm_stats.json b/ckpt-25000/assets/robodyna/norm_stats.json new file mode 100644 index 0000000000000000000000000000000000000000..4adf51e58e9d1855d3aaa71c20114f9a1a7dabc8 --- /dev/null +++ b/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 + ] + } + } +} \ No newline at end of file diff --git a/ckpt-25000/params/_METADATA b/ckpt-25000/params/_METADATA new file mode 100644 index 0000000000000000000000000000000000000000..4abbf23f798c447e399397ac5525b74fb19e3df6 --- /dev/null +++ b/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} \ No newline at end of file diff --git a/ckpt-25000/params/_sharding b/ckpt-25000/params/_sharding new file mode 100644 index 0000000000000000000000000000000000000000..9f1e9754b674dbc417074caf482599cddc003da2 --- /dev/null +++ b/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}]]}}"} \ No newline at end of file diff --git a/ckpt-25000/params/array_metadatas/process_0 b/ckpt-25000/params/array_metadatas/process_0 new file mode 100644 index 0000000000000000000000000000000000000000..126972f76293b18a1a491d31c0182e3c428d6945 --- /dev/null +++ b/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}}]} \ No newline at end of file diff --git a/ckpt-25000/params/d/ef69a052365d6b7369b666efc81d52fd b/ckpt-25000/params/d/ef69a052365d6b7369b666efc81d52fd new file mode 100644 index 0000000000000000000000000000000000000000..26edde26c3eb00822d22fc70414e3dda82816719 Binary files /dev/null and b/ckpt-25000/params/d/ef69a052365d6b7369b666efc81d52fd differ diff --git a/ckpt-25000/params/manifest.ocdbt b/ckpt-25000/params/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..67cecd2bf21e50c2388bd17759a91ffca069e924 Binary files /dev/null and b/ckpt-25000/params/manifest.ocdbt differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/0cfa7eeae7aef76ad5e96fe05f98a323 b/ckpt-25000/params/ocdbt.process_0/d/0cfa7eeae7aef76ad5e96fe05f98a323 new file mode 100644 index 0000000000000000000000000000000000000000..6b634cb0637d2fd1aca44126edecc5c2e675a4da Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/d/0cfa7eeae7aef76ad5e96fe05f98a323 differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/135cef970fe363af882d5bc49563f37b b/ckpt-25000/params/ocdbt.process_0/d/135cef970fe363af882d5bc49563f37b new file mode 100644 index 0000000000000000000000000000000000000000..05be8712c4712672c91189617db1294d31115faf Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/d/135cef970fe363af882d5bc49563f37b differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/1d9cd396aea3a0cf32deb4ba7415adc6 b/ckpt-25000/params/ocdbt.process_0/d/1d9cd396aea3a0cf32deb4ba7415adc6 new file mode 100644 index 0000000000000000000000000000000000000000..f72c1130901e453f3388db7fba8db163bf1451a8 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/2f113304c7a6aa52b440b7789543990f b/ckpt-25000/params/ocdbt.process_0/d/2f113304c7a6aa52b440b7789543990f new file mode 100644 index 0000000000000000000000000000000000000000..9ef6c411d1d7344ae50190fa41f483aaba4f85a5 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/3756e21310bc007794d3260e4315ce01 b/ckpt-25000/params/ocdbt.process_0/d/3756e21310bc007794d3260e4315ce01 new file mode 100644 index 0000000000000000000000000000000000000000..48dfdf487334ab7e5adc0f5f889d785352ee312a --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/383f471a0dc0bbb778b32cede406bdf4 b/ckpt-25000/params/ocdbt.process_0/d/383f471a0dc0bbb778b32cede406bdf4 new file mode 100644 index 0000000000000000000000000000000000000000..267b38dd76675fb5ca5a0e9ac89ce236510f1f1f --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/3fb585923941d5ebbeda9609f2414820 b/ckpt-25000/params/ocdbt.process_0/d/3fb585923941d5ebbeda9609f2414820 new file mode 100644 index 0000000000000000000000000000000000000000..241a54c381080d9846206135aa7402e897afb323 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/5761589a4e70b05f29f8cd27df7f89da b/ckpt-25000/params/ocdbt.process_0/d/5761589a4e70b05f29f8cd27df7f89da new file mode 100644 index 0000000000000000000000000000000000000000..75d9146169ca890730b8f27ba8a0bbc247ddb93f Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/d/5761589a4e70b05f29f8cd27df7f89da differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/642cfa8dc1baeeb752b25846dd7adccc b/ckpt-25000/params/ocdbt.process_0/d/642cfa8dc1baeeb752b25846dd7adccc new file mode 100644 index 0000000000000000000000000000000000000000..a55a5f31325146d761d7c27be23dd60435a52f3c --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/6e5037f59b74d62fdc7499e356b644e8 b/ckpt-25000/params/ocdbt.process_0/d/6e5037f59b74d62fdc7499e356b644e8 new file mode 100644 index 0000000000000000000000000000000000000000..fb7973d1f95ca12a6d1eccb77dc29a6d10327930 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/77cc68f1580fb6fe84d8afc9f53fc590 b/ckpt-25000/params/ocdbt.process_0/d/77cc68f1580fb6fe84d8afc9f53fc590 new file mode 100644 index 0000000000000000000000000000000000000000..fb27a19926761509f857bae6e27dc54e46fc49f2 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/8135367ddcd46976691419ee1171a0a5 b/ckpt-25000/params/ocdbt.process_0/d/8135367ddcd46976691419ee1171a0a5 new file mode 100644 index 0000000000000000000000000000000000000000..4273856d71f3c034a3eae03114fcb5865c5f7c14 Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/d/8135367ddcd46976691419ee1171a0a5 differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/814ea7572de2bf1d46b47e1f7e766f3b b/ckpt-25000/params/ocdbt.process_0/d/814ea7572de2bf1d46b47e1f7e766f3b new file mode 100644 index 0000000000000000000000000000000000000000..a36097ccc138ce7a74473fbe21db6099460419a5 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/844a9f6ed46ed50d43d01eda1e1b79e3 b/ckpt-25000/params/ocdbt.process_0/d/844a9f6ed46ed50d43d01eda1e1b79e3 new file mode 100644 index 0000000000000000000000000000000000000000..92103c1c6f033c3eb064cba150ce3efbc2c32259 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/84608024d12e3db50a429e221c114f0f b/ckpt-25000/params/ocdbt.process_0/d/84608024d12e3db50a429e221c114f0f new file mode 100644 index 0000000000000000000000000000000000000000..6c535d96eb4d82490804d6ffe386a7aa5f5efb41 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/8d499e6b9adad4a17c6168d3003a9248 b/ckpt-25000/params/ocdbt.process_0/d/8d499e6b9adad4a17c6168d3003a9248 new file mode 100644 index 0000000000000000000000000000000000000000..4727dbea768fadc5038a1f1d326930823c58295a --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/945a0bb81e976a65bf98a12de75daad1 b/ckpt-25000/params/ocdbt.process_0/d/945a0bb81e976a65bf98a12de75daad1 new file mode 100644 index 0000000000000000000000000000000000000000..283b3a6575fc25699bfa515a3838435430d85683 Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/d/945a0bb81e976a65bf98a12de75daad1 differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/9a4b698a6474ff2170b0041bd0278bde b/ckpt-25000/params/ocdbt.process_0/d/9a4b698a6474ff2170b0041bd0278bde new file mode 100644 index 0000000000000000000000000000000000000000..644b6962a9c98d01f14f1e4417c56e6f33454e9c --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/9a5e633956844d5ae85b3b8b137e0eef b/ckpt-25000/params/ocdbt.process_0/d/9a5e633956844d5ae85b3b8b137e0eef new file mode 100644 index 0000000000000000000000000000000000000000..7713e77b5dbbd9b07323c9f73ca734e21a0aa3ba --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/a504e8af8a397251f252ea602774d558 b/ckpt-25000/params/ocdbt.process_0/d/a504e8af8a397251f252ea602774d558 new file mode 100644 index 0000000000000000000000000000000000000000..5443e47d0a71e1eddf93e6fcb53cece86353a2dc --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/b28b2e01b16dd245900db5f76b3c769d b/ckpt-25000/params/ocdbt.process_0/d/b28b2e01b16dd245900db5f76b3c769d new file mode 100644 index 0000000000000000000000000000000000000000..b4a4f19b626dd94b59f02e20e84bcb37282ac267 Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/d/b28b2e01b16dd245900db5f76b3c769d differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/b3aa0e59038363c810430bcd75926994 b/ckpt-25000/params/ocdbt.process_0/d/b3aa0e59038363c810430bcd75926994 new file mode 100644 index 0000000000000000000000000000000000000000..a4abb205188d297be64bc29eaac83cdd455c85f9 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/b686b6f5922b31086c318788cc36c207 b/ckpt-25000/params/ocdbt.process_0/d/b686b6f5922b31086c318788cc36c207 new file mode 100644 index 0000000000000000000000000000000000000000..10bf83d4dd88cb10381a19d2c7626bc9c7e3a4f6 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/b7b32cf0341d6cd18f1cbedb38e302fe b/ckpt-25000/params/ocdbt.process_0/d/b7b32cf0341d6cd18f1cbedb38e302fe new file mode 100644 index 0000000000000000000000000000000000000000..932515fc30bdee70275af1825f12a3f8500490e2 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/b83f5e590fde16d2ffcdcd6e831c1afa b/ckpt-25000/params/ocdbt.process_0/d/b83f5e590fde16d2ffcdcd6e831c1afa new file mode 100644 index 0000000000000000000000000000000000000000..693e0f8549a3314b64e88bd5fdf45925570eb659 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/c404c61702c03623a07974424801c937 b/ckpt-25000/params/ocdbt.process_0/d/c404c61702c03623a07974424801c937 new file mode 100644 index 0000000000000000000000000000000000000000..fca947b6b545f5479b293a97ea9b5ed8d6b5bb7d --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/cba6f2686dfdc394f5bf01f73d73b3d9 b/ckpt-25000/params/ocdbt.process_0/d/cba6f2686dfdc394f5bf01f73d73b3d9 new file mode 100644 index 0000000000000000000000000000000000000000..a676a21858a7d8a5c93abe1d5118c09b3e5bc64c --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/cee0e52b99a22c65aaacaf3d2e7ce8b6 b/ckpt-25000/params/ocdbt.process_0/d/cee0e52b99a22c65aaacaf3d2e7ce8b6 new file mode 100644 index 0000000000000000000000000000000000000000..e85d95eeed6e347cf41e3634f434a75d60123819 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/d7ec042a1992012f36ed497f98483cb4 b/ckpt-25000/params/ocdbt.process_0/d/d7ec042a1992012f36ed497f98483cb4 new file mode 100644 index 0000000000000000000000000000000000000000..0b0154f3b22e15b716acc7846fce032559ce0469 Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/d/d7ec042a1992012f36ed497f98483cb4 differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/e33291732b948f87f2326174174a0a18 b/ckpt-25000/params/ocdbt.process_0/d/e33291732b948f87f2326174174a0a18 new file mode 100644 index 0000000000000000000000000000000000000000..2f845613058f2cf0aa8f64be55d09967cf76ef29 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/e7c7d296a9a38fbc70379e1de2674538 b/ckpt-25000/params/ocdbt.process_0/d/e7c7d296a9a38fbc70379e1de2674538 new file mode 100644 index 0000000000000000000000000000000000000000..e56984b86a3d1df3197ba577dcaef870851074c8 Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/d/e7c7d296a9a38fbc70379e1de2674538 differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/ee1ebec8f8e30f57ff9b5363921c6c2f b/ckpt-25000/params/ocdbt.process_0/d/ee1ebec8f8e30f57ff9b5363921c6c2f new file mode 100644 index 0000000000000000000000000000000000000000..3a2b234f451cf0616a8ef1b4c77e3580b9f4f34c Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/d/ee1ebec8f8e30f57ff9b5363921c6c2f differ diff --git a/ckpt-25000/params/ocdbt.process_0/d/f14c5cb70c664815cf05e4b8902bf828 b/ckpt-25000/params/ocdbt.process_0/d/f14c5cb70c664815cf05e4b8902bf828 new file mode 100644 index 0000000000000000000000000000000000000000..025c3ebc12debd5ecdd89e16f4230434492f685f --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/f26316144892aa02c7dc5e289d6870b1 b/ckpt-25000/params/ocdbt.process_0/d/f26316144892aa02c7dc5e289d6870b1 new file mode 100644 index 0000000000000000000000000000000000000000..1dd3a80ac6abd84d7f26145a89c742d1712fcc4f --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/d/fef0dc9c901358f5ea9051ea10470d6a b/ckpt-25000/params/ocdbt.process_0/d/fef0dc9c901358f5ea9051ea10470d6a new file mode 100644 index 0000000000000000000000000000000000000000..9357f0adbc1a22b78b4ae993c3301d16e8930908 --- /dev/null +++ b/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 diff --git a/ckpt-25000/params/ocdbt.process_0/manifest.ocdbt b/ckpt-25000/params/ocdbt.process_0/manifest.ocdbt new file mode 100644 index 0000000000000000000000000000000000000000..ae1fefd0233c363d0aa72541718074c53b23ca7a Binary files /dev/null and b/ckpt-25000/params/ocdbt.process_0/manifest.ocdbt differ diff --git a/training_config/build_robodyna_lerobot.py b/training_config/build_robodyna_lerobot.py new file mode 100644 index 0000000000000000000000000000000000000000..9166002efc0ccf023dec7d26edc9e0413d3e4458 --- /dev/null +++ b/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() diff --git a/training_config/compute_norm_stats_fast.py b/training_config/compute_norm_stats_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..06e858a36cb4abfec377bcc2e2a0c3efe3ad9c53 --- /dev/null +++ b/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() diff --git a/training_config/pi05_robodyna_config.patch b/training_config/pi05_robodyna_config.patch new file mode 100644 index 0000000000000000000000000000000000000000..4a3dec34351370bb8b8bd0b4872c1699bb837f04 --- /dev/null +++ b/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} diff --git a/training_config/pi05_robodyna_workers16.patch b/training_config/pi05_robodyna_workers16.patch new file mode 100644 index 0000000000000000000000000000000000000000..35dabe541ba7dc01098a7ea6b7dcdbf01777c045 --- /dev/null +++ b/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", diff --git a/training_config/train_30k.sbatch b/training_config/train_30k.sbatch new file mode 100644 index 0000000000000000000000000000000000000000..0b775908088078a90825974ef53266670a835e8b --- /dev/null +++ b/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 diff --git a/training_config/train_config.py b/training_config/train_config.py new file mode 100644 index 0000000000000000000000000000000000000000..5d62cdc7115a0e964fe0a7a1fe8d82e87fd65029 --- /dev/null +++ b/training_config/train_config.py @@ -0,0 +1,39 @@ +# build_robodyna_lerobot.py. HF_LEROBOT_HOME must be /work/markhsp/datasets. +_CONFIGS.append(TrainConfig( + name="pi05_robodyna", + wandb_enabled=False, + num_workers=16, + 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, +))