From cea4cfc89b2b87c99d97ce2c2975c4ab431a18df Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 21:42:20 +0000 Subject: [PATCH 1/2] tests: follow transformers 5's is_offline_mode() _hf_offline monkeypatched transformers.utils.hub._is_offline_mode, a module flag that 5.x replaced with is_offline_mode(). setattr on a name the module no longer has raises, so every test standing on the harness errored out before it reached the network: 12 of them, including the whole PI train-step matrix. Verified failing identically on an unmodified tree, so this is the venv's transformers 5.17 meeting a 4.x assumption, not a regression from this stack. Probe for whichever spelling the installed version has rather than pinning either, since the lockfile still says 4.53.2. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Xh3wMw9pXYPHUSGjU7qmix --- tests/fixtures/train_harness.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/fixtures/train_harness.py b/tests/fixtures/train_harness.py index 6bdb415f8b..025692925f 100644 --- a/tests/fixtures/train_harness.py +++ b/tests/fixtures/train_harness.py @@ -73,7 +73,12 @@ def _hf_offline(monkeypatch) -> None: monkeypatch.setenv("HF_HUB_OFFLINE", "1") monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) - monkeypatch.setattr(transformers.utils.hub, "_is_offline_mode", True) + # transformers <= 4 caches the flag in `_is_offline_mode`; 5.x reads it + # through `is_offline_mode()` instead. Whichever this version has. + if hasattr(transformers.utils.hub, "_is_offline_mode"): + monkeypatch.setattr(transformers.utils.hub, "_is_offline_mode", True) + else: + monkeypatch.setattr(transformers.utils.hub, "is_offline_mode", lambda: True) def write_fixtures( From 3c557ccc471c9eac364bab58636f690ad8fddedc Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 21:42:51 +0000 Subject: [PATCH 2/2] pi: initialize from PaliGemma with a scratch action expert An arm for training pi0.5 from the released VLM instead of the pi0.5 base checkpoint, to separate what the VLM pretraining buys from what pi0.5's own robot pretraining buys. The prefix (SigLIP tower, Gemma-2B, the projector) comes from google/paligemma-3b-pt-224; everything pi0.5 adds on top -- the Gemma-300M action expert, action_in/out_proj, the adaRMS time MLPs -- stays at its init. pytorch_weight_path and paligemma_weight_path are mutually exclusive rather than ordered by precedence: silently preferring one would make a typo in the other look like it worked. Setting neither now warns instead of loading nothing quietly. The loader streams shard by shard and then fails if any parameter went untouched or any checkpoint tensor went unused, because a partial load of a 3.6B model is invisible in the loss for a long time. Two places where the release does not line up with what openpi builds, both handled in the key remap: - transformers 5 flattened SiglipVisionModel, so vision_tower.vision_model.* is vision_tower.* there; below 5 (the pin, 4.53.2) the hop is still in the module tree. The loader maps onto whichever the live model has. Mapping to the 5.x spelling unconditionally failed 21 SigLIP tensors on the pin, which the key-string tests could not see; a unit test now loads a tiny real PaliGemmaForConditionalGeneration under the installed transformers, and passes on both 4.53.2 and 5.17.0. - The checkpoint pads the vocabulary to 257216, 64 rows past the 257152 openpi allocates ( plus 63 unused). embed_tokens and lm_head are truncated; openpi reaches the image slots through the vision tower, never through an embed_tokens lookup, so the dropped rows are unreachable. google/paligemma-3b-pt-224 is license-gated. With paligemma_allow_mirror: true, resolve_paligemma_dir falls back on a 401/403 (not on a rate limit or an outage) to ungated mirrors whose shards are sha256-checked against the canonical repo's LFS metadata first -- a mirror that has drifted is refused, not used. It is off by default: integrity is checked, but a research arm's weights of record should come from the canonical repo unless a run opts in, and this cluster's HF token is authorized for it. Only *.safetensors are pulled. openpi's own loader source is not an option: its GCS bucket 403s anonymously. pi0.5_paligemma_init is an overlay, not a runnable config, so a keypoint or 6D arm is four lines. It changes the init source only; the constant LR stays, since pi0.5_base documents dropping the cosine schedule on purpose. openpi_compat carries three shims for a transformers 5 venv, which openpi's pin (4.53.2, also this repo's) does not expect; all should be deleted once the fork's transformers_replace patch is ported forward. The third restores openpi's fp32 keep-list, which spells the pre-flattening vision paths, so on 5.x the SigLIP patch and position embeddings ran in bf16. On 5.x only this arm loads: the pi0.5 base checkpoint stores the nested SigLIP names and its strict load_model stops at startup. pi05.md says which setup is supported and what differs. Verified with a full trainHydra run to exit 0 on an H100 (4 epochs, all four validation heads, checkpoint and videos written) and the real 12 GB checkpoint into the real model in tests/integration, both on transformers 5.17, plus the unit tests on 4.53.2 and 5.17.0. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Xh3wMw9pXYPHUSGjU7qmix --- README.md | 1 + egomimic/algo/pi.py | 36 +- egomimic/hydra_configs/model/pi0.5_base.yaml | 11 + .../pi0.5_bc_mecka_6d_paligemma_init.yaml | 15 + .../model/pi0.5_paligemma_init.yaml | 30 ++ egomimic/models/openpi_compat.py | 129 +++++ egomimic/models/paligemma_init.py | 296 ++++++++++++ egomimic/robot/rollout.py | 19 +- egomimic/trainHydra.py | 19 +- pi05.md | 117 +++++ tests/integration/test_paligemma_init_real.py | 253 ++++++++++ tests/unit/test_checkpoint_resume.py | 18 + tests/unit/test_openpi_compat.py | 115 +++++ tests/unit/test_paligemma_init.py | 454 ++++++++++++++++++ tests/unit/test_rollout_transform_mode.py | 25 + 15 files changed, 1518 insertions(+), 20 deletions(-) create mode 100644 egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml create mode 100644 egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml create mode 100644 egomimic/models/openpi_compat.py create mode 100644 egomimic/models/paligemma_init.py create mode 100644 tests/integration/test_paligemma_init_real.py create mode 100644 tests/unit/test_openpi_compat.py create mode 100644 tests/unit/test_paligemma_init.py diff --git a/README.md b/README.md index ea399a3751..f2f57eada7 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,7 @@ Every cluster-specific default in the configs can be overridden without editing | `EGOVERSE_CACHE_DIR` | `paths.cache_dir`; norm stats cached by content under `/norm_stats/` | `$EGOVERSE_LOG_DIR/cache` | | `EGOVERSE_ROOT` | `paths.root_dir` | the directory you launch from | | `EGOVERSE_PI05_WEIGHTS` | pi0.5 base weights directory (`model.robomimic_model.config.pytorch_weight_path`) | lab PACE path | +| `EGOVERSE_PALIGEMMA_WEIGHTS` | PaliGemma checkpoint for the from-scratch pi arm (`...config.paligemma_weight_path`), a local dir or an HF repo id | `google/paligemma-3b-pt-224` | | `WANDB_ENTITY`, `WANDB_PROJECT` | `logger.wandb.entity` / `.project` | `rl2-group` / `zarr_test` | ### Other Settings diff --git a/egomimic/algo/pi.py b/egomimic/algo/pi.py index 6f384bb536..27aec4ced3 100644 --- a/egomimic/algo/pi.py +++ b/egomimic/algo/pi.py @@ -16,6 +16,8 @@ from transformers import AutoTokenizer from egomimic.algo.algo import Algo +from egomimic.models import openpi_compat +from egomimic.models.paligemma_init import load_paligemma_weights, select_init_source from egomimic.models.preprocess_pi_obs import ( PI_CAMERA_SLOTS, _concat_proprio, @@ -38,6 +40,10 @@ logger.setLevel(logging.INFO) logger.propagate = True # Explicitly enable propagation (default, but ensures it works) +# openpi targets transformers 4.53 (this repo's pin); these shims let it run on a +# 5.x venv too. See egomimic/models/openpi_compat.py. +openpi_compat.apply() + class PI(Algo): """ """ @@ -176,7 +182,15 @@ def __init__( self.model = openpi.models_pytorch.pi0_pytorch.PI0Pytorch(model_cfg) - if self.config.pytorch_weight_path is not None: + # Three mutually exclusive starting points, in decreasing order of what + # is pretrained: the full pi0.5 base checkpoint, PaliGemma only (the + # action expert and the action/time projections stay at their init), or + # nothing at all. + paligemma_weight_path = getattr(self.config, "paligemma_weight_path", None) + init_source = select_init_source( + self.config.pytorch_weight_path, paligemma_weight_path + ) + if init_source == "pi05_base": model_path = os.path.join( self.config.pytorch_weight_path, "model.safetensors" ) @@ -184,21 +198,23 @@ def __init__( raise FileNotFoundError( f"Pretrained weight file not found: {model_path}" ) - target = ( - self.model.module - if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) - else self.model - ) - safetensors.torch.load_model(target, model_path) + safetensors.torch.load_model(self.model, model_path) logger.info( "Loaded pretrained weights from %s (%d parameters)", model_path, - sum(p.numel() for p in target.parameters()), + sum(p.numel() for p in self.model.parameters()), + ) + elif init_source == "paligemma": + load_paligemma_weights( + self.model, + paligemma_weight_path, + allow_mirror=getattr(self.config, "paligemma_allow_mirror", False), ) else: logger.warning( - "No pytorch_weight_path: no base weights loaded. The weights must " - "come from a checkpoint; otherwise this trains from scratch." + "Neither pytorch_weight_path nor paligemma_weight_path is set: no " + "base weights loaded. The weights must come from a checkpoint; " + "otherwise this trains from scratch, VLM included." ) # openpi's PyTorch port hard-codes 32-wide action projections # regardless of ``Pi0Config.action_dim``. Wider action spaces (the diff --git a/egomimic/hydra_configs/model/pi0.5_base.yaml b/egomimic/hydra_configs/model/pi0.5_base.yaml index 4a05b646ee..283ec7b4bd 100644 --- a/egomimic/hydra_configs/model/pi0.5_base.yaml +++ b/egomimic/hydra_configs/model/pi0.5_base.yaml @@ -5,6 +5,17 @@ robomimic_model: pytorch_training_precision: bfloat16 # directory holding model.safetensors of the pi0.5 base ($EGOVERSE_PI05_WEIGHTS overrides) pytorch_weight_path: ${oc.env:EGOVERSE_PI05_WEIGHTS,/storage/home/hcoda1/4/paphiwetsa3/r-dxu345-0/projects/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch} + # The alternative to pytorch_weight_path: a local directory or an HF repo id + # holding the paligemma-3b-*-224 safetensors. Only the VLM prefix (SigLIP + # tower, multimodal projector, Gemma-2B) is loaded from it; everything pi0.5 + # adds -- the action expert, the action in/out projections, the adaRMS time + # MLPs -- trains from scratch. Setting both paths is an error. Flip it on + # with the model/pi0.5_paligemma_init.yaml overlay. + paligemma_weight_path: null + # google/paligemma-3b-pt-224 is license-gated. true = without an HF token + # that has accepted the Gemma terms, fall back to a third-party mirror whose + # shards sha256-match the canonical repo; false (default) = raise. + paligemma_allow_mirror: false model: pi05: true # Width of the model's action vector. 32 is openpi's pretrained width diff --git a/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml new file mode 100644 index 0000000000..91abc72867 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml @@ -0,0 +1,15 @@ +defaults: + - pi0.5_bc_mecka_6d + - pi0.5_paligemma_init + - _self_ + +# pi0.5_bc_mecka_6d with the pi0.5 base checkpoint swapped for a PaliGemma-only +# init -- the ablation that asks how much of the mecka 6D result comes from +# pi0.5's robot pretraining rather than from the VLM. Everything else (the +# cartesian 6D action layout, the action converter, the prompt blocks, the +# optimizer) is the parent's; see pi0.5_paligemma_init.yaml for the two keys +# this arm changes and for the lr knobs worth sweeping. +# +# Pairs with the same data configs as the parent (data=mecka_all_6d, +# mecka_fold_freeform*_6d) and evaluator=eval_pi_wristframe_6d. Norm stats are +# data-config-specific, not model-specific, so this arm reuses the parent's. diff --git a/egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml b/egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml new file mode 100644 index 0000000000..4e56103988 --- /dev/null +++ b/egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml @@ -0,0 +1,30 @@ +# OVERLAY, not a runnable model config. Compose it AFTER any pi0.5_* config to +# swap the pi0.5 base checkpoint for a PaliGemma-only init: the VLM prefix comes +# from the public paligemma-3b-pt-224 release and the action expert, the action +# in/out projections and the adaRMS time MLPs start from their PyTorch init. +# +# defaults: +# - pi0.5_bc_mecka_6d +# - pi0.5_paligemma_init +# - _self_ +# +# It sets only the two init keys, so it never clobbers the config it is layered +# on (action_dim, ac_keys, domains, the prompt flags all survive). Per-run, +# without a new file: +# model.robomimic_model.config.pytorch_weight_path=null \ +# model.robomimic_model.config.paligemma_weight_path=google/paligemma-3b-pt-224 +# +# The optimizer is left at pi0.5_base's fine-tuning recipe (constant 3e-5 after +# 2000 warmup steps). A randomly initialized action expert may well want a +# higher peak and a longer warmup; that is a per-run decision, e.g. +# model.optimizer.lr=1e-4 model.scheduler.num_warmup_steps=5000 +# +# Multi-rank runs should pre-fetch the checkpoint once, or every rank races on +# the same Hugging Face download lock: +# python -c "from egomimic.models.paligemma_init import resolve_paligemma_dir as r; print(r('google/paligemma-3b-pt-224'))" +robomimic_model: + config: + pytorch_weight_path: null + # $EGOVERSE_PALIGEMMA_WEIGHTS overrides: point it at a local snapshot dir to + # skip the Hub entirely. + paligemma_weight_path: ${oc.env:EGOVERSE_PALIGEMMA_WEIGHTS,google/paligemma-3b-pt-224} diff --git a/egomimic/models/openpi_compat.py b/egomimic/models/openpi_compat.py new file mode 100644 index 0000000000..ce6f5e15a3 --- /dev/null +++ b/egomimic/models/openpi_compat.py @@ -0,0 +1,129 @@ +"""Adapters that let openpi's PyTorch pi0 run against transformers 5. + +`external/openpi` pins `transformers==4.53.2` and ships a `transformers_replace` +patch written against it (its own `check.py` asserts that version), which is also +this repo's pin. A venv on transformers 5 works too, but there the stock modules +have moved on, so a few of openpi's call sites no longer line up. Nothing here +runs below transformers 5, and each shim is a no-op when it finds openpi's own +spelling already in place -- so porting the fork's patch forward makes this +module inert. Delete it then. + +Applied from `egomimic.algo.pi`, which is the only importer of openpi. +""" + +import logging + +import torch + +logger = logging.getLogger(__name__) + +_APPLIED = False + + +def _patch_embed_image() -> bool: + """`get_image_features` returns a `BaseModelOutputWithPooling`, not a tensor. + + transformers 5's `PaliGemmaModel.get_image_features` returns the vision + output object with the projected features stashed on `pooler_output`, while + openpi's `embed_prefix` reads `.shape` off the result. Unwrap it. + """ + from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel + + original = PaliGemmaWithExpertModel.embed_image + + def embed_image(self, image: torch.Tensor) -> torch.Tensor: + features = original(self, image) + if isinstance(features, torch.Tensor): + return features + return features.pooler_output + + PaliGemmaWithExpertModel.embed_image = embed_image + return True + + +def _patch_paligemma_bc_properties() -> bool: + """`paligemma.language_model` / `.vision_tower` / `.multi_modal_projector`. + + transformers 5 moved those submodules under `PaliGemmaModel` and dropped the + back-compat properties that used to forward to them from + `PaliGemmaForConditionalGeneration`. openpi still reads all three off the + outer model. Put the forwarders back. + """ + from transformers import PaliGemmaForConditionalGeneration + + added = [] + for name in ("language_model", "vision_tower", "multi_modal_projector"): + if hasattr(PaliGemmaForConditionalGeneration, name): + continue + setattr( + PaliGemmaForConditionalGeneration, + name, + property(lambda self, _name=name: getattr(self.model, _name)), + ) + added.append(name) + return bool(added) + + +# openpi keeps SigLIP's patch and position embeddings in float32 by substring on +# `vision_tower.vision_model.embeddings.*`; the flattened 5.x spelling of the same +# three tensors. +_FP32_KEEP_FLAT = ( + "vision_tower.embeddings.patch_embedding.", + "vision_tower.embeddings.position_embedding.", +) + + +def _patch_fp32_keep_list() -> bool: + """openpi's float32 keep-list misses the SigLIP embeddings. + + `to_bfloat16_for_selected_params` casts everything to bf16 and then restores + float32 on names containing `vision_tower.vision_model.embeddings.*`. + transformers 5 flattened `SiglipVisionModel`, so nothing matches and the + patch embedding and position embedding stay bf16 -- the tensors openpi + protects because small updates vanish in bf16. Re-apply the keep under the + flattened names; with the nested names present they match nothing. + """ + from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel + + original = PaliGemmaWithExpertModel.to_bfloat16_for_selected_params + + def to_bfloat16_for_selected_params(self, precision="bfloat16"): + original(self, precision) + if precision != "bfloat16": + return + for name, param in self.named_parameters(): + if any(selector in name for selector in _FP32_KEEP_FLAT): + param.data = param.data.to(dtype=torch.float32) + + PaliGemmaWithExpertModel.to_bfloat16_for_selected_params = ( + to_bfloat16_for_selected_params + ) + return True + + +def apply() -> None: + """Install every shim once. Safe to call repeatedly.""" + global _APPLIED + if _APPLIED: + return + _APPLIED = True + + import transformers + + major = int(transformers.__version__.split(".")[0]) + if major < 5: + return + + shims = ( + ("embed_image", _patch_embed_image), + ("paligemma_bc_properties", _patch_paligemma_bc_properties), + ("fp32_keep_list", _patch_fp32_keep_list), + ) + applied = [name for name, fn in shims if fn()] + if applied: + logger.info( + "openpi targets transformers 4.53 and this environment has %s; " + "applied compatibility shims: %s", + transformers.__version__, + ", ".join(applied), + ) diff --git a/egomimic/models/paligemma_init.py b/egomimic/models/paligemma_init.py new file mode 100644 index 0000000000..55799580ea --- /dev/null +++ b/egomimic/models/paligemma_init.py @@ -0,0 +1,296 @@ +"""PaliGemma-only initialization for the pi0.5 chassis. + +``egomimic.algo.pi.PI`` normally starts from the full pi0.5 base checkpoint. +This module is the other entry point: the VLM half (SigLIP tower, multimodal +projector, Gemma-2B prefix) comes from the public PaliGemma release, and +everything pi0.5 adds on top -- the action expert, the action in/out +projections, the adaRMS time MLPs -- keeps its PyTorch init. + +openpi builds the prefix as a stock ``PaliGemmaForConditionalGeneration`` whose +config is the ``gemma_2b`` variant (width 2048, depth 18, head_dim 256) next to +a SigLIP-so400m/14 tower, which is exactly the ``paligemma-3b-*-224`` +architecture, so the released checkpoint drops straight in. +""" + +import hashlib +import logging +import os +import re +from pathlib import Path + +import torch +from safetensors import safe_open + +logger = logging.getLogger(__name__) + +CANONICAL_REPO = "google/paligemma-3b-pt-224" + +# `CANONICAL_REPO` is license-gated: a machine without an HF token that has +# accepted the Gemma terms gets a 401 on the weight shards (the repo *metadata* +# stays anonymously readable, which is what makes the check below possible). +# These repos carry the same three shards. Every one is sha256-checked against +# the canonical repo's file metadata before it is used, so a mirror cannot +# substitute different weights -- it is only a transport. Opt-in +# (`paligemma_allow_mirror`): integrity is checked, but the weights of record +# should come from the canonical repo unless a run says otherwise. +MIRROR_REPOS = ( + "leo009/paligemma-3b-pt-224", + "hehe156/paligemma-3b-pt-224", +) + +# Only the shards: the loader reads nothing else and openpi builds the config in +# code, so an unverified json would only sit in the HF cache. +_ALLOW_PATTERNS = ["*.safetensors"] + +# Released PaliGemma checkpoint keys -> the module tree. These are what +# transformers applies via +# ``PaliGemmaForConditionalGeneration._checkpoint_conversion_mapping``; spelled +# out here because we stream the shards ourselves rather than going through +# ``from_pretrained`` (which would build a second full-size model alongside the +# one openpi already made). The vision rule is the transformers >= 5 spelling, +# which flattened ``SiglipVisionModel``; below 5 the ``vision_model.`` hop is +# still there, and ``_module_name`` picks whichever the live tree has. +# +# ``load_paligemma_weights`` fails loudly if these leave any parameter +# uncovered, so a future rename surfaces as an error, not a silent partial load. +_KEY_REMAP = ( + (re.compile(r"^language_model\.model\."), "model.language_model."), + (re.compile(r"^language_model\.lm_head\."), "lm_head."), + (re.compile(r"^vision_tower\.vision_model\."), "model.vision_tower."), + (re.compile(r"^vision_tower\."), "model.vision_tower."), + (re.compile(r"^multi_modal_projector\."), "model.multi_modal_projector."), +) + + +# The released checkpoint pads the token embedding from PaliGemma's 257152-token +# vocabulary up to 257216 rows. Row 257152 is -- which openpi never looks +# up, because ``embed_prefix`` concatenates SigLIP features directly instead of +# scattering them into an placeholder -- and the 63 rows after it are +# unused padding. openpi builds the table at exactly 257152, so these two +# tensors are truncated on their vocabulary axis; every real text token keeps its +# row, since the padding is at the end. +_VOCAB_TRUNCATABLE = frozenset( + {"model.language_model.embed_tokens.weight", "lm_head.weight"} +) + + +def _remap_key(key: str) -> str: + for pattern, replacement in _KEY_REMAP: + new, n = pattern.subn(replacement, key, count=1) + if n: + return new + return key + + +def _module_name(key: str, own: dict) -> str: + """Where checkpoint ``key`` lives in ``own``, the model's named tensors.""" + name = _remap_key(key) + if name not in own and key.startswith("vision_tower.vision_model."): + nested = "model." + key + if nested in own: + return nested + return name + + +def _fit(name: str, tensor: torch.Tensor, param: torch.Tensor) -> torch.Tensor: + if tensor.shape == param.shape: + return tensor + if ( + name in _VOCAB_TRUNCATABLE + and tensor.ndim == param.ndim == 2 + and tensor.shape[0] > param.shape[0] + and tensor.shape[1] == param.shape[1] + ): + logger.info( + "Truncating %s from %d to %d vocabulary rows (padding tokens).", + name, + tensor.shape[0], + param.shape[0], + ) + return tensor[: param.shape[0]] + raise ValueError( + f"Shape mismatch for {name}: checkpoint {tuple(tensor.shape)} " + f"vs model {tuple(param.shape)}" + ) + + +def select_init_source( + pytorch_weight_path: str | None, paligemma_weight_path: str | None +) -> str: + """Which set of starting weights a PI run uses: the full pi0.5 base + checkpoint (``"pi05_base"``), PaliGemma with a fresh action expert + (``"paligemma"``), or nothing at all (``"none"``).""" + if pytorch_weight_path is not None and paligemma_weight_path is not None: + raise ValueError( + "pytorch_weight_path and paligemma_weight_path are both set. The " + "pi0.5 base checkpoint already contains a trained PaliGemma; set " + "pytorch_weight_path=null to start from PaliGemma with a fresh " + "action expert." + ) + if pytorch_weight_path is not None: + return "pi05_base" + if paligemma_weight_path is not None: + return "paligemma" + return "none" + + +def _shard_digests(repo: str) -> dict[str, str]: + """``{shard filename: sha256}`` from a repo's file metadata.""" + from huggingface_hub import HfApi + + info = HfApi().model_info(repo, files_metadata=True) + digests = {} + for sibling in info.siblings: + if not sibling.rfilename.endswith(".safetensors"): + continue + lfs = getattr(sibling, "lfs", None) + sha = ( + lfs.get("sha256") if isinstance(lfs, dict) else getattr(lfs, "sha256", None) + ) + if sha: + digests[sibling.rfilename] = sha + return digests + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(16 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def _verify_shards(directory: str, expected: dict[str, str]) -> None: + root = Path(directory) + for name, sha in expected.items(): + shard = root / name + if not shard.is_file(): + raise FileNotFoundError(f"{shard} is missing from the mirror snapshot") + actual = _sha256(shard) + if actual != sha: + raise ValueError( + f"{shard} does not match {CANONICAL_REPO}: expected sha256 {sha}, " + f"got {actual}. Refusing to initialize from it." + ) + + +def resolve_paligemma_dir(source: str, *, allow_mirror: bool = False) -> str: + """Local directory holding the PaliGemma checkpoint named by ``source``. + + ``source`` is a local path (returned as-is) or a Hugging Face repo id. When + it is the gated canonical repo, this machine is not authorized for it (401 + or 403, not a rate limit or an outage) and ``allow_mirror`` is set, fall + back to a sha256-verified mirror. + """ + if os.path.isdir(source): + return source + + from huggingface_hub import snapshot_download + from huggingface_hub.errors import GatedRepoError, HfHubHTTPError + + try: + return snapshot_download(source, allow_patterns=_ALLOW_PATTERNS) + except HfHubHTTPError as gated: + status = getattr(getattr(gated, "response", None), "status_code", None) + unauthorized = isinstance(gated, GatedRepoError) or status in (401, 403) + if not (allow_mirror and unauthorized and source == CANONICAL_REPO): + raise + logger.warning( + "%s is gated for this machine (%s) and paligemma_allow_mirror is set. " + "Falling back to a mirror, sha256-verified against %s. Set HF_TOKEN " + "to a token that has accepted the Gemma terms to pull the canonical " + "repo instead.", + CANONICAL_REPO, + type(gated).__name__, + CANONICAL_REPO, + ) + + expected = _shard_digests(CANONICAL_REPO) + if not expected: + raise RuntimeError( + f"Cannot read shard checksums for {CANONICAL_REPO}; refusing to use " + "an unverified mirror." + ) + + failures = [] + for mirror in MIRROR_REPOS: + try: + if _shard_digests(mirror) != expected: + raise ValueError("shard checksums differ from the canonical repo") + directory = snapshot_download(mirror, allow_patterns=_ALLOW_PATTERNS) + _verify_shards(directory, expected) + except Exception as exc: # try the next mirror, report them all if none works + failures.append(f"{mirror}: {type(exc).__name__}: {exc}") + continue + logger.warning( + "Initializing PaliGemma from mirror %s (all %d shards sha256-match %s).", + mirror, + len(expected), + CANONICAL_REPO, + ) + return directory + + raise RuntimeError( + f"{CANONICAL_REPO} is gated and no verified mirror worked:\n " + + "\n ".join(failures) + ) + + +def load_paligemma_weights(model, source: str, *, allow_mirror: bool = False) -> str: + """Copy a PaliGemma checkpoint into ``model``'s VLM prefix, in place. + + ``model`` is an ``openpi.models_pytorch.pi0_pytorch.PI0Pytorch``. Only + ``paligemma_with_expert.paligemma`` is touched; the action expert and the + action/time projections are left at their init. Returns the directory the + weights came from. + """ + directory = resolve_paligemma_dir(source, allow_mirror=allow_mirror) + target = model.paligemma_with_expert.paligemma + + # remove_duplicate=False so a tied lm_head.weight is reachable under both + # names; the checkpoint only carries one of them. + params = dict(target.named_parameters(remove_duplicate=False)) + # Buffers are accepted if the checkpoint ships them, but are not required: + # rotary inv_freq, Gemma's embed_scale and the SigLIP position_ids are all + # derived from the config at construction time and are absent from the + # release. + own = dict(target.named_buffers(remove_duplicate=False)) + own.update(params) + + shards = sorted(Path(directory).glob("*.safetensors")) + if not shards: + raise FileNotFoundError(f"No .safetensors under {directory}") + + loaded, unexpected = set(), [] + for shard in shards: + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping + name = _module_name(key, own) + param = own.get(name) + if param is None: + unexpected.append(key) + continue + tensor = _fit(name, handle.get_tensor(key), param) + with torch.no_grad(): + param.copy_(tensor.to(param.dtype)) + loaded.add(name) + + # lm_head is tied to embed_tokens and, either way, dead weight here: no code + # path in openpi's PyTorch pi0 reads it. + tied = set(getattr(type(target), "_tied_weights_keys", None) or ()) + missing = {n for n in params if n not in loaded} - tied + if missing or unexpected: + raise RuntimeError( + f"PaliGemma load from {directory} did not cover the prefix exactly: " + f"{len(missing)} missing (e.g. {sorted(missing)[:5]}), " + f"{len(unexpected)} unexpected (e.g. {unexpected[:5]})." + ) + + logger.info( + "Initialized the VLM prefix from PaliGemma at %s (%d tensors, %d parameters). " + "The action expert and the action/time projections stay randomly initialized.", + directory, + len(loaded), + sum(p.numel() for p in target.parameters()), + ) + return directory diff --git a/egomimic/robot/rollout.py b/egomimic/robot/rollout.py index 649357b21a..f148f4e58b 100644 --- a/egomimic/robot/rollout.py +++ b/egomimic/robot/rollout.py @@ -317,12 +317,21 @@ def _patch_checkpoint_paths(cls, ckpt_path): robomimic = cfg.get("model", {}).get("robomimic_model", {}) config = robomimic.get("config", {}) old_path = config.get("pytorch_weight_path") - if old_path is None or old_path == cls.LOCAL_WEIGHT_PATH: + # The state_dict holds every weight; a PaliGemma-init checkpoint would + # otherwise fetch the gated Hub repo on the robot. + paligemma = config.get("paligemma_weight_path") + if paligemma is None and ( + old_path is None or old_path == cls.LOCAL_WEIGHT_PATH + ): return ckpt_path - print( - f"[rollout] Patching pytorch_weight_path: {old_path} -> {cls.LOCAL_WEIGHT_PATH}" - ) - config["pytorch_weight_path"] = cls.LOCAL_WEIGHT_PATH + if paligemma is not None: + print(f"[rollout] Patching paligemma_weight_path: {paligemma} -> None") + config["paligemma_weight_path"] = None + else: + print( + f"[rollout] Patching pytorch_weight_path: {old_path} -> {cls.LOCAL_WEIGHT_PATH}" + ) + config["pytorch_weight_path"] = cls.LOCAL_WEIGHT_PATH ckpt["hyper_parameters"]["config_tree"] = OmegaConf.create(cfg) patched_path = ckpt_path + ".patched" _torch.save(ckpt, patched_path) diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index 6f9ed07e8f..e94203b810 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -44,6 +44,10 @@ _PI_WEIGHT_KEY = "model.robomimic_model.config.pytorch_weight_path" +_PI_BASE_WEIGHT_KEYS = ( + _PI_WEIGHT_KEY, + "model.robomimic_model.config.paligemma_weight_path", +) def _data_transform_modes(cfg: DictConfig) -> dict[str, str]: @@ -78,13 +82,18 @@ def _build_model_config_tree(cfg: DictConfig) -> DictConfig: modes = _data_transform_modes(cfg) if modes: tree.data_transform_modes = modes - has_weights = OmegaConf.select(tree, _PI_WEIGHT_KEY, default=None) is not None - if has_weights and _weights_from_checkpoint(cfg): + set_keys = [ + k + for k in _PI_BASE_WEIGHT_KEYS + if OmegaConf.select(tree, k, default=None) is not None + ] + if set_keys and _weights_from_checkpoint(cfg): log.info( - f"Loading every weight from {cfg.ckpt_path}: {_PI_WEIGHT_KEY}=null " - "for the model (the base safetensors need not exist here)" + f"Loading every weight from {cfg.ckpt_path}: {', '.join(set_keys)}=null " + "for the model (the base weights need not exist here)" ) - OmegaConf.update(tree, _PI_WEIGHT_KEY, None) + for k in set_keys: + OmegaConf.update(tree, k, None) return tree diff --git a/pi05.md b/pi05.md index ddcd12b643..a679425c25 100644 --- a/pi05.md +++ b/pi05.md @@ -25,3 +25,120 @@ uv run --active examples/convert_jax_model_to_pytorch.py \ --checkpoint_dir ../../egomimic/algo/pi_checkpoints/pi05_base \ --output_path ../../egomimic/algo/pi_checkpoints/pi05_base_pytorch ``` + +## Starting from PaliGemma instead of the pi0.5 base + +`model/pi0.5_paligemma_init.yaml` is an overlay that swaps the pi0.5 base +checkpoint for the public PaliGemma release: the VLM prefix (SigLIP tower, +multimodal projector, Gemma-2B) is loaded from `paligemma-3b-pt-224`, and +everything pi0.5 adds on top — the 300M action expert, the action in/out +projections, the adaRMS time MLPs — trains from scratch. That is the ablation +that separates "what the VLM knows" from "what pi0.5's robot pretraining added". + +A ready-made arm ships for the mecka 6D configs: + +```bash +python egomimic/trainHydra.py --config-name=train_zarr_cartesian_pi \ + data=mecka_all_6d \ + model=pi0.5_bc_mecka_6d_paligemma_init \ + evaluator=eval_pi_wristframe_6d +``` + +Layer the overlay onto any other `pi0.5_*` config the same way (`defaults: +[pi0.5_bc_aria, pi0.5_paligemma_init, _self_]`), or set the two keys per run: + +```bash +model.robomimic_model.config.pytorch_weight_path=null \ +model.robomimic_model.config.paligemma_weight_path=google/paligemma-3b-pt-224 +``` + +Setting both paths raises: the pi0.5 base checkpoint already contains a trained +PaliGemma, so letting one silently win would make a run's provenance +unreadable from its config. + +### Getting the weights + +`google/paligemma-3b-pt-224` is license-gated. Either authenticate + +```bash +huggingface-cli login # with a token that has accepted the Gemma terms +``` + +or set `paligemma_allow_mirror: true` to let `egomimic/models/paligemma_init.py` +fall back to a third-party mirror on a 401/403 — it compares every shard's +sha256 against the canonical repo's file metadata (readable without a token) +before using it, so a mirror is transport only and cannot substitute different +weights. It is off by default so the weights of record come from the canonical +repo. `$EGOVERSE_PALIGEMMA_WEIGHTS` points the arm at a local snapshot +directory and skips the Hub entirely. + +Pre-fetch once before a multi-rank run, or every rank races on the same +download lock: + +```bash +python -c "from egomimic.models.paligemma_init import resolve_paligemma_dir as r; print(r('google/paligemma-3b-pt-224'))" +``` + +### What the loader checks + +The released checkpoint does not drop in unchanged, so the load is strict about +the two places it differs and fails loudly anywhere else: + +- the checkpoint's `vision_tower.vision_model.*` keeps its `vision_model.` hop + below transformers 5 and loses it from 5 on (`SiglipVisionModel` was + flattened); the loader maps onto whichever the model has. +- the release pads the token embedding from PaliGemma's 257152-token vocabulary + to 257216 rows (`` plus 63 unused). openpi builds the table at exactly + 257152, so `embed_tokens` / `lm_head` are truncated on their vocabulary axis; + the padding is at the end, so every real text token keeps its row. + +Any parameter the checkpoint does not cover, and any checkpoint tensor with no +home in the module tree, raises — a partial load that quietly leaves half the +VLM random is the failure mode this exists to prevent. + +### Hyperparameters + +The overlay changes the init source and nothing else: the optimizer stays +`pi0.5_base`'s fine-tuning recipe (constant 3e-5 after 2000 warmup steps). A +randomly initialized action expert may want a higher peak and a longer warmup — +that is a per-run decision, e.g. + +```bash +model.optimizer.lr=1e-4 model.scheduler.num_warmup_steps=5000 +``` + +## transformers 4.53 vs 5.x + +`external/openpi` pins `transformers==4.53.2` and its `transformers_replace` +patch is written against that version (`models/siglip/check.py` asserts it). +So does `pyproject.toml`, and the install steps copy that patch into the venv: +4.53.2 plus the patch is the supported setup, what CI installs and what the +shared `emimic` venv carries. Which transformers a run gets is whatever its venv +has, so check before comparing numbers across machines. + +A venv on transformers 5.x (the PaliGemma-init arm was first run on 5.17 with a +partial hand-port of the patch: the Gemma modules carry openpi's adaRMS changes, +but `modeling_paligemma.py` and `modeling_siglip.py` are effectively stock 5.x) +also works, but three of openpi's assumptions do not survive it: + +- `PaliGemmaModel.get_image_features` now returns a `BaseModelOutputWithPooling` + with the projected features on `pooler_output`; openpi's `embed_prefix` reads + `.shape` off the result. +- transformers 5 dropped the back-compat properties that forwarded + `PaliGemmaForConditionalGeneration.language_model` / `.vision_tower` / + `.multi_modal_projector` to the inner `PaliGemmaModel`; openpi reads all three + off the outer model. +- openpi's float32 keep-list names the SigLIP patch and position embeddings + `vision_tower.vision_model.embeddings.*`; with `SiglipVisionModel` flattened + nothing matches and they would train in bf16. + +`egomimic/models/openpi_compat.py` bridges all three, applied from +`egomimic.algo.pi`. Each shim is guarded on what it finds and stands down +entirely below transformers 5, so porting the fork's patch forward makes the +module a no-op — delete it then. + +On 5.x only the PaliGemma-init arm loads. The pi0.5 base checkpoint +(`pi05_base_pytorch/model.safetensors`) stores the nested +`vision_tower.vision_model.*` names and `safetensors.torch.load_model` is +strict, so the default init path stops at startup on missing SigLIP keys. +Train from the pi0.5 base on the pinned 4.53.2. diff --git a/tests/integration/test_paligemma_init_real.py b/tests/integration/test_paligemma_init_real.py new file mode 100644 index 0000000000..e74135417f --- /dev/null +++ b/tests/integration/test_paligemma_init_real.py @@ -0,0 +1,253 @@ +"""The real PaliGemma release, loaded into the real openpi prefix. + +The unit tier proves the remapping and the coverage rules on a miniature module +tree; this is the part that can only be checked against the actual 3B +checkpoint -- that openpi's prefix and `paligemma-3b-pt-224` line up tensor for +tensor, that the values land unchanged, and that nothing outside the prefix +moves. It builds a 3.6B model and reads ~12 GB of weights, so it lives here. + + # the load itself: CPU, needs the checkpoint (or Hub access to fetch it) + python -m pytest --integration tests/integration/test_paligemma_init_real.py -q + # plus the flow-matching step, on a GPU node + python -m pytest --integration -m gpu tests/integration/test_paligemma_init_real.py -q +""" + +from __future__ import annotations + +import glob + +import pytest +import torch +from fixtures.train_harness import pi_unavailable +from safetensors import safe_open + +from egomimic.models.paligemma_init import ( + CANONICAL_REPO, + load_paligemma_weights, + resolve_paligemma_dir, +) + +pytestmark = pytest.mark.skipif( + pi_unavailable() is not None, reason=pi_unavailable() or "" +) + +PG = "paligemma_with_expert.paligemma." +EX = "paligemma_with_expert.gemma_expert." + +# These start random and must still be random afterwards: the whole point of +# this init is that pi0.5's own parameters are NOT pretrained. +FROM_SCRATCH = ( + EX + "model.layers.0.self_attn.q_proj.weight", + EX + "model.layers.17.mlp.up_proj.weight", + EX + "model.layers.0.input_layernorm.dense.weight", + EX + "model.norm.dense.weight", + "action_in_proj.weight", + "action_out_proj.weight", + "time_mlp_in.weight", + "time_mlp_out.weight", +) + +# module parameter -> the checkpoint key it must come from, one per remapping rule +PRETRAINED = { + PG + + "model.language_model.layers.0.self_attn.q_proj.weight": "language_model.model.layers.0.self_attn.q_proj.weight", + PG + + "model.language_model.layers.17.mlp.down_proj.weight": "language_model.model.layers.17.mlp.down_proj.weight", + PG + "model.language_model.norm.weight": "language_model.model.norm.weight", + PG + + "model.language_model.embed_tokens.weight": "language_model.model.embed_tokens.weight", + PG + + "model.vision_tower.encoder.layers.0.self_attn.q_proj.weight": "vision_tower.vision_model.encoder.layers.0.self_attn.q_proj.weight", + PG + + "model.vision_tower.encoder.layers.26.mlp.fc2.bias": "vision_tower.vision_model.encoder.layers.26.mlp.fc2.bias", + PG + + "model.vision_tower.embeddings.patch_embedding.weight": "vision_tower.vision_model.embeddings.patch_embedding.weight", + PG + + "model.multi_modal_projector.linear.weight": "multi_modal_projector.linear.weight", +} + + +@pytest.fixture(scope="module") +def checkpoint_dir() -> str: + try: + return resolve_paligemma_dir(CANONICAL_REPO) + except Exception as exc: # gated and offline, or no verified mirror reachable + pytest.skip(f"cannot resolve {CANONICAL_REPO}: {type(exc).__name__}: {exc}") + + +def _flat(named) -> dict: + """Parameters under their transformers-5 names, which the tables above use: + below 5 the SigLIP tree keeps a ``vision_model.`` hop.""" + return { + k.replace("vision_tower.vision_model.", "vision_tower."): v for k, v in named + } + + +@pytest.fixture(scope="module") +def loaded(checkpoint_dir): + """The real pi0.5 chassis with only its VLM prefix initialized.""" + import openpi.models.pi0_config + import openpi.models_pytorch.pi0_pytorch as pi0t + + # Production reaches openpi only through `egomimic.algo.pi`, which applies + # the transformers-5 shims at import. This builds PI0Pytorch directly, so + # take the same entry point rather than an unpatched openpi. + import egomimic.algo.pi # noqa: F401 + + config = openpi.models.pi0_config.Pi0Config( + dtype="bfloat16", + action_dim=32, + action_horizon=100, + max_token_len=180, + paligemma_variant="gemma_2b", + action_expert_variant="gemma_300m", + pi05=True, + ) + torch.manual_seed(0) + model = pi0t.PI0Pytorch(config) + before = _flat((k, v.detach().clone()) for k, v in model.named_parameters()) + load_paligemma_weights(model, checkpoint_dir) + return model, before, _flat(model.named_parameters()) + + +def test_pi05s_own_parameters_are_left_at_their_init(loaded): + _, before, after = loaded + for name in FROM_SCRATCH: + assert torch.equal(after[name], before[name]), f"{name} was overwritten" + + +def test_the_prefix_moved_off_its_init(loaded): + _, before, after = loaded + for name in PRETRAINED: + assert not torch.equal(after[name], before[name]), f"{name} was not loaded" + + +def test_the_prefix_matches_the_checkpoint_bit_for_bit(loaded, checkpoint_dir): + _, _, after = loaded + wanted = set(PRETRAINED.values()) + found = {} + for shard in sorted(glob.glob(checkpoint_dir + "/*.safetensors")): + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): # noqa: SIM118 - safe_open is not a Mapping + if key in wanted: + found[key] = handle.get_tensor(key) + assert set(found) == wanted, f"checkpoint is missing {wanted - set(found)}" + + for name, key in PRETRAINED.items(): + param = after[name] + reference = found[key] + if reference.shape != param.shape: # padded vocabulary + reference = reference[: param.shape[0]] + assert torch.equal( + param, reference.to(param.dtype) + ), f"value mismatch at {name}" + + +def test_openpis_mixed_precision_policy_survives_the_load(loaded): + # The norms openpi keeps in fp32 must not be dragged to bf16 by the copy, + # and the matmul weights must not be widened. + _, _, after = loaded + assert ( + after[PG + "model.language_model.layers.0.input_layernorm.weight"].dtype + is torch.float32 + ) + assert after[PG + "model.language_model.norm.weight"].dtype is torch.float32 + assert ( + after[PG + "model.language_model.layers.0.self_attn.q_proj.weight"].dtype + is torch.bfloat16 + ) + assert ( + after[PG + "model.vision_tower.encoder.layers.0.mlp.fc1.weight"].dtype + is torch.bfloat16 + ) + # openpi's keep-list, on either transformers (openpi_compat on 5.x) + for name in ( + "patch_embedding.weight", + "patch_embedding.bias", + "position_embedding.weight", + ): + assert ( + after[PG + "model.vision_tower.embeddings." + name].dtype is torch.float32 + ) + + +# ------------------------------------------------------------------ a real step + + +def _observation(batch: int, device, max_token_len: int): + """A minimal openpi observation: three camera slots, the wrists masked out + the way human data leaves them, and a prompt of real token ids.""" + from egomimic.models.preprocess_pi_obs import PI_CAMERA_SLOTS, _SimpleObservation + + image = torch.rand(batch, 3, 224, 224, device=device) * 2 - 1 + present = { + "base_0_rgb": True, + "left_wrist_0_rgb": False, + "right_wrist_0_rgb": False, + } + mask = torch.ones(batch, max_token_len, dtype=torch.bool, device=device) + return _SimpleObservation( + images={slot: image.clone() for slot in PI_CAMERA_SLOTS}, + image_masks={ + slot: torch.full((batch,), present[slot], dtype=torch.bool, device=device) + for slot in PI_CAMERA_SLOTS + }, + state=torch.zeros(batch, 32, device=device), + tokenized_prompt=torch.randint( + 0, 1000, (batch, max_token_len), device=device, dtype=torch.long + ), + tokenized_prompt_mask=mask, + token_ar_mask=mask.clone(), + token_loss_mask=mask.clone(), + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_one_flow_matching_step_trains_both_halves(loaded): + """The point of the arm: a finite loss whose gradient reaches the frozen-in + PaliGemma prefix AND the action expert that started from noise. A load that + left half the model detached would still produce a loss here, so check the + gradients, not just the number.""" + model, _, _ = loaded + device = torch.device("cuda") + model = model.to(device) + model.train() + + observation = _observation(1, device, model.config.max_token_len) + actions = torch.randn( + 1, model.config.action_horizon, model.config.action_dim, device=device + ) + + torch.manual_seed(0) + losses = model.forward(observation, actions) + loss = ( + torch.stack(losses).mean() + if isinstance(losses, (list, tuple)) + else losses.mean() + ) + assert torch.isfinite(loss), f"non-finite flow-matching loss: {loss}" + + loss.backward() + probes = { + "prefix (pretrained)": model.paligemma_with_expert.paligemma.model.language_model.layers[ + 0 + ].self_attn.q_proj.weight, + "vision tower (pretrained)": model.paligemma_with_expert.paligemma.model.vision_tower.encoder.layers[ + 0 + ].self_attn.q_proj.weight, + "action expert (from scratch)": model.paligemma_with_expert.gemma_expert.model.layers[ + 0 + ].self_attn.q_proj.weight, + "action_out_proj (from scratch)": model.action_out_proj.weight, + } + for label, param in probes.items(): + assert param.grad is not None, f"no gradient reached {label}" + assert torch.isfinite(param.grad).all(), f"non-finite gradient at {label}" + assert param.grad.abs().sum() > 0, f"zero gradient at {label}" + + # and a step actually moves them + before = {k: v.detach().clone() for k, v in probes.items()} + torch.optim.SGD(model.parameters(), lr=1e-3).step() + for label, param in probes.items(): + assert not torch.equal(param, before[label]), f"{label} did not move" diff --git a/tests/unit/test_checkpoint_resume.py b/tests/unit/test_checkpoint_resume.py index 0e233e51ee..d35a9f18c0 100644 --- a/tests/unit/test_checkpoint_resume.py +++ b/tests/unit/test_checkpoint_resume.py @@ -67,6 +67,24 @@ def test_explicit_ckpt_path_nulls_base_weights_for_model_only( assert cfg.model.robomimic_model.config.pytorch_weight_path is not None +def test_explicit_ckpt_path_nulls_paligemma_init_too(compose_resolve, tmp_path): + """A PaliGemma-init run resumed from a checkpoint must not fetch the gated + PaliGemma repo again.""" + ckpt = _touch(tmp_path / "runA" / "last.ckpt") + cfg = _pi_cfg( + compose_resolve, + tmp_path, + "model.robomimic_model.config.pytorch_weight_path=null", + "model.robomimic_model.config.paligemma_weight_path=/some/paligemma", + f"ckpt_path={ckpt}", + ) + train_hydra._prepare_checkpoint_resume(cfg) + config = train_hydra._build_model_config_tree(cfg).model.robomimic_model.config + assert config.pytorch_weight_path is None + assert config.paligemma_weight_path is None + assert cfg.model.robomimic_model.config.paligemma_weight_path == "/some/paligemma" + + @pytest.mark.parametrize("ckpt_path", ["last", "best", "hpc", "missing.ckpt"]) def test_ckpt_path_that_is_not_a_file_keeps_base_weights( compose_resolve, tmp_path, ckpt_path diff --git a/tests/unit/test_openpi_compat.py b/tests/unit/test_openpi_compat.py new file mode 100644 index 0000000000..951a90fa6a --- /dev/null +++ b/tests/unit/test_openpi_compat.py @@ -0,0 +1,115 @@ +"""openpi pins transformers 4.53 and ships a patch written against it; a venv +on transformers 5 needs these shims. The shims in egomimic/models/openpi_compat.py bridge the +gap and must disappear on their own once the fork's patch is ported forward.""" + +import pytest +import torch + +pytest.importorskip("openpi.models_pytorch.gemma_pytorch") + +from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel # noqa: E402 + +from egomimic.models import openpi_compat # noqa: E402 + + +@pytest.fixture(autouse=True) +def _restore_embed_image(): + original = PaliGemmaWithExpertModel.embed_image + cast = PaliGemmaWithExpertModel.to_bfloat16_for_selected_params + applied = openpi_compat._APPLIED + yield + PaliGemmaWithExpertModel.embed_image = original + PaliGemmaWithExpertModel.to_bfloat16_for_selected_params = cast + openpi_compat._APPLIED = applied + + +class _Outputs: + """What transformers 5's ``get_image_features`` hands back: the vision + output object with the projected features on ``pooler_output``.""" + + def __init__(self, features): + self.last_hidden_state = torch.zeros_like(features) + self.pooler_output = features + + +class _Prefix: + def __init__(self, result): + self._result = result + + class _Inner: + def __init__(self, result): + self._result = result + + def get_image_features(self, image): + return self._result + + @property + def model(self): + return self._Inner(self._result) + + +def _embed(result): + fake = object.__new__(PaliGemmaWithExpertModel) + fake.paligemma = _Prefix(result) + return PaliGemmaWithExpertModel.embed_image(fake, torch.zeros(1, 3, 224, 224)) + + +def test_a_wrapped_output_is_unwrapped_to_its_features(): + # Without this, openpi's embed_prefix reads `.shape` off the output object. + features = torch.randn(2, 256, 2048) + openpi_compat._patch_embed_image() + assert torch.equal(_embed(_Outputs(features)), features) + + +def test_a_plain_tensor_passes_straight_through(): + # openpi's own patched transformers already returns the tensor; the shim + # must not change that behaviour. + features = torch.randn(2, 256, 2048) + openpi_compat._patch_embed_image() + assert torch.equal(_embed(features), features) + + +def test_apply_is_a_no_op_the_second_time(caplog): + openpi_compat._APPLIED = False + openpi_compat.apply() + first = PaliGemmaWithExpertModel.embed_image + openpi_compat.apply() + assert PaliGemmaWithExpertModel.embed_image is first + + +def test_the_shim_stands_down_on_the_transformers_openpi_pins(monkeypatch): + import transformers + + original = PaliGemmaWithExpertModel.embed_image + monkeypatch.setattr(transformers, "__version__", "4.53.2") + openpi_compat._APPLIED = False + openpi_compat.apply() + assert PaliGemmaWithExpertModel.embed_image is original + + +def _siglip_embeddings(nested: bool) -> torch.nn.Module: + root = torch.nn.Module() + tower = torch.nn.Module() + parent = tower + if nested: + tower.vision_model = torch.nn.Module() + parent = tower.vision_model + parent.embeddings = torch.nn.Module() + parent.embeddings.patch_embedding = torch.nn.Conv2d(3, 4, 2) + parent.embeddings.position_embedding = torch.nn.Embedding(4, 4) + parent.encoder = torch.nn.Linear(4, 4) + root.vision_tower = tower + return root + + +@pytest.mark.parametrize("nested", [True, False]) +def test_the_siglip_embeddings_stay_float32(nested): + """openpi's keep-list matches the nested 4.x names only; with the shim the + flattened 5.x names are kept too, and everything else still goes bf16.""" + openpi_compat._patch_fp32_keep_list() + module = _siglip_embeddings(nested) + PaliGemmaWithExpertModel.to_bfloat16_for_selected_params(module, "bfloat16") + dtypes = {n: p.dtype for n, p in module.named_parameters()} + for name, dtype in dtypes.items(): + expected = torch.float32 if "embeddings" in name else torch.bfloat16 + assert dtype == expected, name diff --git a/tests/unit/test_paligemma_init.py b/tests/unit/test_paligemma_init.py new file mode 100644 index 0000000000..4e79db9db5 --- /dev/null +++ b/tests/unit/test_paligemma_init.py @@ -0,0 +1,454 @@ +"""PaliGemma-only init: the VLM prefix comes from the public PaliGemma release +and everything pi0.5 adds on top starts from its PyTorch init. + +No network and no 3B model: the loader is exercised against a miniature module +tree, and against a tiny real ``PaliGemmaForConditionalGeneration`` built by +whichever transformers is installed (the vision tree differs across 4.x and +5.x), with safetensors written on the fly; the download resolution is driven +through stubs. The one thing that needs the real thing -- +that the released checkpoint covers openpi's prefix exactly -- is asserted by +``load_paligemma_weights`` itself (it raises on any missing or unexpected +tensor), so a rename upstream cannot pass silently. +""" + +import pytest +import torch +from safetensors.torch import save_file +from torch import nn + +from egomimic.models import paligemma_init as pgi + +CANONICAL = pgi.CANONICAL_REPO + + +# ---------------------------------------------------------------- key remapping + + +# Spelled exactly as google/paligemma-3b-pt-224 ships them. +@pytest.mark.parametrize( + ("checkpoint_key", "module_name"), + [ + ( + "language_model.model.layers.0.self_attn.q_proj.weight", + "model.language_model.layers.0.self_attn.q_proj.weight", + ), + ( + "language_model.model.embed_tokens.weight", + "model.language_model.embed_tokens.weight", + ), + ("language_model.model.norm.weight", "model.language_model.norm.weight"), + ("language_model.lm_head.weight", "lm_head.weight"), + # transformers >= 5 flattened SiglipVisionModel: the `.vision_model.` hop + # in the checkpoint has no counterpart in the module tree. + ( + "vision_tower.vision_model.encoder.layers.26.mlp.fc2.bias", + "model.vision_tower.encoder.layers.26.mlp.fc2.bias", + ), + ( + "vision_tower.vision_model.embeddings.patch_embedding.weight", + "model.vision_tower.embeddings.patch_embedding.weight", + ), + ( + "multi_modal_projector.linear.bias", + "model.multi_modal_projector.linear.bias", + ), + ], +) +def test_checkpoint_keys_map_onto_the_module_tree(checkpoint_key, module_name): + assert pgi._remap_key(checkpoint_key) == module_name + + +def test_unknown_prefix_is_left_alone(): + assert pgi._remap_key("action_in_proj.weight") == "action_in_proj.weight" + + +# ------------------------------------------------------------- vocab truncation + + +def test_padded_vocabulary_is_truncated_to_the_model_width(): + # The release pads 257152 -> 257216; the extra rows are + padding. + tensor, param = torch.arange(12.0).reshape(6, 2), torch.zeros(4, 2) + fitted = pgi._fit("model.language_model.embed_tokens.weight", tensor, param) + assert fitted.shape == param.shape + assert torch.equal(fitted, tensor[:4]) # padding is at the END + + +def test_truncation_is_not_applied_to_other_tensors(): + with pytest.raises(ValueError, match="Shape mismatch"): + pgi._fit( + "model.vision_tower.encoder.layers.0.mlp.fc1.weight", + torch.zeros(6, 2), + torch.zeros(4, 2), + ) + + +def test_a_narrower_checkpoint_is_never_padded_out(): + with pytest.raises(ValueError, match="Shape mismatch"): + pgi._fit("lm_head.weight", torch.zeros(3, 2), torch.zeros(4, 2)) + + +# ------------------------------------------------------------------ init source + + +def test_pi05_base_wins_when_only_it_is_set(): + assert pgi.select_init_source("/ckpt/pi05_base_pytorch", None) == "pi05_base" + + +def test_paligemma_when_the_base_checkpoint_is_off(): + assert pgi.select_init_source(None, CANONICAL) == "paligemma" + + +def test_neither_is_reported_as_no_weights(): + assert pgi.select_init_source(None, None) == "none" + + +def test_both_at_once_is_refused(): + # The pi0.5 base already contains a trained PaliGemma; silently letting one + # win would make the run's provenance unreadable from its config. + with pytest.raises(ValueError, match="both set"): + pgi.select_init_source("/ckpt/pi05_base_pytorch", CANONICAL) + + +# ------------------------------------------------------------ the loader itself + + +class _Prefix(nn.Module): + """Shaped like the part of ``PaliGemmaForConditionalGeneration`` the loader + walks: a ``model.`` subtree, a tied ``lm_head`` and a derived buffer.""" + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, vocab=4, width=3): + super().__init__() + self.model = nn.Module() + self.model.language_model = nn.Module() + self.model.language_model.embed_tokens = nn.Embedding(vocab, width) + self.model.vision_tower = nn.Module() + self.model.vision_tower.post_layernorm = nn.Linear(width, width) + self.lm_head = nn.Linear(width, vocab, bias=False) + self.lm_head.weight = self.model.language_model.embed_tokens.weight # tied + # config-derived, absent from every released checkpoint + self.model.language_model.register_buffer("inv_freq", torch.ones(width)) + + +class _Model(nn.Module): + def __init__(self): + super().__init__() + self.paligemma_with_expert = nn.Module() + self.paligemma_with_expert.paligemma = _Prefix() + self.action_in_proj = nn.Linear(3, 3) + + +def _write_checkpoint(tmp_path, tensors, shards=2): + items = list(tensors.items()) + size = -(-len(items) // shards) + for i in range(shards): + chunk = dict(items[i * size : (i + 1) * size]) + if chunk: + save_file(chunk, str(tmp_path / f"model-{i}.safetensors")) + return str(tmp_path) + + +def _full_checkpoint(vocab=4, width=3): + torch.manual_seed(1) + return { + "language_model.model.embed_tokens.weight": torch.randn(vocab, width), + "vision_tower.vision_model.post_layernorm.weight": torch.randn(width, width), + "vision_tower.vision_model.post_layernorm.bias": torch.randn(width), + } + + +def test_the_prefix_loads_and_nothing_else_moves(tmp_path): + model = _Model() + before = model.action_in_proj.weight.detach().clone() + tensors = _full_checkpoint() + directory = _write_checkpoint(tmp_path, tensors) + + assert pgi.load_paligemma_weights(model, directory) == directory + + prefix = model.paligemma_with_expert.paligemma + assert torch.equal( + prefix.model.language_model.embed_tokens.weight, + tensors["language_model.model.embed_tokens.weight"], + ) + assert torch.equal( + prefix.model.vision_tower.post_layernorm.bias, + tensors["vision_tower.vision_model.post_layernorm.bias"], + ) + # everything outside the prefix is untouched + assert torch.equal(model.action_in_proj.weight, before) + + +def test_the_tied_head_follows_the_embedding(tmp_path): + # The checkpoint carries no lm_head; tying must still leave it consistent. + model = _Model() + tensors = _full_checkpoint() + pgi.load_paligemma_weights(model, _write_checkpoint(tmp_path, tensors)) + prefix = model.paligemma_with_expert.paligemma + assert torch.equal( + prefix.lm_head.weight, prefix.model.language_model.embed_tokens.weight + ) + + +def test_a_derived_buffer_is_not_required(tmp_path): + model = _Model() + pgi.load_paligemma_weights(model, _write_checkpoint(tmp_path, _full_checkpoint())) + assert torch.equal( + model.paligemma_with_expert.paligemma.model.language_model.inv_freq, + torch.ones(3), + ) + + +def test_a_partial_checkpoint_is_refused(tmp_path): + # Silently leaving part of the VLM at its random init is the failure this + # whole module exists to avoid. + tensors = _full_checkpoint() + del tensors["vision_tower.vision_model.post_layernorm.weight"] + with pytest.raises(RuntimeError, match="missing"): + pgi.load_paligemma_weights(_Model(), _write_checkpoint(tmp_path, tensors)) + + +def test_an_unknown_tensor_is_refused(tmp_path): + tensors = _full_checkpoint() + tensors["vision_tower.vision_model.renamed_upstream.weight"] = torch.zeros(3) + with pytest.raises(RuntimeError, match="unexpected"): + pgi.load_paligemma_weights(_Model(), _write_checkpoint(tmp_path, tensors)) + + +def test_an_empty_directory_is_refused(tmp_path): + with pytest.raises(FileNotFoundError, match="safetensors"): + pgi.load_paligemma_weights(_Model(), str(tmp_path)) + + +def test_a_padded_vocabulary_checkpoint_loads(tmp_path): + # The real release is 64 rows wider than openpi's table. + tensors = _full_checkpoint() + wide = torch.randn(6, 3) + tensors["language_model.model.embed_tokens.weight"] = wide + pgi.load_paligemma_weights(_Model(), _write_checkpoint(tmp_path, tensors)) + + +# ------------------------------------------------------- download / mirror path + + +def _tiny_paligemma(): + from transformers import PaliGemmaConfig, PaliGemmaForConditionalGeneration + + config = PaliGemmaConfig( + text_config={ + "model_type": "gemma", + "num_hidden_layers": 1, + "hidden_size": 32, + "intermediate_size": 32, + "num_attention_heads": 1, + "num_key_value_heads": 1, + "head_dim": 32, + "vocab_size": 300, + }, + vision_config={ + "num_hidden_layers": 1, + "hidden_size": 16, + "intermediate_size": 16, + "num_attention_heads": 1, + "image_size": 28, + "patch_size": 14, + "projection_dim": 32, + }, + image_token_index=299, + projection_dim=32, + ) + return PaliGemmaForConditionalGeneration(config) + + +def _release_key(name: str) -> str: + """The module-tree name as google/paligemma-3b-pt-224 spells it.""" + for old, new in ( + ("model.language_model.", "language_model.model."), + ("lm_head.", "language_model.lm_head."), + ("model.vision_tower.vision_model.", "vision_tower.vision_model."), + ("model.vision_tower.", "vision_tower.vision_model."), + ("model.multi_modal_projector.", "multi_modal_projector."), + ): + if name.startswith(old): + return new + name[len(old) :] + return name + + +def test_a_real_paligemma_loads_on_this_transformers(tmp_path): + """The miniature tree above is spelled the 5.x way; a real model is spelled + whichever way the installed transformers spells it, which is how a 4.53 + venv came to fail 21 SigLIP tensors while the key tests all passed.""" + source = _tiny_paligemma() + released = { + _release_key(n): p.detach().clone().contiguous() + for n, p in source.named_parameters() + } + target = _tiny_paligemma() + model = nn.Module() + model.paligemma_with_expert = nn.Module() + model.paligemma_with_expert.paligemma = target + pgi.load_paligemma_weights(model, _write_checkpoint(tmp_path, released)) + for name, param in target.named_parameters(): + assert torch.equal(param, dict(source.named_parameters())[name]), name + + +def test_a_nested_vision_tree_is_found(): + own = {"model.vision_tower.vision_model.post_layernorm.bias": None} + key = "vision_tower.vision_model.post_layernorm.bias" + assert ( + pgi._module_name(key, own) + == "model.vision_tower.vision_model.post_layernorm.bias" + ) + assert pgi._module_name(key, {}) == "model.vision_tower.post_layernorm.bias" + + +def test_a_local_directory_is_used_as_is(tmp_path): + assert pgi.resolve_paligemma_dir(str(tmp_path)) == str(tmp_path) + + +def _gated_error(repo_id): + """The 401 huggingface_hub raises for a license-gated repo.""" + import requests + from huggingface_hub.errors import GatedRepoError + + response = requests.Response() + response.status_code = 401 + return GatedRepoError(f"401 for {repo_id}", response=response) + + +def _stub_hub(monkeypatch, *, gated, digests, downloads): + import huggingface_hub + + def snapshot_download(repo_id, **_): + if repo_id in gated: + raise _gated_error(repo_id) + return downloads[repo_id] + + monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot_download) + monkeypatch.setattr(pgi, "_shard_digests", lambda repo: digests[repo]) + monkeypatch.setattr(pgi, "_verify_shards", lambda directory, expected: None) + + +def test_only_the_shards_are_downloaded(monkeypatch, tmp_path): + import huggingface_hub + + seen = [] + + def snapshot_download(repo_id, **kwargs): + seen.append(kwargs.get("allow_patterns")) + return str(tmp_path) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot_download) + pgi.resolve_paligemma_dir(CANONICAL) + assert seen == [["*.safetensors"]] + + +def test_the_canonical_repo_is_preferred(monkeypatch, tmp_path): + _stub_hub( + monkeypatch, gated=set(), digests={}, downloads={CANONICAL: str(tmp_path)} + ) + assert pgi.resolve_paligemma_dir(CANONICAL) == str(tmp_path) + + +def test_a_gated_canonical_repo_falls_back_to_a_matching_mirror(monkeypatch, tmp_path): + digests = {r: {"model.safetensors": "abc"} for r in (CANONICAL, *pgi.MIRROR_REPOS)} + _stub_hub( + monkeypatch, + gated={CANONICAL}, + digests=digests, + downloads={m: str(tmp_path) for m in pgi.MIRROR_REPOS}, + ) + assert pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=True) == str(tmp_path) + + +def test_the_mirror_is_opt_in(monkeypatch, tmp_path): + from huggingface_hub.errors import GatedRepoError + + digests = {r: {"model.safetensors": "abc"} for r in (CANONICAL, *pgi.MIRROR_REPOS)} + _stub_hub( + monkeypatch, + gated={CANONICAL}, + digests=digests, + downloads={m: str(tmp_path) for m in pgi.MIRROR_REPOS}, + ) + with pytest.raises(GatedRepoError): + pgi.resolve_paligemma_dir(CANONICAL) + + +def test_a_rate_limit_is_not_a_reason_to_use_a_mirror(monkeypatch): + import huggingface_hub + import requests + from huggingface_hub.errors import HfHubHTTPError + + response = requests.Response() + response.status_code = 429 + + def snapshot_download(repo_id, **_): + raise HfHubHTTPError("429", response=response) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot_download) + with pytest.raises(HfHubHTTPError): + pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=True) + + +def test_a_mirror_whose_shards_differ_is_rejected(monkeypatch, tmp_path): + # A mirror is transport only: different bytes means it is not PaliGemma. + digests = {CANONICAL: {"model.safetensors": "abc"}} + digests.update({m: {"model.safetensors": "TAMPERED"} for m in pgi.MIRROR_REPOS}) + _stub_hub( + monkeypatch, + gated={CANONICAL}, + digests=digests, + downloads={m: str(tmp_path) for m in pgi.MIRROR_REPOS}, + ) + with pytest.raises(RuntimeError, match="no verified mirror"): + pgi.resolve_paligemma_dir(CANONICAL, allow_mirror=True) + + +def test_a_non_canonical_repo_never_falls_back(monkeypatch, tmp_path): + # Only the gated canonical repo has verified stand-ins; anything else the + # user names must fail loudly rather than silently load other weights. + from huggingface_hub.errors import GatedRepoError + + _stub_hub( + monkeypatch, gated={"someone/private-paligemma"}, digests={}, downloads={} + ) + with pytest.raises(GatedRepoError): + pgi.resolve_paligemma_dir("someone/private-paligemma", allow_mirror=True) + + +# -------------------------------------------------------------------- the config + + +def test_the_overlay_flips_the_init_source_and_keeps_its_parent(compose_resolve): + cfg = compose_resolve( + "train_zarr_cartesian_pi", ["model=pi0.5_bc_mecka_6d_paligemma_init"] + ) + model = cfg.model.robomimic_model + assert model.config.pytorch_weight_path is None + assert model.config.paligemma_weight_path == CANONICAL + assert ( + pgi.select_init_source( + model.config.pytorch_weight_path, model.config.paligemma_weight_path + ) + == "paligemma" + ) + # the two keys are all the overlay touches: pi0.5_bc_mecka_6d's own choices + # (cartesian 6D action, no Embodiment prompt block) must survive it + assert model.ac_keys.human_bimanual == "actions_cartesian" + assert model.embodiment_label is False + assert model.config.model.action_dim == 32 + assert ( + model.action_converters.rules.HUMAN_BIMANUAL._target_ + == "egomimic.utils.action_utils.HumanBimanualCartesian6D" + ) + + +def test_the_base_pi_config_still_starts_from_pi05_base(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian_pi", ["model=pi0.5_bc_mecka_6d"]) + config = cfg.model.robomimic_model.config + assert config.paligemma_weight_path is None + assert ( + pgi.select_init_source(config.pytorch_weight_path, config.paligemma_weight_path) + == "pi05_base" + ) diff --git a/tests/unit/test_rollout_transform_mode.py b/tests/unit/test_rollout_transform_mode.py index 9442d1a062..40a19572c2 100644 --- a/tests/unit/test_rollout_transform_mode.py +++ b/tests/unit/test_rollout_transform_mode.py @@ -99,3 +99,28 @@ def test_unrecorded_legacy_checkpoint_uses_predictions_as_is(monkeypatch): monkeypatch.setattr(PolicyRollout, "_load_policy", lambda self: policy) r = PolicyRollout("both", "ckpt", 1, True) assert r.transform_mode is None and r.revert_transform_list is None + + +def test_patch_nulls_paligemma_init(tmp_path): + """The checkpoint's state_dict holds every weight, so a PaliGemma-init + policy must not fetch the gated repo on the robot.""" + import torch + + tree = { + "model": { + "robomimic_model": { + "config": { + "pytorch_weight_path": None, + "paligemma_weight_path": "google/paligemma-3b-pt-224", + } + } + } + } + path = tmp_path / "a.ckpt" + torch.save({"hyper_parameters": {"config_tree": tree}}, path) + patched = PolicyRollout._patch_checkpoint_paths(str(path)) + config = torch.load(patched, weights_only=False)["hyper_parameters"]["config_tree"][ + "model" + ]["robomimic_model"]["config"] + assert config["paligemma_weight_path"] is None + assert config["pytorch_weight_path"] is None