diff --git a/README.md b/README.md index ea399a375..f2f57eada 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 6f384bb53..27aec4ced 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 4a05b646e..283ec7b4b 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 000000000..91abc7286 --- /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 000000000..4e5610398 --- /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 000000000..ce6f5e15a --- /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 000000000..55799580e --- /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 649357b21..f148f4e58 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 6f9ed07e8..e94203b81 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 ddcd12b64..a679425c2 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/fixtures/train_harness.py b/tests/fixtures/train_harness.py index 6bdb415f8..025692925 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( diff --git a/tests/integration/test_paligemma_init_real.py b/tests/integration/test_paligemma_init_real.py new file mode 100644 index 000000000..e74135417 --- /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 0e233e51e..d35a9f18c 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 000000000..951a90fa6 --- /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 000000000..4e79db9db --- /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 9442d1a06..40a19572c 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