diff --git a/README.md b/README.md index ddf88be1..86acd4ce 100644 --- a/README.md +++ b/README.md @@ -282,7 +282,7 @@ Advanced/SO-100/internal commands stay callable directly (`config`, `calibrate`, ### Install notes -- `[monolithic]` extra is required for the cos=+1.000000 verified export path (pins transformers==5.3.0) +- `[monolithic]` extra is required for the cos=+1.000000 verified export path (pins transformers==5.5.4, the release LeRobot 0.6.1 supports) - CPU-only: `pip install 'fastcrest-tether[serve,onnx,monolithic]'` - GPU install needs the FULL cuDNN 9 system library (not just the pip wheel). Easiest path: NVIDIA's container `docker run --gpus all -it nvcr.io/nvidia/tensorrt:24.10-py3`, then `apt-get install -y clang` (for lerobot→evdev), then the pip install - `tether serve` errors loudly if cuDNN can't load — no silent CPU fallback diff --git a/docs/research/verification/lerobot-06-export-port-20261012.md b/docs/research/verification/lerobot-06-export-port-20261012.md new file mode 100644 index 00000000..f85642c6 --- /dev/null +++ b/docs/research/verification/lerobot-06-export-port-20261012.md @@ -0,0 +1,260 @@ +# LeRobot 0.6.1 export port, 2026-10-12 + +The customer is a Tether user exporting a pi0, pi0.5 or SnapFlow policy with +LeRobot 0.6.1. The required outcome is that each patched prefix, denoise and +inference path executes the same math as that version's native forward. +The checkout started clean on `feat/lerobot-0.6.1` at +`b2d447f2cb6d1841c0bbfaa97c0891e5b95595a5`, PR #355. Changes are uncommitted. +The existing failed GPU run is recorded in Studio's +`docs/research/verification/lerobot-06-parity-20261012.md`; it is reused as +failure evidence, not as successful parity evidence. + +## Source comparison and root cause + +Compared the installed 0.5.1 and 0.6.1 wheel sources in +`/Users/romirjain/dev/studio-wt-lerobot06/.sol-scratch/lerobot06/venv-051d` +and `venv-061t`, including both `pi0/modeling_pi0.py` and +`pi05/modeling_pi05.py`. The version-pin commit `d7361f8` made the old +private-method calls fail. The tiny CPU pi0 regression reproduced the exact +`AttributeError` before the port. + +0.5.1's `_prepare_attention_masks_4d` at pi0 line 623 and pi0.5 line 619 +became `lerobot.policies.common.vla_utils.prepare_attention_masks_4d` at +0.6.1 line 93. Both modeling modules import and call that shared function. +This is extraction of duplicated VLA utilities. The default computation is +unchanged: add the singleton head dimension, then select `0.0` or +`OPENPI_ATTENTION_MASK_VALUE` with `torch.where`. The optional `dtype` +conversion in 0.6.1 only runs when requested. Tether uses the same default +as native LeRobot, including float32 additive masks. + +Two other numerical changes matter: + +- Both native denoise methods now use `clone_past_key_values`, a shared + tensor-clone helper, instead of `copy.deepcopy`. The prefix cache remains + unchanged between steps. Teacher export retains its equivalent frozen + cache update. SnapFlow export now freezes updates too, because graph + tracing can evaluate a single-step forward repeatedly. SnapFlow pi0 + training uses the tensor-clone helper to preserve both the prefix cache + and gradients; pi0.5 retains its existing detached training-cache copy. +- Both native image embedding methods removed the 0.5.1 multiplication by + `sqrt(hidden_size)`. Tether's pi0 export shim still multiplied, producing + a maximum action difference of 0.5559507608 in the tiny three-step fixture. + The shared exporter now uses the native image embedding unchanged, and + the standalone export scripts no longer add that factor. + +0.6.1 also extracts the sampling loop into +`lerobot.policies.common.flow_matching.euler_integrate`. Its time schedule +and arithmetic remain `dt=-1/num_steps`, `t=1+step*dt`, `x=x+dt*v`. +The decomposed baked wrapper already performs those same operations. + +## Ported paths + +- pi0 and pi0.5 monolithic teacher denoise patches. +- Decomposed pi0.5 prefix, per-step expert and baked Euler expert, including + SnapFlow and state-out students. Student loading installs the same + export-friendly denoise patch as monolithic export. +- SnapFlow pi0, pi0.5 and state-out denoise/prefix/inference overrides. +- SnapFlow monolithic export for both families, including explicit state + forwarding for the state-out variant. +- Both SnapFlow training-prefix adapters. +- Standalone pi0/pi0.5 monolithic export scripts, layer-pruning profile, + state-out smoke and state-out stage-two distillation. Their LeRobot pins + now match the 0.6.1 helper they call. These Modal scripts were inspected + and syntax-checked, not executed on GPU. + +## Complete private LeRobot call audit + +Searched `src/` and `scripts/` for private calls and imports, followed the +receiver types, and checked each owner against the installed 0.6.1 source. +The executable regression also checks the surviving signatures and rejects +removed-mask calls or new private LeRobot imports across those directories. + +| Owner and private symbol | 0.6.1 status | Source and disposition | +| --- | --- | --- | +| `PI0Pytorch._prepare_attention_masks_4d` | Removed; all callers ported | Native shared-helper calls at `modeling_pi0.py:616,677,727` | +| `PI05Pytorch._prepare_attention_masks_4d` | Removed; all callers ported | Native shared-helper calls at `modeling_pi05.py:586,644,692` | +| `PI0Pytorch._apply_checkpoint` | Present | `modeling_pi0.py:466`, `(self, func, *args, **kwargs)` | +| `PI05Pytorch._apply_checkpoint` | Present | `modeling_pi05.py:462`, same signature | +| `PI0Policy._preprocess_images` | Present | `modeling_pi0.py:983`, `(self, batch)` returning image and mask lists | +| `PI05Policy._preprocess_images` | Present | `modeling_pi05.py:952`, same signature | + +There are no underscore-prefixed imports from LeRobot in these directories. +`config._attn_implementation` belongs to Transformers, is present, and is +assigned by LeRobot 0.6.1 itself. `lerobot.policies.pi_gemma` and its +`create_causal_mask` still exist. Python `__class__` and cache `__dict__` +accesses are not private LeRobot callable symbols. + +The original broken pi0 call came from `768a536d`, pi0.5 from `6ffeeff5`, +decomposed prefix from `bd0f7738`, and SnapFlow subclass from `0be7388f`. +The separate legacy import `lerobot.common.policies...` in +`scripts/modal_test_export.py` predates this migration and remains outside +these patched export paths. + +## CPU verification + +`tests/test_lerobot06_patched_forward.py` creates real random LeRobot models +with two language/expert layers, width 32, eight attention heads, a two-layer +SigLIP vision encoder, vocabulary 64, 16x16 images and a three-action chunk. +Only constructor configuration sizes are replaced. Embeddings, attention, +normalization, projections, caches, losses and integration use real library +code. No checkpoint download or large weights are involved. + +Native reference outputs are captured before installing export patches. +Tests compare native full training-forward losses, denoise velocities and +three-step teacher inference. SnapFlow checks zero initialization and a +nonzero learned target-time correction against unpatched LeRobot forward +methods with the student's intended suffix embedding. State-out includes +its explicit state token. Both training adapters are exercised. The +real decomposed teacher and student wrappers compare prefix caches, +per-step velocities and baked integration against native outputs. Their +exported prefix, per-step and baked graphs are executed too. Torch graphs for steps +and full image/language-to-action inference are executed and compared too. +Every float32 comparison uses `rtol=1e-6`, `atol=1e-6`, with masked cameras +and padded language tokens in a batch of two. Subprocesses isolate global +export patches from the rest of pytest. + +The requested 0.6.1 environment's actual interpreter is +`/Users/romirjain/dev/studio-wt-lerobot06/.sol-scratch/lerobot06/venv-061t/bin/python`. +It has Python 3.12.12, LeRobot 0.6.1, Transformers 5.3.0 and Torch 2.7.1. +The CPU Torch graphs are diagnostic evidence; production monolithic export +still correctly requires Torch >=2.8. No new GPU/ONNX Runtime acceptance +receipt was produced, and existing qualification verdicts are unchanged. + +```bash +PYTHONPATH="$PWD/src" \ + /Users/romirjain/dev/studio-wt-lerobot06/.sol-scratch/lerobot06/venv-061t/bin/python \ + -m pytest -q tests/test_lerobot06_patched_forward.py \ + tests/test_snapflow_backend.py tests/test_monolithic_causal_mask_shim.py + +PYTHONPATH="$PWD/src:/Users/romirjain/dev/.sol-out/tether-mcpdeps:/Users/romirjain/dev/.sol-out/tether-asyncdeps" \ + /Users/romirjain/dev/studio-wt-s11/.sol-scratch/venv/bin/python -m pytest +``` + +The final focused run in `venv-061t` passed: **48 passed in 71.20 seconds**, +including all seven new native/graph/private-API cases, the existing SnapFlow +backend tests and causal-mask tests. Raw output is +`/tmp/tether-lerobot06-final-focused.log`. + +The full requested suite ran in 206.12 seconds: **4,169 passed, 133 skipped, +21 failed and 10 errors**. All 31 failures/errors are sandbox-denied TCP +socket binds: two fake HTTP servers and 29 ZeroMQ cases. Each reports +`Operation not permitted`; there were no other failures. The zero-failure +full-suite target therefore requires an environment permitted to bind local +TCP sockets. No test was weakened or skipped to conceal those results. +Raw output is `/tmp/tether-lerobot06-full-suite.log`. + +After the final loader and test changes, the affected tests in the S11 +interpreter passed: **56 passed, 11 skipped**. LeRobot-specific tests skip +there because that environment lacks LeRobot; they execute in `venv-061t`. +The affected core exporter/distillation files and new regression pass Ruff, +and all edited Python files pass Ruff's undefined-name check. `git diff +--check` is clean. + +## Adversarial review (Claude, 2026-10-12) + +- Line-by-line against the LeRobot 0.6.1 wheel (`venv-061t`): `vla_utils.prepare_attention_masks_4d` + (`vla_utils.py:93-102`) is called with no `dtype`, as native does at `modeling_pi0.py:677,727` + and `modeling_pi05.py:692`. The patched pi0/pi0.5 denoise steps (`monolithic.py` pi0 patch, + `_apply_pi05_denoise_step_patch`) match native `denoise_step` (`modeling_pi0.py:705-743`, + `modeling_pi05.py:671-708`) term by term; the two deliberate differences are (a) `F.pad`+AND + instead of `torch.cat` for the 2D mask, equal by construction when no pad deficit, and (b) a + frozen `DynamicLayer.update` instead of `clone_past_key_values`; both feed attention + `cat(prefix_kv, suffix_kv)` and leave the prefix cache unchanged. The decomposed prefix wrapper + matches native `sample_actions` prefix (`modeling_pi05.py:640-656`). 0.6.1 `embed_image` + (`modeling_pi0.py:287-296`) has no `sqrt(hidden_size)` factor (0.5.1 had it at `:446`), so + dropping Tether's shim is correct. `euler_integrate` (`flow_matching.py:93-120`) is + `dt=-1/n, x+=dt*v`. VERIFIED. +- The CPU references are native: `native_step = cls.denoise_step`, `expected_*` from a + `copy.deepcopy` of the model, all computed before `apply_export_patches()` and the SnapFlow + installers run. SnapFlow references call the unpatched class `denoise_step` with only the + student's suffix embedding substituted. VERIFIED (`tests/test_lerobot06_patched_forward.py:136-232`). +- Mutation checks (file restored by copy, shasum `7d653f5b...aef49` equal before and after each): + M1 disable the denoise-phase cache freeze in the pi0 patch: 6 of 7 cases fail (pi0: 3-step + actions max abs 0.118 at line 249). M2 make the pi0.5 patch ignore the prefix pad mask: pi05 + and decomposed fail (velocity max abs 0.012 at line 244). Both mutations killed. VERIFIED. +- Re-run after review: 48 passed in 77.7 s (`venv-061t`). VERIFIED. + +## Correction 2026-10-12 (Claude, after the GPU re-run): native 0.6.1 is mis-scaled on transformers 5.3 + +The claim above that dropping the `embed_image` sqrt(hidden_size) shim is correct holds only +for transformers >=5.4. It is **wrong for the pinned transformers 5.3.0**. Evidence: + +- LeRobot 0.6.1 declares `transformers>=5.4.0,<5.6.0` (`lerobot-0.6.1.dist-info/METADATA:79`). +- transformers 5.3.0 `PaliGemmaModel.get_image_features` divides by sqrt(hidden_size) + (`modeling_paligemma.py:273`); 5.5.4 does not. LeRobot 0.5.1 multiplied back + (`modeling_pi0.py:446`), 0.6.1 does not (`modeling_pi0.py:287-296`). +- Language embeddings likewise: 0.5.1 `lang_emb * sqrt(dim)` (`modeling_pi0.py:669`); 0.6.1 + relies on 5.4+'s `GemmaTextScaledWordEmbedding` (5.5.4 `modeling_gemma.py:50-61`), which 5.3.0 + lacks, and `PiGemmaModel.forward` (`pi_gemma.py:250-290`) has no normalizer. +- GPU re-run at `5fb92df` (Studio `docs/research/verification/lerobot-06-parity-rerun-20261012.md`): + pi0 export vs the receipt harness reference, cosine 0.465, max_abs 0.403, **failed**. The + harness (`scripts/local_pi0_monolithic_parity.py:39-56`, `local_pi05_...:48`) still applies the + sqrt shim to the reference, the exporter no longer does. + +The CPU tests compare against native 0.6.1 under 5.3, which is itself mis-scaled, so they +passed. The mask port itself stands (review above). Not merged. Fix requires an owner +decision: move transformers to 5.4/5.5 as LeRobot 0.6.1 requires (re-qualify the exporters), or +keep 5.3 and restore both openpi scalings for transformers <5.4. + +## Resolution 2026-10-12 (Claude): transformers 5.5.4, independent golden reference + +Owner decision: move to the transformers release LeRobot 0.6.1 supports. Latest 5.5.x on PyPI +is 5.5.4 (5.6.x is outside `<5.6.0`). Changes on this branch: + +- Pins: `[monolithic]`, `[native]`, `[eval-local]`, `[jetson]` -> `transformers==5.5.4`; base + `transformers>=4.40,<5.6`; the `[tool.uv]` 5.3.0 override is removed (LeRobot's own range now + holds). `QUALIFIED_TRANSFORMERS_VERSION` (`src/tether/_lerobot_version.py`) drives the exporter + gate, which now refuses 5.3.0 by name (`tests/test_cli_export_end_to_end.py`). Every extra + resolves with `uv pip compile` for linux x86_64 and darwin arm64 without an override; the + dragged pins are tokenizers 0.22.2, huggingface-hub 1.33.0, safetensors 0.8.0, accelerate 1.15.0 + (within transformers 5.5.4's `tokenizers>=0.22,<=0.23`, `huggingface-hub>=1.5,<2.0`). VERIFIED. +- Receipt harnesses (`scripts/local_pi0_monolithic_parity.py`, `local_pi05_...`) drop the + `embed_image` sqrt shim and the `create_causal_mask` kwarg shim: the reference is native + LeRobot 0.6.1 under 5.5.4. +- Tether's own PaliGemma text embedding (`PaliGemmaBackbone.embed_text_openpi`, used by + `models/vlas/pi0.py`, `pi05.py`, `runtime/fast_inference/pi05.py`) now looks the table up + directly. 5.5.4's `GemmaTextScaledWordEmbedding` already multiplies by sqrt(hidden) + (`modeling_gemma.py:50-61,381-382`), so calling the module and scaling again would have + scaled text by `hidden` in the Triton fast path on 5.5.4. VERIFIED by source; regression test + `tests/test_paligemma_backbone.py::test_embed_text_openpi_is_independent_of_embedding_module_scaling`. +- Parity vehicles: `transformers==5.5.4`, `retries=0`, phase subprocess timeouts instead of the + SIGKILL watchdog, restart guard, provenance-checked `--reuse-export`, pi0.5 on A100-40GB. + +### Independent reference (golden) + +`tests/_lerobot_pi_golden.py` + `tests/fixtures/lerobot_pi_golden/{pi0,pi05}.pt`: tiny random +pi0/pi0.5 (fixed seed, weights stored in the fixture), captured on the qualified OLD stack +(LeRobot 0.5.1 + transformers 5.3.0 + torch 2.7.1 + Tether `cdb94dc`). The new stack is compared +to it, not to itself. Results, float32, B=2 with masked cameras and padded language: + +| Stack | prefix_embs max_abs | actions (10 steps) max_abs | patched export forward | verdict | +|---|---|---|---|---| +| 0.6.1 + 5.5.4 (venv-061) pi0 | 2.98e-08 | 1.79e-07 | 1.79e-07 | pass | +| 0.6.1 + 5.5.4 pi0.5 | 2.98e-08 | 2.38e-07 | 2.38e-07 | pass | +| 0.6.1 + 5.3.0 (venv-061t) pi0, negative control | 0.318, norm ratio 0.176777 = 1/sqrt(32) | 0.396 | 0.396 | fail | +| 0.6.1 + 5.3.0 pi0.5, negative control | same | 0.0175 (cosine 0.99998) | same | fail | + +Expected difference: a few float32 ulps, not bit identity. The old stack divides image features +by sqrt(hidden) inside transformers 5.3 and LeRobot 0.5.1 multiplies back (`modeling_pi0.py:446`); +5.4+ does neither, and language scaling moved from LeRobot (`modeling_pi0.py:669`, 0.5.1) into +`GemmaTextScaledWordEmbedding`. Tolerance `rtol=atol=1e-5`. Note the pi0.5 negative control has +action cosine 0.99998, which a cosine >= 0.999 gate would pass: the prefix-embedding check +and the absolute tolerance are what catch the bug at tiny width. + +The same fixture drives the real monolithic toolchain (`apply_export_patches`, +`torch_export_patches(patch_transformers=True)`, `torch.export`, `torch.onnx.export` opset 19, +Where-dtype fix, ONNX Runtime CPU) on torch 2.11.0 + transformers 5.5.4 + onnx-diagnostic 0.9.3: +pi0 max_abs 1.25e-06, pi0.5 1.19e-07 against the OLD-stack golden at num_steps=10. The ADR-era +"5.4+ q_length regression" did not reproduce on this toolchain. VERIFIED (CPU, tiny model); +GPU receipts on the real checkpoints decide. + +### Suites + +- S11 interpreter, full suite: 4204 passed, 138 skipped, 0 failed. +- LeRobot 0.6.1 + transformers 5.5.4 + torch 2.7.1 (Studio eval runtime, `venv-061`), full suite: + 4262 passed, 80 skipped, 0 failed. +- torch 2.11.0 + transformers 5.5.4 export venv, LeRobot/export subset (golden, patched-forward, + CLI export, causal-mask shim, SnapFlow backend, Pad-free mask): 73 passed, 0 failed. + +Raw logs: Studio `.sol-scratch/lerobot06/tf55/`; probes and record: Studio +`docs/research/verification/lerobot-06-transformers55-20261012/`. diff --git a/pyproject.toml b/pyproject.toml index cab571f8..15067e64 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ # compiled against numpy 1.x ABI; numpy 2.x's ABI break corrupts ORT + # JAX wheels on aarch64. On x86_64 (Modal A100, customer desktops) + # arm64 (M-series Mac dev hosts) numpy 2.x is fine, AND the - # [monolithic] extra's lerobot==0.5.1 dep REQUIRES numpy>=2.0 — pinning + # [monolithic] extra's lerobot==0.6.1 dep REQUIRES numpy>=2.0 — pinning # numpy<2 globally makes lerobot uninstallable on Modal (caught by the # 2026-05-20 lift #4 smoke fire). Split via PEP-508 platform marker. "numpy>=1.24.0,<2.0; platform_machine == 'aarch64'", @@ -39,12 +39,11 @@ dependencies = [ # loading for the auto-detect path. Apr-14 install-path verification caught # that omitting these breaks `tether export` even with [serve,gpu]. "huggingface_hub>=0.20.0", - # Base allows 4.40 → 5.3.x. [monolithic] extra pins ==5.3.0 exactly (5.4+ has - # a q_length regression that breaks onnx-diagnostic patches). Without the - # upper bound, pip would pull 5.4/5.5 here and conflict with [monolithic]. - # Without raising the upper bound past 5.0, [monolithic] can't install - # alongside the base (the v0.2.1 customer-dogfood finding). - "transformers>=4.40,<5.4", + # Base allows 4.40 → 5.5.x. [monolithic] pins ==5.5.4 exactly, inside + # LeRobot 0.6.1's declared >=5.4,<5.6 (on 5.3 its pi0/pi0.5 embeddings are + # mis-scaled; Studio ADR 0003 correction 2026-10-12). The upper bound keeps + # pip from pulling an unqualified 5.6+ here. + "transformers>=4.40,<5.6", "onnx>=1.15.0", "onnxscript>=0.1.0", # TetherClient (Phase 0.5 customer-sdk) — sync + async HTTP client. @@ -141,32 +140,42 @@ comply = [] # gate in GOALS.yaml. Heavy (pulls torch, torchvision, accelerate, etc.) — # that's why it's a separate extra, not a base dep. # -# lerobot 0.5.x requires Python >=3.12. Jetson (JetPack 6) ships 3.10, +# lerobot 0.6.x requires Python >=3.12. Jetson (JetPack 6) ships 3.10, # so we gate lerobot behind a version marker — pip skips it on 3.10/3.11 # instead of hard-failing with ResolutionImpossible. Native parity # verification should be run on a Python 3.12+ host. native = [ - "lerobot==0.5.1; python_version >= '3.12'", + "transformers==5.5.4", + "lerobot==0.6.1; python_version >= '3.12'", + "draccus==0.11.6; python_version >= '3.12'", "num2words", + "accelerate>=1.14,<2", + "diffusers==0.39.0", ] # Monolithic ONNX export — cos=1.0 verified production path. Runs locally # via `tether export --monolithic` (the DEFAULT as of v0.2.1). Requires -# transformers==5.3.0 EXACTLY (5.4+ has a q_length regression in -# masking_utils.sdpa_mask that breaks onnx-diagnostic patches). Base pin -# `transformers>=4.40,<5.4` permits this extra to pin-to-exact. +# transformers==5.5.4 EXACTLY: LeRobot 0.6.1 needs >=5.4,<5.6, and the +# exporter is qualified on this release (tests/_lerobot_pi_golden.py runs the +# full onnx-diagnostic export against the pre-migration golden). # -# lerobot 0.5.x requires Python >=3.12. On Jetson (Python 3.10) the +# lerobot 0.6.x requires Python >=3.12. On Jetson (Python 3.10) the # monolithic export can't run locally — export on a cloud/desktop host # (Python 3.12+) and `tether serve` the resulting ONNX on the Jetson # with `pip install 'fastcrest-tether[serve]'` only. monolithic = [ - "lerobot==0.5.1; python_version >= '3.12'", - "transformers==5.3.0", + "lerobot==0.6.1; python_version >= '3.12'", + "draccus==0.11.6; python_version >= '3.12'", + "transformers==5.5.4", + # onnx-diagnostic's export patches reference a torch 2.8+ symbol; + # _require_monolithic_deps refuses older torch with the same floor. + "torch>=2.8", "onnx-diagnostic>=0.9", "onnxscript>=0.1", "optree", "scipy", "num2words", + "accelerate>=1.14,<2", + "diffusers==0.39.0", ] # Task-success evaluation (heavy sim deps). Gated behind this extra so the # base install stays light. Use with `tether bench --benchmark libero_10`. @@ -177,8 +186,14 @@ eval = [ "gymnasium", "h5py", ] +# Spell out the SmolVLA dependencies (LeRobot's transformers-dep extra, +# >=5.4,<5.6) at the exact qualified transformers. eval-local = [ - "lerobot[smolvla]==0.5.1; python_version >= '3.12' and sys_platform == 'linux'", + "transformers==5.5.4; python_version >= '3.12' and sys_platform == 'linux'", + "num2words>=0.5.14,<0.6; python_version >= '3.12' and sys_platform == 'linux'", + "accelerate>=1.14,<2; python_version >= '3.12' and sys_platform == 'linux'", + "lerobot==0.6.1; python_version >= '3.12' and sys_platform == 'linux'", + "draccus==0.11.6; python_version >= '3.12' and sys_platform == 'linux'", "peft; python_version >= '3.12' and sys_platform == 'linux'", "mujoco==3.3.2; sys_platform == 'linux'", "robosuite==1.4.1; sys_platform == 'linux'", @@ -212,13 +227,14 @@ tracing = [ # pip install 'fastcrest-tether[lerobot]' # tether serve bundle/ --embodiment so_arm100 --port /dev/ttyUSB0 # -# lerobot 0.5.x requires Python >=3.12 (we skip on 3.10/3.11 to avoid pip's +# lerobot 0.6.x requires Python >=3.12 (we skip on 3.10/3.11 to avoid pip's # ResolutionImpossible; users on JetPack 6 Python 3.10 can still construct an # adapter + run the conversion math, they just can't open a serial bus until # they upgrade to Python 3.12 or install scservo_sdk directly via the [so100] # extra). lerobot = [ - "lerobot>=0.5.1; python_version >= '3.12'", + "lerobot==0.6.1; python_version >= '3.12'", + "draccus==0.11.6; python_version >= '3.12'", ] # SO-100 hardware extra — installs the scservo_sdk Python package on the # Raspberry Pi / host wired to the arm. Separate from [lerobot] because the @@ -234,9 +250,10 @@ so100 = [ # RTC-enabled deployments. Skeleton + config validation work without this # extra; only the actual processor construction needs it. # -# lerobot 0.5.x requires Python >=3.12 (skipped on 3.10/3.11). +# lerobot 0.6.x requires Python >=3.12 (skipped on 3.10/3.11). rtc = [ - "lerobot==0.5.1; python_version >= '3.12'", + "lerobot==0.6.1; python_version >= '3.12'", + "draccus==0.11.6; python_version >= '3.12'", ] # Curate format-converter extras. Phase 1 ships LeRobot v3 + HDF5 with the # core install (parquet via pyarrow which huggingface_hub already pulls; h5py @@ -270,7 +287,7 @@ curate-rlds = [ # Jetson / ARM64 edge stack. # - Skips [gpu] x86_64 wheels (onnxruntime-gpu PyPI has no aarch64 builds) -# - Skips lerobot (0.5.1 requires Python 3.12; JetPack 6 ships 3.10) +# - Skips lerobot (0.6.1 requires Python 3.12; JetPack 6 ships 3.10) # - User must pre-install Jetson Zoo onnxruntime-gpu from pypi.jetson-ai-lab.io # and mount host CUDA via --runtime=nvidia. jetson = [ @@ -281,7 +298,7 @@ jetson = [ "jsonschema>=4.0", "prometheus-client>=0.19.0", # monolithic export stack minus lerobot - "transformers==5.3.0", + "transformers==5.5.4", "onnx-diagnostic>=0.9", "onnxscript>=0.1", "optree", diff --git a/scripts/local_pi05_monolithic_parity.py b/scripts/local_pi05_monolithic_parity.py index b802f513..3a38ca26 100644 --- a/scripts/local_pi05_monolithic_parity.py +++ b/scripts/local_pi05_monolithic_parity.py @@ -34,36 +34,6 @@ sys.modules[_mod] = _stub -def _patch_pi05_for_transformers() -> None: - from lerobot.policies.pi05 import modeling_pi05 - from lerobot.policies import pi_gemma - from transformers import masking_utils - - def patched_embed_image(self, image): - out_dtype = image.dtype - if image.dtype != torch.float32: - image = image.to(torch.float32) - output = self.paligemma.model.get_image_features(image) - features = output.pooler_output if hasattr(output, "pooler_output") else output - features = features * self.paligemma.config.text_config.hidden_size**0.5 - if features.dtype != out_dtype: - features = features.to(out_dtype) - return features - - modeling_pi05.PaliGemmaWithExpertModel.embed_image = patched_embed_image - - original = masking_utils.create_causal_mask - - def causal_mask_shim(*args, **kwargs): - if "inputs_embeds" in kwargs and "input_embeds" not in kwargs: - kwargs["input_embeds"] = kwargs.pop("inputs_embeds") - return original(*args, **kwargs) - - masking_utils.create_causal_mask = causal_mask_shim - if hasattr(pi_gemma, "create_causal_mask"): - pi_gemma.create_causal_mask = causal_mask_shim - - def _repo_commit() -> str: root = Path(__file__).resolve().parents[1] try: @@ -160,7 +130,8 @@ def main(argv: list[str] | None = None) -> int: if args.num_steps <= 0: raise ValueError("--num-steps must be positive.") - _patch_pi05_for_transformers() + # Reference is native LeRobot 0.6.1 under the qualified transformers + # (5.5.4): no embedding or mask shims (Studio ADR 0003 correction 2026-10-12). import onnxruntime as ort from huggingface_hub import snapshot_download diff --git a/scripts/local_pi0_monolithic_parity.py b/scripts/local_pi0_monolithic_parity.py index 871e3f71..052cb3e3 100644 --- a/scripts/local_pi0_monolithic_parity.py +++ b/scripts/local_pi0_monolithic_parity.py @@ -35,44 +35,6 @@ sys.modules[_mod] = _stub -def _patch_pi0_for_transformers_457() -> None: - from lerobot.policies.pi0 import modeling_pi0 - - def patched_embed_image(self, image): - out_dtype = image.dtype - if image.dtype != torch.float32: - image = image.to(torch.float32) - image_outputs = self.paligemma.model.get_image_features(image) - features = ( - image_outputs.pooler_output - if hasattr(image_outputs, "pooler_output") - else image_outputs - ) - features = features * self.paligemma.config.text_config.hidden_size**0.5 - if features.dtype != out_dtype: - features = features.to(out_dtype) - return features - - modeling_pi0.PaliGemmaWithExpertModel.embed_image = patched_embed_image - - -def _patch_create_causal_mask_kwarg() -> None: - from transformers import masking_utils - - original = masking_utils.create_causal_mask - - def shim(*args, **kwargs): - if "inputs_embeds" in kwargs and "input_embeds" not in kwargs: - kwargs["input_embeds"] = kwargs.pop("inputs_embeds") - return original(*args, **kwargs) - - masking_utils.create_causal_mask = shim - from lerobot.policies import pi_gemma - - if hasattr(pi_gemma, "create_causal_mask"): - pi_gemma.create_causal_mask = shim - - def _repo_commit() -> str: root = Path(__file__).resolve().parents[1] try: @@ -167,8 +129,8 @@ def main(argv: list[str] | None = None) -> int: if args.num_steps <= 0: raise ValueError("--num-steps must be positive.") - _patch_pi0_for_transformers_457() - _patch_create_causal_mask_kwarg() + # Reference is native LeRobot 0.6.1 under the qualified transformers + # (5.5.4): no embedding or mask shims (Studio ADR 0003 correction 2026-10-12). import onnxruntime as ort from huggingface_hub import snapshot_download diff --git a/scripts/modal_gr00t_lerobot_probe.py b/scripts/modal_gr00t_lerobot_probe.py index 15b75f47..34402a28 100644 --- a/scripts/modal_gr00t_lerobot_probe.py +++ b/scripts/modal_gr00t_lerobot_probe.py @@ -1,6 +1,6 @@ -"""Quick probe — what does lerobot 0.5.1 actually expose for GR00T? +"""Quick probe — what does lerobot 0.6.1 actually expose for GR00T? -lerobot-0.5.1 on Modal may or may not have GR00T support. The cloned +lerobot-0.6.1 on Modal may or may not have GR00T support. The cloned local lerobot is `main` (newer). Need to reconcile. Prints the actual module contents of lerobot's groot subpackage. @@ -23,7 +23,7 @@ "numpy", "safetensors>=0.4.0", "huggingface_hub", - "lerobot==0.5.1", + "lerobot==0.6.1", "num2words", ) ) @@ -52,7 +52,7 @@ def probe(): # Attempt to import groot_n1 + list symbols try: - from lerobot.policies.groot import groot_n1 + from lerobot.policies.groot import groot_n1_7 as groot_n1 print(f"\ngroot_n1 symbols (top-level, non-dunder):") for name in sorted(dir(groot_n1)): if not name.startswith("_"): @@ -64,7 +64,7 @@ def probe(): # Attempt action_head try: - from lerobot.policies.groot.action_head import flow_matching_action_head as fma + from lerobot.policies.groot.action_head import cross_attention_dit as fma print(f"\nflow_matching_action_head symbols:") for name in sorted(dir(fma)): if not name.startswith("_"): diff --git a/scripts/modal_groot_export_parity_pinned.py b/scripts/modal_groot_export_parity_pinned.py index 05d30434..2783269c 100644 --- a/scripts/modal_groot_export_parity_pinned.py +++ b/scripts/modal_groot_export_parity_pinned.py @@ -46,7 +46,7 @@ ) HF_CACHE_PATH = "/root/.cache/huggingface" -PARITY_OUT_PATH = "/parity_out" +PARITY_OUT_PATH = "/parity_out/lerobot-061-20261011" def _hf_secret(): @@ -63,24 +63,26 @@ def _hf_secret(): modal.Image.debian_slim(python_version="3.12") .apt_install("git") .pip_install( - "torch", + # torch>=2.8: onnx-diagnostic 0.9.x torch_export_patches needs it (see pi0 vehicle). + "torch>=2.8,<2.12", "safetensors>=0.4.0", "huggingface_hub", - # transformers==5.3.0 EXACTLY: the monolithic exporter refuses + # transformers==5.5.4 EXACTLY: the monolithic exporter refuses # anything else. - "transformers==5.3.0", + "transformers==5.5.4", "onnx>=1.16", "onnxruntime>=1.20", "onnxscript>=0.1", # monolithic dep group (export_gr00t_monolithic gates on these): - "lerobot==0.5.1", + "lerobot==0.6.1", "num2words", + "diffusers==0.39.0", "onnx-diagnostic>=0.9", "optree", "scipy", "accelerate", - "draccus", - "numpy", + "draccus==0.11.6", + "numpy==2.2.6", "Pillow", "typer", "rich", @@ -122,7 +124,7 @@ def _hf_secret(): timeout=3600, volumes={ HF_CACHE_PATH: _hf_cache_volume, - PARITY_OUT_PATH: _parity_out_volume, + "/parity_out": _parity_out_volume, }, secrets=[_hf_secret()], ) @@ -130,7 +132,7 @@ def run_export_parity( model_revision: str = GROOT_REVISION, tether_revision: str = TETHER_REVISION, embodiment_id: int = 0, - reuse_export: bool = True, + reuse_export: bool = False, watchdog_seconds: float = 3300, ) -> dict: import subprocess @@ -322,8 +324,13 @@ def _progress(stage: str) -> None: @app.local_entrypoint() -def main() -> None: - result = run_export_parity.remote() +def main( + tether_revision: str = TETHER_REVISION, + watchdog_seconds: float = 1800, +) -> None: + result = run_export_parity.remote( + tether_revision=tether_revision, watchdog_seconds=watchdog_seconds + ) print("\n=== RESULT ===") for key, value in result.items(): print(f" {key}: {value}") diff --git a/scripts/modal_layer_pruning_profile.py b/scripts/modal_layer_pruning_profile.py index 024cbd73..0e6f910e 100644 --- a/scripts/modal_layer_pruning_profile.py +++ b/scripts/modal_layer_pruning_profile.py @@ -65,8 +65,8 @@ def _build_bust() -> str: modal.Image.debian_slim(python_version="3.12") .apt_install("git") .pip_install( - "lerobot==0.5.1", - "transformers==5.3.0", + "lerobot==0.6.1", + "transformers==5.5.4", "num2words", "safetensors>=0.4.0", "scipy", @@ -164,6 +164,7 @@ def hook(_module, _inp, output): images = [img, img, img] img_masks = [mask, mask, mask] from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d prefix_embs, prefix_pad_masks, prefix_att_masks = ( policy.model.embed_prefix( images, img_masks, lang_tokens, lang_masks, @@ -172,7 +173,7 @@ def hook(_module, _inp, output): prefix_att_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 prefix_att_2d_4d = ( - policy.model._prepare_attention_masks_4d(prefix_att_2d) + prepare_attention_masks_4d(prefix_att_2d) ) language_model.config._attn_implementation = "eager" _ = policy.model.paligemma_with_expert.forward( diff --git a/scripts/modal_pi05_export_parity_pinned.py b/scripts/modal_pi05_export_parity_pinned.py index c54ed686..9445ad01 100644 --- a/scripts/modal_pi05_export_parity_pinned.py +++ b/scripts/modal_pi05_export_parity_pinned.py @@ -2,7 +2,7 @@ Runs the receipt-grade harness `scripts/local_pi05_monolithic_parity.py` on Modal Linux + CUDA, following the proven pi0 pattern (exact 40-hex model -commit pin, no floating HEAD, plus an in-function self-kill watchdog): +commit pin, no floating HEAD, bounded GPU time): - `lerobot/pi05_base` is pinned to PI05_BASE_REVISION, the last revision before "Add relative action processor steps (#6)" (2026-06-03) added a @@ -14,9 +14,8 @@ - The Tether implementation under test is pinned to TETHER_REVISION (an exact 40-hex commit): the container clones the repo and checks out that commit, and the harness records `git rev-parse HEAD` from that checkout. -- A self-kill wall-clock watchdog bounds GPU burn: `timeout=` on the - decorator is the platform's promise; a hung download or stuck CUDA call - can sit in the container without tripping it. +- GPU burn is bounded by per-phase subprocess timeouts, the Modal function + timeout and retries=0, plus a restart guard (see EXPORT_TIMEOUT_SECONDS). Usage: modal run scripts/modal_pi05_export_parity_pinned.py @@ -47,7 +46,7 @@ ) HF_CACHE_PATH = "/root/.cache/huggingface" -PARITY_OUT_PATH = "/parity_out" +PARITY_OUT_PATH = "/parity_out/lerobot-061-tf554-20261012" def _hf_secret(): @@ -61,9 +60,14 @@ def _hf_secret(): modal.Image.debian_slim(python_version="3.12") .apt_install("git") .pip_install( - # Install lerobot first — it pins torch>=2.7 and hub>=1.0 - "lerobot==0.5.1", + # torch>=2.8: every onnx-diagnostic 0.9.x patches + # torch._export.non_strict_utils._constrain_user_specified_dimhint_range, + # which torch 2.7.1 lacks, so torch_export_patches() raises before export. + # lerobot 0.6.1 caps torch <2.12. (Studio's 2.7.1 pin is an MPS eval pin.) + "torch>=2.8,<2.12", + "lerobot==0.6.1", "num2words", + "diffusers==0.39.0", "safetensors>=0.4.0", "onnx>=1.16", "onnxruntime>=1.20", @@ -71,13 +75,13 @@ def _hf_secret(): "onnx-diagnostic>=0.9", "optree", # onnx-diagnostic soft-dep "scipy", # onnx-diagnostic patches_transformers_qwen2_5 transitively - "numpy", + "numpy==2.2.6", "accelerate", - "draccus", + "draccus==0.11.6", ) - # transformers==5.3.0 EXACTLY: the monolithic exporter refuses anything - # else (5.4+ has a q_length regression in masking_utils.sdpa_mask). - .pip_install("transformers==5.3.0") + # transformers==5.5.4 EXACTLY: LeRobot 0.6.1 needs >=5.4,<5.6 and the + # monolithic exporter gate refuses anything but the qualified release. + .pip_install("transformers==5.5.4") .env({ "HF_HOME": HF_CACHE_PATH, "TRANSFORMERS_CACHE": f"{HF_CACHE_PATH}/transformers", @@ -113,13 +117,30 @@ def _hf_secret(): ) +# Phase budgets from measured runs: pi0 fresh export 1776.5 s and harness +# incl. setup 274.9 s on A100-40GB (Studio +# docs/research/verification/lerobot-06-parity-rerun-20261012.md); pi0.5 +# fresh export plus failed session about 2190 s on A10G +# (docs/parity/parity-pi05-refix-report.md). About 1.5x and 3x margin. Each phase runs in a subprocess with this timeout, so an +# overrun ends the input as an ordinary failure. The previous in-container +# SIGKILL watchdog surfaced as a runner crash, which Modal treats as an +# internal failure and re-schedules regardless of the retry policy. +EXPORT_TIMEOUT_SECONDS = 2700 +PARITY_TIMEOUT_SECONDS = 900 +SETUP_MARGIN_SECONDS = 600 +MODAL_TIMEOUT_SECONDS = EXPORT_TIMEOUT_SECONDS + PARITY_TIMEOUT_SECONDS + SETUP_MARGIN_SECONDS + + @app.function( image=gpu_image, - gpu="A10G", - timeout=3600, + # A100-40GB: on A10G the fp32 reference and the ~13 GB ORT session do + # not co-reside (docs/parity/parity-pi05-refix-report.md, run 1 OOM). + gpu="A100-40GB", + timeout=MODAL_TIMEOUT_SECONDS, + retries=0, volumes={ HF_CACHE_PATH: _hf_cache_volume, - PARITY_OUT_PATH: _parity_out_volume, + "/parity_out": _parity_out_volume, }, secrets=[_hf_secret()], ) @@ -127,40 +148,25 @@ def run_export_parity( model_revision: str = PI05_BASE_REVISION, tether_revision: str = TETHER_REVISION, num_steps: int = 10, - reuse_export: bool = True, - watchdog_seconds: float = 3300, + reuse_export: bool = False, + export_timeout_seconds: float = EXPORT_TIMEOUT_SECONDS, + parity_timeout_seconds: float = PARITY_TIMEOUT_SECONDS, + run_id: str = "", ) -> dict: import subprocess import sys import time from pathlib import Path + import json as _json import os as _os - import signal as _signal - import threading as _threading - - _WATCHDOG_SECONDS = float( - _os.environ.get("PARITY_WATCHDOG_SECONDS", str(watchdog_seconds)) - ) - - def _self_kill() -> None: - print( - f"[watchdog] deadline of {_WATCHDOG_SECONDS:.0f}s reached; killing this " - "container rather than burning GPU time on a hung step.", - flush=True, - ) - _os.kill(_os.getpid(), _signal.SIGKILL) - - _watchdog = _threading.Timer(_WATCHDOG_SECONDS, _self_kill) - _watchdog.daemon = True - _watchdog.start() _started_at = time.monotonic() def _progress(stage: str) -> None: print(f"[progress] {time.monotonic() - _started_at:7.1f}s {stage}", flush=True) - _progress("container up, watchdog armed") + _progress("container up") summary: dict = { "model_source": "lerobot/pi05_base", "model_revision": model_revision, @@ -169,6 +175,20 @@ def _progress(stage: str) -> None: "num_steps": num_steps, } + # Restart guard: if Modal re-schedules this input after a container + # crash, the second attempt finds its own marker and exits at once + # instead of starting another GPU export. + if run_id: + marker = Path(PARITY_OUT_PATH) / "pi05" / "attempts" / run_id + _parity_out_volume.reload() + if marker.exists(): + _progress("restarted attempt detected; refusing to run again") + summary["verdict"] = "aborted-restarted-attempt" + return summary + marker.parent.mkdir(parents=True, exist_ok=True) + marker.write_text(str(time.time())) + _parity_out_volume.commit() + # Fail fast if the CUDA EP cannot be created (a broken CUDA/cuDNN # image otherwise surfaces 10 minutes later as a no-fallback gate # failure after the reference has already run). @@ -236,46 +256,58 @@ def _progress(stage: str) -> None: snap = snapshot_download("lerobot/pi05_base", revision=model_revision) summary["snapshot_dir"] = snap - from tether.exporters.monolithic import export_pi05_monolithic - export_dir = Path(PARITY_OUT_PATH) / "pi05" / "export" - existing = export_dir / "model.onnx" - # model.onnx may be graph-only with weights in model.onnx.data: count - # the whole artifact dir, which is also what the receipt hashes. - existing_bytes = ( - sum(f.stat().st_size for f in export_dir.glob("*") if f.is_file()) - if existing.is_file() - else 0 + provenance_path = export_dir / "tether-export-provenance.json" + harness_env = dict(_os.environ) + harness_env["PYTHONPATH"] = str(pin_dir / "src") + _os.pathsep + harness_env.get( + "PYTHONPATH", "" ) - if reuse_export and existing_bytes > 10**9: - export_result = { - "status": "reused", - "onnx_path": str(existing), - "size_mb": existing_bytes / 1e6, - "num_steps": num_steps, - } - _progress(f"reusing retained export ({export_result['size_mb']:.1f}MB)") + from importlib.metadata import version as _v + + provenance = { + "tether_revision": tether_revision, + "model_revision": model_revision, + "num_steps": num_steps, + **{p: _v(p) for p in ("lerobot", "transformers", "torch", "onnx-diagnostic", "onnxscript")}, + } + summary["export_provenance"] = provenance + retained = ( + _json.loads(provenance_path.read_text()) if provenance_path.is_file() else None + ) + if reuse_export and retained == provenance: + export_result = {"status": "reused", "onnx_path": str(export_dir / "model.onnx")} + _progress("reusing retained export (provenance matches)") else: - _progress("export_pi05_monolithic from pinned snapshot") - export_result = export_pi05_monolithic(snap, export_dir, num_steps=num_steps) + if reuse_export: + _progress(f"not reusing: retained provenance {retained} != {provenance}") + _progress("export_pi05_monolithic from pinned snapshot (subprocess)") + code = ( + "import json, sys; from tether.exporters.monolithic import export_pi05_monolithic; " + "r = export_pi05_monolithic(sys.argv[1], sys.argv[2], num_steps=int(sys.argv[3])); " + "print('EXPORT_RESULT ' + json.dumps(r), flush=True)" + ) + try: + proc = subprocess.run( + [sys.executable, "-c", code, snap, str(export_dir), str(num_steps)], + cwd=str(pin_dir), env=harness_env, capture_output=True, text=True, + timeout=export_timeout_seconds, + ) + except subprocess.TimeoutExpired as exc: + _parity_out_volume.commit() + raise RuntimeError( + f"export exceeded {export_timeout_seconds:.0f}s; stopped" + ) from exc + print(proc.stdout[-4000:], flush=True) + print(proc.stderr[-4000:], flush=True) + if proc.returncode != 0: + raise RuntimeError(f"export failed with return code {proc.returncode}") + line = next(x for x in proc.stdout.splitlines() if x.startswith("EXPORT_RESULT ")) + export_result = _json.loads(line.split(" ", 1)[1]) + provenance_path.write_text(_json.dumps(provenance, sort_keys=True) + "\n") + _parity_out_volume.commit() _progress(f"export ok: {export_result.get('size_mb', 0):.1f}MB") summary["export"] = export_result - # Release the export phase's cached GPU blocks before the parity - # subprocess: the harness places the reference on CUDA, and the - # parent's caching-allocator reservation would otherwise OOM it - # (observed: parent held 22.03/22.06 GiB, child got 16 MiB free). - import gc as _gc - - _gc.collect() - try: - import torch as _torch - - if _torch.cuda.is_available(): - _torch.cuda.empty_cache() - except Exception as exc: # noqa: BLE001 -- diagnostic only - print(f"[warn] cuda cache release failed: {exc!r}", flush=True) - _progress("parent GPU cache released") # 3. Receipt-grade parity with CUDA, CPU fallback disabled. receipt_path = Path(PARITY_OUT_PATH) / "pi05" / "pi05-export-parity-receipt.json" @@ -290,19 +322,18 @@ def _progress(stage: str) -> None: "--receipt", str(receipt_path), ] _progress("running receipt-grade parity (CUDA, no CPU fallback)") - harness_env = dict(_os.environ) - harness_env["PYTHONPATH"] = str(pin_dir / "src") + _os.pathsep + harness_env.get( - "PYTHONPATH", "" - ) - proc = subprocess.run( - cmd, cwd=str(pin_dir), capture_output=True, text=True, env=harness_env - ) + try: + proc = subprocess.run( + cmd, cwd=str(pin_dir), capture_output=True, text=True, env=harness_env, + timeout=parity_timeout_seconds, + ) + except subprocess.TimeoutExpired as exc: + _parity_out_volume.commit() + raise RuntimeError(f"parity harness exceeded {parity_timeout_seconds:.0f}s; stopped") from exc print(proc.stdout[-6000:], flush=True) print(proc.stderr[-3000:], flush=True) summary["harness_returncode"] = proc.returncode - import json as _json - if receipt_path.is_file(): receipt = _json.loads(receipt_path.read_text()) summary["verdict"] = receipt.get("verdict") @@ -318,14 +349,29 @@ def _progress(stage: str) -> None: summary["verdict"] = "no-receipt" _parity_out_volume.commit() - _watchdog.cancel() _progress("done") return summary @app.local_entrypoint() -def main() -> None: - result = run_export_parity.remote() +def main( + tether_revision: str = TETHER_REVISION, + reuse_export: bool = False, + num_steps: int = 10, + export_timeout_seconds: float = EXPORT_TIMEOUT_SECONDS, + parity_timeout_seconds: float = PARITY_TIMEOUT_SECONDS, + run_id: str = "", +) -> None: + import uuid + + result = run_export_parity.remote( + tether_revision=tether_revision, + reuse_export=reuse_export, + num_steps=num_steps, + export_timeout_seconds=export_timeout_seconds, + parity_timeout_seconds=parity_timeout_seconds, + run_id=run_id or uuid.uuid4().hex, + ) print("\n=== RESULT ===") for key, value in result.items(): print(f" {key}: {value}") diff --git a/scripts/modal_pi05_monolithic_export.py b/scripts/modal_pi05_monolithic_export.py index b8222245..efa15b6f 100644 --- a/scripts/modal_pi05_monolithic_export.py +++ b/scripts/modal_pi05_monolithic_export.py @@ -38,7 +38,7 @@ def _hf_secret(): modal.Image.debian_slim(python_version="3.12") .pip_install( # Install lerobot first — it pins torch>=2.7 and hub>=1.0 - "lerobot==0.5.1", + "lerobot==0.6.1", "num2words", "safetensors>=0.4.0", "onnx>=1.16", @@ -165,7 +165,6 @@ def _patched_embed_image(self, image): features = out.pooler_output else: features = out - features = features * self.paligemma.config.text_config.hidden_size ** 0.5 if features.dtype != out_dtype: features = features.to(out_dtype) return features @@ -203,6 +202,7 @@ def _ccm_shim(*args, **kwargs): import torch.nn.functional as _F from lerobot.policies.pi05 import modeling_pi05 as _mp0 from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks as _make_att_2d_masks + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d _patched_called = [0] @@ -247,7 +247,7 @@ def _patched_denoise_step(self, prefix_pad_masks, past_key_values, x_t, timestep prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" outputs_embeds, _ = self.paligemma_with_expert.forward( @@ -510,7 +510,6 @@ def _patched_embed_image(self, image): image = image.to(torch.float32) out = self.paligemma.model.get_image_features(image) features = out.pooler_output if hasattr(out, "pooler_output") else out - features = features * self.paligemma.config.text_config.hidden_size ** 0.5 if features.dtype != out_dtype: features = features.to(out_dtype) return features diff --git a/scripts/modal_pi0_export_parity_pinned.py b/scripts/modal_pi0_export_parity_pinned.py index 52d5e10b..6b9f9883 100644 --- a/scripts/modal_pi0_export_parity_pinned.py +++ b/scripts/modal_pi0_export_parity_pinned.py @@ -11,7 +11,8 @@ - A100-40GB: the fp32 reference + ~13 GB session cannot co-reside on A10G (pi0.5 refix OOM'd allocating a 128 MB arena there); this is a tier change only, gates unmodified. -- In-function self-kill watchdog bounds GPU burn. +- GPU burn is bounded by per-phase subprocess timeouts, the Modal function + timeout and retries=0, plus a restart guard (see EXPORT_TIMEOUT_SECONDS). Usage: modal run scripts/modal_pi0_export_parity_pinned.py @@ -39,7 +40,7 @@ ) HF_CACHE_PATH = "/root/.cache/huggingface" -PARITY_OUT_PATH = "/parity_out" +PARITY_OUT_PATH = "/parity_out/lerobot-061-tf554-20261012" def _hf_secret(): @@ -53,9 +54,14 @@ def _hf_secret(): modal.Image.debian_slim(python_version="3.12") .apt_install("git") .pip_install( - # Install lerobot first — it pins torch>=2.7 and hub>=1.0 - "lerobot==0.5.1", + # torch>=2.8: every onnx-diagnostic 0.9.x patches + # torch._export.non_strict_utils._constrain_user_specified_dimhint_range, + # which torch 2.7.1 lacks, so torch_export_patches() raises before export. + # lerobot 0.6.1 caps torch <2.12. (Studio's 2.7.1 pin is an MPS eval pin.) + "torch>=2.8,<2.12", + "lerobot==0.6.1", "num2words", + "diffusers==0.39.0", "safetensors>=0.4.0", "onnx>=1.16", "onnxruntime>=1.20", @@ -63,13 +69,13 @@ def _hf_secret(): "onnx-diagnostic>=0.9", "optree", # onnx-diagnostic soft-dep "scipy", # onnx-diagnostic patches_transformers_qwen2_5 transitively - "numpy", + "numpy==2.2.6", "accelerate", - "draccus", + "draccus==0.11.6", ) - # transformers==5.3.0 EXACTLY: the monolithic exporter refuses anything - # else (5.4+ has a q_length regression in masking_utils.sdpa_mask). - .pip_install("transformers==5.3.0") + # transformers==5.5.4 EXACTLY: LeRobot 0.6.1 needs >=5.4,<5.6 and the + # monolithic exporter gate refuses anything but the qualified release. + .pip_install("transformers==5.5.4") .env({ "HF_HOME": HF_CACHE_PATH, "TRANSFORMERS_CACHE": f"{HF_CACHE_PATH}/transformers", @@ -105,13 +111,28 @@ def _hf_secret(): ) +# Phase budgets from measured runs: pi0 fresh export 1776.5 s and harness +# incl. setup 274.9 s on A100-40GB (Studio +# docs/research/verification/lerobot-06-parity-rerun-20261012.md); pi0.5 +# fresh export plus failed session about 2190 s on A10G +# (docs/parity/parity-pi05-refix-report.md). About 1.5x and 3x margin. Each phase runs in a subprocess with this timeout, so an +# overrun ends the input as an ordinary failure. The previous in-container +# SIGKILL watchdog surfaced as a runner crash, which Modal treats as an +# internal failure and re-schedules regardless of the retry policy. +EXPORT_TIMEOUT_SECONDS = 2700 +PARITY_TIMEOUT_SECONDS = 900 +SETUP_MARGIN_SECONDS = 600 +MODAL_TIMEOUT_SECONDS = EXPORT_TIMEOUT_SECONDS + PARITY_TIMEOUT_SECONDS + SETUP_MARGIN_SECONDS + + @app.function( image=gpu_image, gpu="A100-40GB", - timeout=3600, + timeout=MODAL_TIMEOUT_SECONDS, + retries=0, volumes={ HF_CACHE_PATH: _hf_cache_volume, - PARITY_OUT_PATH: _parity_out_volume, + "/parity_out": _parity_out_volume, }, secrets=[_hf_secret()], ) @@ -120,39 +141,24 @@ def run_export_parity( tether_revision: str = TETHER_REVISION, num_steps: int = 10, reuse_export: bool = False, - watchdog_seconds: float = 3300, + export_timeout_seconds: float = EXPORT_TIMEOUT_SECONDS, + parity_timeout_seconds: float = PARITY_TIMEOUT_SECONDS, + run_id: str = "", ) -> dict: import subprocess import sys import time from pathlib import Path + import json as _json import os as _os - import signal as _signal - import threading as _threading - - _WATCHDOG_SECONDS = float( - _os.environ.get("PARITY_WATCHDOG_SECONDS", str(watchdog_seconds)) - ) - - def _self_kill() -> None: - print( - f"[watchdog] deadline of {_WATCHDOG_SECONDS:.0f}s reached; killing this " - "container rather than burning GPU time on a hung step.", - flush=True, - ) - _os.kill(_os.getpid(), _signal.SIGKILL) - - _watchdog = _threading.Timer(_WATCHDOG_SECONDS, _self_kill) - _watchdog.daemon = True - _watchdog.start() _started_at = time.monotonic() def _progress(stage: str) -> None: print(f"[progress] {time.monotonic() - _started_at:7.1f}s {stage}", flush=True) - _progress("container up, watchdog armed") + _progress("container up") summary: dict = { "model_source": "lerobot/pi0_base", "model_revision": model_revision, @@ -161,6 +167,20 @@ def _progress(stage: str) -> None: "num_steps": num_steps, } + # Restart guard: if Modal re-schedules this input after a container + # crash, the second attempt finds its own marker and exits at once + # instead of starting another GPU export. + if run_id: + marker = Path(PARITY_OUT_PATH) / "pi0" / "attempts" / run_id + _parity_out_volume.reload() + if marker.exists(): + _progress("restarted attempt detected; refusing to run again") + summary["verdict"] = "aborted-restarted-attempt" + return summary + marker.parent.mkdir(parents=True, exist_ok=True) + marker.write_text(str(time.time())) + _parity_out_volume.commit() + # Fail fast if the CUDA EP cannot be created (a broken CUDA/cuDNN # image otherwise surfaces much later as a no-fallback gate failure # after the reference has already run). @@ -228,46 +248,58 @@ def _progress(stage: str) -> None: snap = snapshot_download("lerobot/pi0_base", revision=model_revision) summary["snapshot_dir"] = snap - from tether.exporters.monolithic import export_pi0_monolithic - export_dir = Path(PARITY_OUT_PATH) / "pi0" / "export" - existing = export_dir / "model.onnx" - # model.onnx may be graph-only with weights in model.onnx.data: count - # the whole artifact dir, which is also what the receipt hashes. - existing_bytes = ( - sum(f.stat().st_size for f in export_dir.glob("*") if f.is_file()) - if existing.is_file() - else 0 + provenance_path = export_dir / "tether-export-provenance.json" + harness_env = dict(_os.environ) + harness_env["PYTHONPATH"] = str(pin_dir / "src") + _os.pathsep + harness_env.get( + "PYTHONPATH", "" ) - if reuse_export and existing_bytes > 10**9: - export_result = { - "status": "reused", - "onnx_path": str(existing), - "size_mb": existing_bytes / 1e6, - "num_steps": num_steps, - } - _progress(f"reusing retained export ({export_result['size_mb']:.1f}MB)") + from importlib.metadata import version as _v + + provenance = { + "tether_revision": tether_revision, + "model_revision": model_revision, + "num_steps": num_steps, + **{p: _v(p) for p in ("lerobot", "transformers", "torch", "onnx-diagnostic", "onnxscript")}, + } + summary["export_provenance"] = provenance + retained = ( + _json.loads(provenance_path.read_text()) if provenance_path.is_file() else None + ) + if reuse_export and retained == provenance: + export_result = {"status": "reused", "onnx_path": str(export_dir / "model.onnx")} + _progress("reusing retained export (provenance matches)") else: - _progress("export_pi0_monolithic from pinned snapshot") - export_result = export_pi0_monolithic(snap, export_dir, num_steps=num_steps) + if reuse_export: + _progress(f"not reusing: retained provenance {retained} != {provenance}") + _progress("export_pi0_monolithic from pinned snapshot (subprocess)") + code = ( + "import json, sys; from tether.exporters.monolithic import export_pi0_monolithic; " + "r = export_pi0_monolithic(sys.argv[1], sys.argv[2], num_steps=int(sys.argv[3])); " + "print('EXPORT_RESULT ' + json.dumps(r), flush=True)" + ) + try: + proc = subprocess.run( + [sys.executable, "-c", code, snap, str(export_dir), str(num_steps)], + cwd=str(pin_dir), env=harness_env, capture_output=True, text=True, + timeout=export_timeout_seconds, + ) + except subprocess.TimeoutExpired as exc: + _parity_out_volume.commit() + raise RuntimeError( + f"export exceeded {export_timeout_seconds:.0f}s; stopped" + ) from exc + print(proc.stdout[-4000:], flush=True) + print(proc.stderr[-4000:], flush=True) + if proc.returncode != 0: + raise RuntimeError(f"export failed with return code {proc.returncode}") + line = next(x for x in proc.stdout.splitlines() if x.startswith("EXPORT_RESULT ")) + export_result = _json.loads(line.split(" ", 1)[1]) + provenance_path.write_text(_json.dumps(provenance, sort_keys=True) + "\n") + _parity_out_volume.commit() _progress(f"export ok: {export_result.get('size_mb', 0):.1f}MB") summary["export"] = export_result - # Release the export phase's cached GPU blocks before the parity - # subprocess: the harness places the reference on CUDA, and the - # parent's caching-allocator reservation would otherwise OOM it - # (observed on pi0.5: parent held 22.03/22.06 GiB, child got 16 MiB). - import gc as _gc - - _gc.collect() - try: - import torch as _torch - - if _torch.cuda.is_available(): - _torch.cuda.empty_cache() - except Exception as exc: # noqa: BLE001 -- diagnostic only - print(f"[warn] cuda cache release failed: {exc!r}", flush=True) - _progress("parent GPU cache released") # 3. Receipt-grade parity with CUDA, CPU fallback disabled. receipt_path = Path(PARITY_OUT_PATH) / "pi0" / "pi0-export-parity-receipt.json" @@ -282,19 +314,18 @@ def _progress(stage: str) -> None: "--receipt", str(receipt_path), ] _progress("running receipt-grade parity (CUDA, no CPU fallback)") - harness_env = dict(_os.environ) - harness_env["PYTHONPATH"] = str(pin_dir / "src") + _os.pathsep + harness_env.get( - "PYTHONPATH", "" - ) - proc = subprocess.run( - cmd, cwd=str(pin_dir), capture_output=True, text=True, env=harness_env - ) + try: + proc = subprocess.run( + cmd, cwd=str(pin_dir), capture_output=True, text=True, env=harness_env, + timeout=parity_timeout_seconds, + ) + except subprocess.TimeoutExpired as exc: + _parity_out_volume.commit() + raise RuntimeError(f"parity harness exceeded {parity_timeout_seconds:.0f}s; stopped") from exc print(proc.stdout[-6000:], flush=True) print(proc.stderr[-3000:], flush=True) summary["harness_returncode"] = proc.returncode - import json as _json - if receipt_path.is_file(): receipt = _json.loads(receipt_path.read_text()) summary["verdict"] = receipt.get("verdict") @@ -310,7 +341,6 @@ def _progress(stage: str) -> None: summary["verdict"] = "no-receipt" _parity_out_volume.commit() - _watchdog.cancel() _progress("done") return summary @@ -320,11 +350,19 @@ def main( tether_revision: str = TETHER_REVISION, reuse_export: bool = False, num_steps: int = 10, + export_timeout_seconds: float = EXPORT_TIMEOUT_SECONDS, + parity_timeout_seconds: float = PARITY_TIMEOUT_SECONDS, + run_id: str = "", ) -> None: + import uuid + result = run_export_parity.remote( tether_revision=tether_revision, reuse_export=reuse_export, num_steps=num_steps, + export_timeout_seconds=export_timeout_seconds, + parity_timeout_seconds=parity_timeout_seconds, + run_id=run_id or uuid.uuid4().hex, ) print("\n=== RESULT ===") for key, value in result.items(): diff --git a/scripts/modal_pi0_monolithic_export.py b/scripts/modal_pi0_monolithic_export.py index 34df8896..09aefd64 100644 --- a/scripts/modal_pi0_monolithic_export.py +++ b/scripts/modal_pi0_monolithic_export.py @@ -48,7 +48,7 @@ def _hf_secret(): modal.Image.debian_slim(python_version="3.12") .pip_install( # Install lerobot first — it pins torch>=2.7 and hub>=1.0 - "lerobot==0.5.1", + "lerobot==0.6.1", "num2words", "safetensors>=0.4.0", "onnx>=1.16", @@ -175,7 +175,6 @@ def _patched_embed_image(self, image): features = out.pooler_output else: features = out - features = features * self.paligemma.config.text_config.hidden_size ** 0.5 if features.dtype != out_dtype: features = features.to(out_dtype) return features @@ -213,6 +212,7 @@ def _ccm_shim(*args, **kwargs): import torch.nn.functional as _F from lerobot.policies.pi0 import modeling_pi0 as _mp0 from lerobot.policies.pi0.modeling_pi0 import make_att_2d_masks as _make_att_2d_masks + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d _patched_called = [0] @@ -256,7 +256,7 @@ def _patched_denoise_step(self, state, prefix_pad_masks, past_key_values, x_t, t prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" outputs_embeds, _ = self.paligemma_with_expert.forward( @@ -519,7 +519,6 @@ def _patched_embed_image(self, image): image = image.to(torch.float32) out = self.paligemma.model.get_image_features(image) features = out.pooler_output if hasattr(out, "pooler_output") else out - features = features * self.paligemma.config.text_config.hidden_size ** 0.5 if features.dtype != out_dtype: features = features.to(out_dtype) return features @@ -750,7 +749,6 @@ def _patched_embed_image(self, image): image = image.to(torch.float32) out = self.paligemma.model.get_image_features(image) features = out.pooler_output if hasattr(out, "pooler_output") else out - features = features * self.paligemma.config.text_config.hidden_size ** 0.5 if features.dtype != out_dtype: features = features.to(out_dtype) return features diff --git a/scripts/modal_stateout_distill_stage2.py b/scripts/modal_stateout_distill_stage2.py index 2a5228a9..ee02bcb6 100644 --- a/scripts/modal_stateout_distill_stage2.py +++ b/scripts/modal_stateout_distill_stage2.py @@ -68,8 +68,8 @@ def _bust(): modal.Image.debian_slim(python_version="3.12") .apt_install("git", "ffmpeg") .pip_install( - "lerobot==0.5.1", - "transformers==5.3.0", + "lerobot==0.6.1", + "transformers==5.5.4", "num2words", "safetensors>=0.4.0", "numpy", @@ -111,6 +111,7 @@ def distill( import torch import torch.nn.functional as F from lerobot.policies.pi05.modeling_pi05 import PI05Policy, PI05Pytorch, make_att_2d_masks + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d from tether.distill.snapflow_pi0_model import ( load_snapflow_student, enable_snapflow_state_out, @@ -246,7 +247,7 @@ def distill( ) prefix_att_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 - prefix_att_2d_masks_4d = student.model._prepare_attention_masks_4d(prefix_att_2d) + prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d) student.model.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" _, past_key_values = student.model.paligemma_with_expert.forward( attention_mask=prefix_att_2d_masks_4d, diff --git a/scripts/modal_stateout_smoke.py b/scripts/modal_stateout_smoke.py index 0c2b4974..b0281a5e 100644 --- a/scripts/modal_stateout_smoke.py +++ b/scripts/modal_stateout_smoke.py @@ -58,8 +58,8 @@ def _bust(): modal.Image.debian_slim(python_version="3.12") .apt_install("git") .pip_install( - "lerobot==0.5.1", - "transformers==5.3.0", + "lerobot==0.6.1", + "transformers==5.5.4", "num2words", "safetensors>=0.4.0", "numpy", @@ -167,9 +167,10 @@ def smoke( images, img_masks, lang_tokens, lang_masks, ) from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d prefix_att_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 - prefix_att_2d_masks_4d = policy.model._prepare_attention_masks_4d(prefix_att_2d) + prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d) policy.model.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" _, past_key_values = policy.model.paligemma_with_expert.forward( diff --git a/src/tether/_lerobot_version.py b/src/tether/_lerobot_version.py new file mode 100644 index 00000000..c1d0d01a --- /dev/null +++ b/src/tether/_lerobot_version.py @@ -0,0 +1,7 @@ +"""LeRobot runtime qualified for dataset admission and export (ADR 0003).""" + +QUALIFIED_LEROBOT_VERSION = "0.6.1" +# LeRobot 0.6.1 declares transformers>=5.4,<5.6 and its pi0/pi0.5 code relies +# on 5.4+ embedding scaling; on 5.3 the image and language embeddings are +# 1/sqrt(hidden) too small (Studio ADR 0003, correction 2026-10-12). +QUALIFIED_TRANSFORMERS_VERSION = "5.5.4" diff --git a/src/tether/cli.py b/src/tether/cli.py index 9e6055cd..249bd5aa 100644 --- a/src/tether/cli.py +++ b/src/tether/cli.py @@ -237,7 +237,7 @@ def export( "path, one ONNX file). Opt into --decomposed only if you specifically need " "the 5-stage export for debugging; --decomposed is the older path with " "known correctness gaps. Monolithic requires `pip install " - "'fastcrest-tether[monolithic]'` (pins transformers==5.3.0).", + "'fastcrest-tether[monolithic]'` (pins transformers==5.5.4).", ), export_mode: str = typer.Option( "auto", @@ -363,7 +363,7 @@ def export( err_console.print(f"[red]{exc}[/red]", markup=False) err_console.print( "\nFix: pip install 'fastcrest-tether[monolithic]' " - "(pins transformers==5.3.0; use a clean venv to avoid " + "(pins transformers==5.5.4; use a clean venv to avoid " "the base transformers<5.0 conflict)", style="cyan", markup=False, ) @@ -459,7 +459,7 @@ def export( err_console.print(f"[red]{exc}[/red]", markup=False) err_console.print( "\nFix: pip install 'fastcrest-tether[monolithic]' " - "(pins transformers==5.3.0; use a clean venv to avoid " + "(pins transformers==5.5.4; use a clean venv to avoid " "the base transformers<5.0 conflict)", style="cyan", markup=False, ) @@ -2392,7 +2392,7 @@ def serve( "lerobot's RTCProcessor so the robot keeps executing the tail " "of one chunk while the next chunk is being computed. 2-3× " "effective throughput on Jetson-class latency. Requires " - "`pip install fastcrest-tether[rtc]` (pulls lerobot==0.5.1).", + "`pip install fastcrest-tether[rtc]` (pulls lerobot==0.6.1).", ), rtc_execution_horizon: int = typer.Option( 10, @@ -5606,7 +5606,7 @@ def go( console.print( "\n`tether go` needs the monolithic export extras to deploy this model.\n" "Fix: pip install 'fastcrest-tether[monolithic]'\n" - "(pins transformers==5.3.0; use a clean venv to avoid the " + "(pins transformers==5.5.4; use a clean venv to avoid the " "base transformers<5.0 conflict)", style="cyan", markup=False, ) diff --git a/src/tether/diagnostics/check_vlm_tokenization.py b/src/tether/diagnostics/check_vlm_tokenization.py index a005442e..386c789b 100644 --- a/src/tether/diagnostics/check_vlm_tokenization.py +++ b/src/tether/diagnostics/check_vlm_tokenization.py @@ -123,8 +123,8 @@ def _run(model_path: str, **kwargs) -> CheckResult: actual=f"AutoTokenizer raised {type(e).__name__}: {e}", remediation=( f"Export has tokenizer files but they are malformed. Check the " - f"export step's logs OR pin transformers==5.3.0 per ADR " - f"2026-04-17 onnx-export-gotchas. See {GH_ISSUE}." + f"export step's logs OR install the qualified transformers " + f"(pip install 'fastcrest-tether[monolithic]'). See {GH_ISSUE}." ), duration_ms=0.0, github_issue=GH_ISSUE, @@ -150,7 +150,7 @@ def _run(model_path: str, **kwargs) -> CheckResult: remediation=( f"Tokenizer/model vocab_size mismatch — tokens > {vocab_size} will " f"crash the ONNX path. Likely cause: transformers version drift. " - f"Pin transformers==5.3.0 per ADR 2026-04-17." + f"Install the qualified transformers: pip install 'fastcrest-tether[monolithic]'." ), duration_ms=0.0, github_issue=GH_ISSUE, diff --git a/src/tether/distill/pi05_state_out_processor.py b/src/tether/distill/pi05_state_out_processor.py index 6e0e5ef4..466be057 100644 --- a/src/tether/distill/pi05_state_out_processor.py +++ b/src/tether/distill/pi05_state_out_processor.py @@ -32,7 +32,7 @@ def _register_step() -> type: import (matching the pattern in snapflow_pi0_model.py). """ from lerobot.processor.pipeline import ProcessorStep, ProcessorStepRegistry - from lerobot.types import EnvTransition, TransitionKey + from lerobot.lerobot_types import EnvTransition, TransitionKey @ProcessorStepRegistry.register(name="pi05_prepare_tokenizer_state_out") @dataclass diff --git a/src/tether/distill/snapflow_pi0_model.py b/src/tether/distill/snapflow_pi0_model.py index 043f8882..d0a8a45c 100644 --- a/src/tether/distill/snapflow_pi0_model.py +++ b/src/tether/distill/snapflow_pi0_model.py @@ -7,11 +7,11 @@ to a correction vector added to the time channel before the action_time MLP. When ``target_time`` is None, behaves exactly like the parent (teacher inference path is unaffected). - 2. ``denoise_step`` does NOT ``copy.deepcopy(past_key_values)`` — the - downstream ``paligemma_with_expert.forward`` is called with - ``use_cache=False`` so past_kv is read-only; deepcopy is strictly - defensive overhead (and breaks on graph-attached tensors during - training — pytorch#103001). + 2. pi0 ``denoise_step`` clones the prefix cache with LeRobot 0.6.1's + tensor-clone helper. DynamicCache updates occur even with + ``use_cache=False``; cloning preserves the prefix and VLM gradients + without ``copy.deepcopy`` on graph-attached tensors. pi0.5 retains + its detached cache copy for the existing expert-only training path. 3. ``denoise_step`` does NOT force-cast ``suffix_out`` to fp32 — uses ``self.action_out_proj.weight.dtype`` instead, so bf16 training stays bf16 end-to-end. @@ -150,10 +150,12 @@ def _resolve_snapflow_pi0_class() -> type: return _SNAPFLOW_PI0_CLASS import torch - from lerobot.policies.pi0.modeling_pi0 import ( - PI0Pytorch, + from lerobot.policies.pi0.modeling_pi0 import PI0Pytorch + from lerobot.policies.common.vla_utils import ( + clone_past_key_values, create_sinusoidal_pos_embedding, make_att_2d_masks, + prepare_attention_masks_4d, ) class SnapFlowPI0Pytorch(PI0Pytorch): @@ -263,8 +265,9 @@ def denoise_step( ): """Parent denoise_step minus the deepcopy and minus the fp32 cast. - - Skips ``copy.deepcopy(past_key_values)``: safe because the - downstream forward uses ``use_cache=False`` (read-only pass). + - Clones the prefix cache using LeRobot's tensor clone helper. + This preserves VLM gradients and prevents DynamicCache updates + from changing the shared prefix between velocity evaluations. - Skips the hardcoded ``suffix_out.to(fp32)``: casts to ``self.action_out_proj.weight.dtype`` instead, so bf16 training stays bf16. @@ -289,9 +292,10 @@ def denoise_step( prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" # noqa: SLF001 + past_key_values = clone_past_key_values(past_key_values) outputs_embeds, _ = self.paligemma_with_expert.forward( attention_mask=full_att_2d_masks_4d, position_ids=position_ids, @@ -343,7 +347,7 @@ def sample_actions_1step( ) prefix_att_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 - prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d) + prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d) self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" # noqa: SLF001 _, past_key_values = self.paligemma_with_expert.forward( @@ -392,10 +396,11 @@ def _resolve_snapflow_pi05_class() -> type: return _SNAPFLOW_PI05_CLASS import torch - from lerobot.policies.pi05.modeling_pi05 import ( - PI05Pytorch, + from lerobot.policies.pi05.modeling_pi05 import PI05Pytorch + from lerobot.policies.common.vla_utils import ( create_sinusoidal_pos_embedding, make_att_2d_masks, + prepare_attention_masks_4d, ) class SnapFlowPI05Pytorch(PI05Pytorch): @@ -533,7 +538,7 @@ def denoise_step( prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" # noqa: SLF001 # Detach past_kv tensors so deepcopy succeeds on graph-attached @@ -583,7 +588,7 @@ def sample_actions_1step( ) prefix_att_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 - prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d) + prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d) self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" # noqa: SLF001 _, past_key_values = self.paligemma_with_expert.forward( @@ -641,11 +646,7 @@ def _resolve_snapflow_pi05_state_out_class() -> type: return _SNAPFLOW_PI05_STATE_OUT_CLASS import torch - import torch.nn.functional as F - from lerobot.policies.pi05.modeling_pi05 import ( - create_sinusoidal_pos_embedding, - make_att_2d_masks, - ) + from lerobot.policies.common.vla_utils import make_att_2d_masks, prepare_attention_masks_4d parent = _resolve_snapflow_pi05_class() @@ -745,7 +746,7 @@ def denoise_step( # type: ignore[override] prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" # noqa: SLF001 past_key_values = _detach_kv_cache_deep(past_key_values) @@ -794,7 +795,7 @@ def sample_actions_1step( # type: ignore[override] ) prefix_att_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 - prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d) + prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d) self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" # noqa: SLF001 _, past_key_values = self.paligemma_with_expert.forward( diff --git a/src/tether/eval/preflight.py b/src/tether/eval/preflight.py index 324386ad..3b09eaa1 100644 --- a/src/tether/eval/preflight.py +++ b/src/tether/eval/preflight.py @@ -172,7 +172,7 @@ def __post_init__(self) -> None: ), "dep-version-conflict": ( "Dependency version mismatch detected. Pin: robosuite==1.4.1, " - "bddl==1.0.1, mujoco==3.3.2, lerobot==0.5.1. Use the bundled " + "bddl==1.0.1, mujoco==3.3.2, lerobot==0.6.1. Use the bundled " "Modal image (--runtime modal) for known-good pins." ), "osmesa-compile-hang": ( diff --git a/src/tether/exporters/decomposed.py b/src/tether/exporters/decomposed.py index 204eb64e..5852d6bf 100644 --- a/src/tether/exporters/decomposed.py +++ b/src/tether/exporters/decomposed.py @@ -38,7 +38,7 @@ ## Dependencies -Same as ``monolithic`` extra: transformers==5.3.0, lerobot==0.5.1, +Same as ``monolithic`` extra: transformers==5.5.4, lerobot==0.6.1, onnx-diagnostic>=0.9. ``apply_export_patches`` from ``monolithic.py`` is invoked here to get the shared denoise-step + cache-freeze patches. @@ -359,6 +359,7 @@ def _load_pi05_policy( apply_export_patches, _force_eager_attn, _apply_pi05_denoise_step_patch, + _install_snapflow_export_denoise_step, ) apply_export_patches() @@ -406,6 +407,8 @@ def _load_pi05_policy( _gc_disable() _force_eager_attn(policy.model) _apply_pi05_denoise_step_patch() # reuses monolithic's F.pad mask fix + if student_checkpoint is not None: + _install_snapflow_export_denoise_step(policy.model, is_pi05=True) logger.info("[decomposed] Loaded in %.1fs", time.time() - t0) return policy @@ -789,6 +792,7 @@ def _build_prefix_class(): import torch import torch.nn as nn from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d class _Pi05PrefixWrapper(nn.Module): def __init__(self, pi05_model): @@ -817,7 +821,7 @@ def forward( ) prefix_att_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 - prefix_att_2d_masks_4d = self.model._prepare_attention_masks_4d(prefix_att_2d) + prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d) self.model.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" # noqa: SLF001 _, past_key_values = self.model.paligemma_with_expert.forward( diff --git a/src/tether/exporters/monolithic.py b/src/tether/exporters/monolithic.py index 26212be3..f5e00e58 100644 --- a/src/tether/exporters/monolithic.py +++ b/src/tether/exporters/monolithic.py @@ -10,8 +10,8 @@ tether export lerobot/smolvla_base --monolithic --output ./smol Requires (pinned in the ``monolithic`` extra): - lerobot==0.5.1 - transformers==5.3.0 (5.4+ has a q_length scalar regression) + lerobot==0.6.1 + transformers==5.5.4 (LeRobot 0.6.1 requires >=5.4,<5.6) onnx-diagnostic>=0.9 onnxscript>=0.1 optree, scipy, num2words @@ -23,12 +23,16 @@ """ from __future__ import annotations +from tether._lerobot_version import QUALIFIED_LEROBOT_VERSION, QUALIFIED_TRANSFORMERS_VERSION + import logging import time from importlib.metadata import PackageNotFoundError, version as _dist_version from pathlib import Path from typing import Any +MIN_EXPORT_TORCH = (2, 8) + logger = logging.getLogger(__name__) @@ -156,15 +160,32 @@ def _require_monolithic_deps() -> None: missing = [] try: import transformers - if transformers.__version__ != "5.3.0": + if transformers.__version__ != QUALIFIED_TRANSFORMERS_VERSION: raise ImportError( f"transformers {transformers.__version__} detected; the " - f"monolithic export requires exactly 5.3.0 (5.4+ has a " - f"q_length regression in masking_utils.sdpa_mask). Install " - f"with: pip install transformers==5.3.0" + f"monolithic export requires exactly {QUALIFIED_TRANSFORMERS_VERSION} " + f"(LeRobot {QUALIFIED_LEROBOT_VERSION} needs >=5.4: on 5.3 its " + f"pi0/pi0.5 embeddings are mis-scaled). Install with: " + f"pip install transformers=={QUALIFIED_TRANSFORMERS_VERSION}" ) except ImportError as e: - missing.append(f"transformers==5.3.0 ({e})") + missing.append(f"transformers=={QUALIFIED_TRANSFORMERS_VERSION} ({e})") + + # onnx-diagnostic 0.8.11 and 0.9.0-0.9.3 unconditionally patch + # torch._export.non_strict_utils._constrain_user_specified_dimhint_range, + # which first exists in torch 2.8.0. On older torch every pi0/pi0.5/ + # SmolVLA/GR00T export dies with an AttributeError on entering + # torch_export_patches(); refuse here instead. + try: + import torch + major_minor = tuple(int(part) for part in str(torch.__version__).split(".")[:2]) + if major_minor < MIN_EXPORT_TORCH: + raise ImportError( + f"torch {torch.__version__} detected; onnx-diagnostic's export " + "patches need torch>=2.8" + ) + except (ImportError, ValueError) as e: + missing.append(f"torch>=2.8 ({e})") try: __import__("lerobot") @@ -175,14 +196,14 @@ def _require_monolithic_deps() -> None: "lerobot imports but package metadata is missing; install with: " "pip install 'fastcrest-tether[monolithic]'" ) from exc - if lerobot_version != "0.5.1": + if lerobot_version != QUALIFIED_LEROBOT_VERSION: raise ImportError( f"lerobot {lerobot_version} detected; monolithic export requires " - "exactly lerobot==0.5.1. Install with: " + f"exactly lerobot=={QUALIFIED_LEROBOT_VERSION}. Install with: " "pip install 'fastcrest-tether[monolithic]'" ) except ImportError as e: - missing.append(f"lerobot==0.5.1 ({e})") + missing.append(f"lerobot=={QUALIFIED_LEROBOT_VERSION} ({e})") for mod_name, pip_name in [ ("onnx_diagnostic", "onnx-diagnostic>=0.9"), @@ -433,24 +454,9 @@ def _safe_where(condition, x=None, y=None, *args, **kwargs): return _orig_where(condition, x, y, *args, **kwargs) torch.where = _safe_where - # pi0-specific: PaliGemmaWithExpertModel.embed_image image_outputs extraction - try: - from lerobot.policies.pi0 import modeling_pi0 - import torch as _t - - def _patched_embed_image(self, image): - out_dtype = image.dtype - if image.dtype != _t.float32: - image = image.to(_t.float32) - out = self.paligemma.model.get_image_features(image) - features = out.pooler_output if hasattr(out, "pooler_output") else out - features = features * self.paligemma.config.text_config.hidden_size ** 0.5 - if features.dtype != out_dtype: - features = features.to(out_dtype) - return features - modeling_pi0.PaliGemmaWithExpertModel.embed_image = _patched_embed_image - except ImportError: - pass + # LeRobot 0.6.1 returns image features in its native embedding scale. + # Keep embed_image unchanged: multiplying by sqrt(hidden_size), as the + # old 0.5.1 shim did, changes every prefix cache and exported action. # =============================================================== # pi0 cos=1.0 at num_steps=10 — the full fix stack (2026-04-19) @@ -485,6 +491,7 @@ def _patched_embed_image(self, image): from transformers.cache_utils import DynamicLayer as _DL from lerobot.policies.pi0 import modeling_pi0 as _mp0 from lerobot.policies.pi0.modeling_pi0 import make_att_2d_masks as _make_att_2d_masks + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d _orig_layer_update = _DL.update @@ -525,7 +532,7 @@ def _patched_denoise_step(self, state, prefix_pad_masks, past_key_values, x_t, t prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" outputs_embeds, _ = self.paligemma_with_expert.forward( @@ -1170,6 +1177,7 @@ def _apply_pi05_denoise_step_patch() -> None: import torch from lerobot.policies.pi05 import modeling_pi05 as _mp05 from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks as _make + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d def _patched(self, prefix_pad_masks, past_key_values, x_t, timestep): # pi0.5: NO state arg (tokenized into language upstream). @@ -1198,7 +1206,7 @@ def _patched(self, prefix_pad_masks, past_key_values, x_t, timestep): prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" outputs_embeds, _ = self.paligemma_with_expert.forward( @@ -1414,12 +1422,14 @@ def forward( def _install_snapflow_export_denoise_step(model: Any, *, is_pi05: bool) -> None: """Replace SnapFlowPI*Pytorch.denoise_step with an export-friendly variant that drops the training-only detach+deepcopy (which doesn't trace - through torch.export.export). The 1-step export path calls denoise_step - exactly once per forward, so past_kv isn't reused across calls — the - defensive copy that matters for multi-step inference isn't needed here. + through torch.export.export). Freeze cache updates during the expert + pass, as for teacher export: tracing can evaluate even a single-step + forward repeatedly, and must preserve the prefix cache each time. """ import torch + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d + if is_pi05: from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks else: @@ -1430,6 +1440,7 @@ def _export_denoise_step(self, prefix_pad_masks, past_key_values, x_t, timestep, target_time=None, state=None): suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix( x_t, timestep, target_time=target_time, + **({"state": state} if hasattr(self, "state_proj") else {}), ) suffix_len = suffix_pad_masks.shape[1] batch_size = prefix_pad_masks.shape[0] @@ -1445,7 +1456,7 @@ def _export_denoise_step(self, prefix_pad_masks, past_key_values, x_t, prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" outputs_embeds, _ = self.paligemma_with_expert.forward( @@ -1479,7 +1490,7 @@ def _export_denoise_step(self, state, prefix_pad_masks, past_key_values, prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 - full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks) + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" outputs_embeds, _ = self.paligemma_with_expert.forward( @@ -1497,7 +1508,15 @@ def _export_denoise_step(self, state, prefix_pad_masks, past_key_values, # Bind onto the class (the SnapFlow subclass) so the instance picks # it up. This overrides the training-time detach+deepcopy variant # for the duration of the export process. - type(model).denoise_step = _export_denoise_step + def _denoise_step_with_flag(self, *args, **kwargs): + previous_phase = _denoise_phase[0] + _denoise_phase[0] = True + try: + return _export_denoise_step(self, *args, **kwargs) + finally: + _denoise_phase[0] = previous_phase + + type(model).denoise_step = _denoise_step_with_flag def export_gr00t_monolithic( diff --git a/src/tether/finetune/backends/snapflow_backend.py b/src/tether/finetune/backends/snapflow_backend.py index 1a1c0cc1..80951f87 100644 --- a/src/tether/finetune/backends/snapflow_backend.py +++ b/src/tether/finetune/backends/snapflow_backend.py @@ -560,6 +560,7 @@ def _build_pi05_adapters(teacher: Any, student: Any): """ import torch from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d from lerobot.utils.constants import ( OBS_LANGUAGE_ATTENTION_MASK, @@ -577,7 +578,7 @@ def _build_prefix_cache_pi05(policy, obs_kwargs): ) prefix_att_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_pos = torch.cumsum(prefix_pad_masks, dim=1) - 1 - prefix_att_4d = m._prepare_attention_masks_4d(prefix_att_2d) + prefix_att_4d = prepare_attention_masks_4d(prefix_att_2d) m.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" m.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" _, past_kv = m.paligemma_with_expert.forward( @@ -681,6 +682,7 @@ def _build_pi_family_adapters( import torch import torch.nn.functional as F from lerobot.policies.pi0.modeling_pi0 import make_att_2d_masks + from lerobot.policies.common.vla_utils import prepare_attention_masks_4d from lerobot.utils.constants import ( OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, @@ -718,7 +720,7 @@ def _build_prefix_cache(policy, obs_kwargs): ) prefix_att_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) prefix_pos = torch.cumsum(prefix_pad_masks, dim=1) - 1 - prefix_att_4d = m._prepare_attention_masks_4d(prefix_att_2d) + prefix_att_4d = prepare_attention_masks_4d(prefix_att_2d) # Mirror lerobot's sample_actions trick: force eager attention # before the prefix forward so bf16 models don't hit the SDPA # dtype-mismatch-on-bias error (modeling_pi0.py:844). @@ -731,10 +733,8 @@ def _build_prefix_cache(policy, obs_kwargs): inputs_embeds=[prefix_embs, None], use_cache=True, ) - # SnapFlowPI0Pytorch.denoise_step skips the `copy.deepcopy(past_kv)` - # that the stock lerobot path does (modeling_pi0.py:918), so we no - # longer need to detach past_kv tensors. Student path now keeps - # past_kv graph-attached → VLM receives gradients. + # SnapFlow pi0 uses LeRobot's tensor clone helper, so the student + # cache remains graph-attached and VLM gradients can flow. return past_kv, prefix_pad_masks, state def _run_denoise_step(policy, x, t, obs_kwargs, target_time=None): diff --git a/src/tether/finetune/local_dataset.py b/src/tether/finetune/local_dataset.py index 35fc4bf4..be50f0db 100644 --- a/src/tether/finetune/local_dataset.py +++ b/src/tether/finetune/local_dataset.py @@ -8,6 +8,8 @@ from __future__ import annotations +from tether._lerobot_version import QUALIFIED_LEROBOT_VERSION + import hashlib import json import math @@ -18,25 +20,25 @@ from typing import Any, NoReturn -ADAPTER_SOURCE = "11c0a0cc0d5f4cd81de9904ec0d08f0305ba9ecec6e5952c25fe435c857afdf3" -WRITER_SOURCE = "1396b9fab7aecddd10006c33c47a487ffdcb54b4" -WRITER_WHEEL = "bbd11021023fde0947b6d1ff1c52fe91c86a28ab09a96359892f3ef7e8866862" +ADAPTER_SOURCE = "af6969a6663d0b288db781ad46736eb62d33a361dbe75cd578da6c163b075307" +WRITER_SOURCE = "7e241bd630a3719a56157a497ce5d08f244784f1" +WRITER_WHEEL = "1894516040c65f80a45bd9741f8174aae90ed5d93da0627ab4f1a85fd8d75e90" WRITER_FILES = { - "lerobot/datasets/lerobot_dataset.py": "03207172bbc82d541f93168017b11e2b3807cf719126bf47d04873257a4cbbd8", - "lerobot/datasets/utils.py": "1dd05df5f1d12721f227fb3601d99c22cd845cb59627b68004d3184c9e8e1486", - "lerobot/datasets/compute_stats.py": "7b0621bbdd36987939742a1803b3491b794db7b0e6393677e9b1a3e1d86aeafc", - "lerobot/datasets/dataset_writer.py": "f377ab064840d45f2dab98da27f5c71b5dfb64f70789fec5c82935eecc925207", - "lerobot/datasets/dataset_reader.py": "c5893ecb6406a6852f01dcc4d6662f6ca2320cd931bc9e4d1cd7dc83dfaa9d6c", - "lerobot/datasets/dataset_metadata.py": "df392d4c0b584177735550251c45877ba7b644377003e3fa08b71536c9969b59", - "lerobot/datasets/io_utils.py": "1b557bab9f1f1445602dfd5dc5bca69a3fe2c34c400ae36d3c69b5c7d4d1120a", - "lerobot/datasets/feature_utils.py": "1477145c9f27051511e661f5189f838d8fdfacdb0c7c8424a64d4c181027ce7a", - "lerobot/datasets/image_writer.py": "99ce7371c50b7fb0931b49228fa5e3ae30302e5559efd7f64dead6334aa9e088", + "lerobot/datasets/lerobot_dataset.py": "07523e2a4f9c53b6a96df9255ad131ee30e707fd8610ec9d2178d3741023ae3a", + "lerobot/datasets/utils.py": "1fd3eddf70baf1c5c7f4e745c78338c580385283fad00e853d49cd0f34dd57e3", + "lerobot/datasets/compute_stats.py": "d546a5c690eff716a10e9b13d70cd5d1ef7fed88f2f9c40c99f2f7f50845cf45", + "lerobot/datasets/dataset_writer.py": "319f4efb53e892ceeb439aae6d25209d14564e3dc91134c7cc2c6c0b5e71ebc5", + "lerobot/datasets/dataset_reader.py": "3a477dae6aa887449f09be25b385485b065861d13c5448b9369a0ba09f01bb5d", + "lerobot/datasets/dataset_metadata.py": "23c927553a53e320161f35f9b36e8b61a446833ca17a4287b9ef5cdc8d19c50f", + "lerobot/datasets/io_utils.py": "af08c7875f34f8469f80c4fd903af12971d6fa96c179be984887b93ef2f6f91c", + "lerobot/datasets/feature_utils.py": "928369cff50f0d663a125abfc67683ccbae0e552ed2eb22edde9e6bcff2aef81", + "lerobot/datasets/image_writer.py": "54cad218c99caad105ea8738507ee87fc8c2858832d31aa5b21d324c97360631", } NORMALIZATION_FILES = { - "lerobot/processor/normalize_processor.py": "faedf40f64e1fcc6aeb287ff757abe6e85983560d82fcde060252d8f75e277d6", - "lerobot/policies/act/configuration_act.py": "11aa9037f981a957d63632561c795fac243e7898ee9302f7afac2f5d45920f7b", - "lerobot/policies/smolvla/configuration_smolvla.py": "6c1f371ccb7ed073774c459387f85626edcf568edbbaa64b4cea1bf70a80feca", - "lerobot/configs/default.py": "e077db856b4ecad9476bdf12fb2e72a0294ae3a302413afa21774a8788ab0d20", + "lerobot/processor/normalize_processor.py": "f0cd88bec4c954a93ba1cc90300473b96eb177762789c1780cf07801fc1af9d6", + "lerobot/policies/act/configuration_act.py": "35a60c0dd3c811b137b5eb035bdbd40d40e2c24ba0cda61756052e6e1c829282", + "lerobot/policies/smolvla/configuration_smolvla.py": "2fb637cb428fa2fdf1d114646dcffaf4728216bfe5b7039d5d0cac4857ffc4e0", + "lerobot/configs/default.py": "e8307a48aa97c1b3fd0722e9a849deb1bc1bb48b63e3109b63683eb7ffe08105", } IDENTITY_FIELDS = ["task", "seed", "episode", "initial_state", "embodiment"] PARTITIONS = ("development", "holdout", "test") @@ -605,13 +607,13 @@ def _verify(dataset_root: Path, expected_manifest_sha256: str, expected_profile: inventory = _inventory(export, receipt) writer = _object(receipt.get("writer"), "writer") _require( - writer.get("version") == "0.5.1" + writer.get("version") == QUALIFIED_LEROBOT_VERSION and writer.get("source_commit") == WRITER_SOURCE and writer.get("wheel_sha256") == WRITER_WHEEL and writer.get("source_files") == WRITER_FILES - and _object(writer.get("runtime_versions"), "writer runtime").get("lerobot") == "0.5.1", + and _object(writer.get("runtime_versions"), "writer runtime").get("lerobot") == QUALIFIED_LEROBOT_VERSION, "writer-pin", - "Export writer/source pins must match qualified LeRobot 0.5.1.", + f"Export writer/source pins must match qualified LeRobot {QUALIFIED_LEROBOT_VERSION}.", ) info = _object(_json(root / "meta/info.json"), "dataset info") stats = _object(_json(root / "meta/stats.json"), "dataset statistics") diff --git a/src/tether/finetune/preflight/schema.py b/src/tether/finetune/preflight/schema.py index 3f9fe4dd..8ecb5f50 100644 --- a/src/tether/finetune/preflight/schema.py +++ b/src/tether/finetune/preflight/schema.py @@ -13,6 +13,7 @@ """ from __future__ import annotations +from tether._lerobot_version import QUALIFIED_LEROBOT_VERSION from tether.finetune.lerobot_v3 import verify_local_dataset import logging @@ -155,7 +156,7 @@ def check_schema(cfg: FinetuneConfig) -> PreflightCheck: expected_modes = {"VISUAL": "IDENTITY", "STATE": "MEAN_STD", "ACTION": "MEAN_STD"} if not isinstance(base_config, dict) or base_config.get("type") != "smolvla": return PreflightCheck("schema", "fail", "Local admission requires a resolvable SmolVLA base config.") - # An omitted field uses only the exact LeRobot 0.5.1 class defaults + # An omitted field uses only the exact LeRobot 0.6.1 class defaults # qualified in the verified receipt. Explicit null/contradictory modes # are not treated as omission. resolved_modes = base_config.get("normalization_mapping", expected_modes) @@ -165,7 +166,7 @@ def check_schema(cfg: FinetuneConfig) -> PreflightCheck: "resolved_normalization_mapping": resolved_modes, "normalization_mapping_source": ( "base config.json" if "normalization_mapping" in base_config - else "pinned LeRobot 0.5.1 SmolVLAConfig defaults" + else f"pinned LeRobot {QUALIFIED_LEROBOT_VERSION} SmolVLAConfig defaults" ), } max_state = base_config.get("max_state_dim") diff --git a/src/tether/finetune/resume.py b/src/tether/finetune/resume.py index 19c4326a..6d981110 100644 --- a/src/tether/finetune/resume.py +++ b/src/tether/finetune/resume.py @@ -1,5 +1,6 @@ from __future__ import annotations +import json import logging from pathlib import Path import subprocess @@ -8,7 +9,9 @@ class ResumeConfigurationError(ValueError): - pass + def __init__(self, message: str, *, code: str = "resume-configuration-invalid"): + self.code = code + super().__init__(message) def resolve_resume_checkpoint(output_dir: Path) -> tuple[Path, Path]: @@ -39,11 +42,31 @@ def resolve_resume_checkpoint(output_dir: Path) -> tuple[Path, Path]: raise ResumeConfigurationError( f"resume checkpoint is missing pretrained_model/train_config.json: {checkpoint}" ) + try: + saved = json.loads(config.read_text()) + except (OSError, ValueError) as exc: + raise ResumeConfigurationError(f"resume train_config.json is unreadable: {config}") from exc + if isinstance(saved, dict) and any( + key in saved + for key in ( + "eval_freq", + "use_rabc", + "rabc_progress_path", + "rabc_kappa", + "rabc_epsilon", + "rabc_head_mode", + ) + ): + raise ResumeConfigurationError( + "cross-version-resume-unsupported: LeRobot 0.5.x training configs " + "cannot resume in the qualified 0.6.1 runtime; start a new run.", + code="cross-version-resume-unsupported", + ) return checkpoint, config def build_resume_command(output_dir: Path) -> tuple[list[str], Path]: - """Build the LeRobot v0.5.1 resume command from the checkpoint's own config. + """Build the LeRobot v0.6.1 resume command from the checkpoint's own config. Do not repeat `--policy.path`, dataset, optimizer, scheduler, seed or other training settings. LeRobot's resume contract loads those values from the diff --git a/src/tether/models/llm/paligemma_backbone.py b/src/tether/models/llm/paligemma_backbone.py index 6c53b54d..f794e7c9 100644 --- a/src/tether/models/llm/paligemma_backbone.py +++ b/src/tether/models/llm/paligemma_backbone.py @@ -110,6 +110,18 @@ def embed_tokens(self) -> nn.Module: before merging in the projected image embeds.""" return self.language_model.embed_tokens + def embed_text_openpi(self, token_ids: torch.Tensor) -> torch.Tensor: + """openpi text embedding: raw table lookup times sqrt(hidden). + + Looks up the weight directly because the module's own scaling moved + between transformers releases: 5.3 Gemma uses a plain nn.Embedding + (scaled later in GemmaModel.forward), 5.4+ uses + GemmaTextScaledWordEmbedding, which already multiplies by sqrt(hidden). + Calling the module and scaling again double-scales on 5.4+. + """ + weight = self.embed_tokens.weight + return torch.nn.functional.embedding(token_ids, weight) * self.text_hidden_size**0.5 + @property def text_hidden_size(self) -> int: """Hidden dim of the language tower (1152 for PaliGemma-3B-pt-224).""" diff --git a/src/tether/models/vlas/pi0.py b/src/tether/models/vlas/pi0.py index 932505d5..fcf184eb 100644 --- a/src/tether/models/vlas/pi0.py +++ b/src/tether/models/vlas/pi0.py @@ -398,8 +398,6 @@ def predict_action( # protocol: pre-scale text externally; image is raw multi_modal_projector # output (no additional scaling needed because the get_image_features # path isn't used — we call multi_modal_projector directly). - text_hidden = self.llm_backbone.text_hidden_size - sqrt_h = text_hidden ** 0.5 image_embeds_list: list[torch.Tensor] = [] for img in images: # SigLIP: [B, 3, 224, 224] → [B, 256, vision_hidden=1152] @@ -410,7 +408,7 @@ def predict_action( image_embeds_list.append(img_emb) # Pre-scale text per lerobot embed_prefix at modeling_pi0.py:669. - text_embs = self.llm_backbone.embed_tokens(lang_tokens) * sqrt_h + text_embs = self.llm_backbone.embed_text_openpi(lang_tokens) # ─── 4. Concat into prefix ────────────────────────────────────── prefix_embs = torch.cat([*image_embeds_list, text_embs], dim=1) diff --git a/src/tether/models/vlas/pi05.py b/src/tether/models/vlas/pi05.py index 50302853..bc108511 100644 --- a/src/tether/models/vlas/pi05.py +++ b/src/tether/models/vlas/pi05.py @@ -304,14 +304,12 @@ def predict_action( # Same scaling protocol as Pi0VLA: PiGemma omits internal sqrt(h) # scaling, so text needs explicit pre-scale; image stays raw # (lerobot patched_embed_image net = ×1). - text_hidden = self.llm_backbone.text_hidden_size - sqrt_h = text_hidden ** 0.5 image_embeds_list: list[torch.Tensor] = [] for img in images: img_emb = self.vision_backbone(img) img_emb = self.llm_backbone.multi_modal_projector(img_emb) image_embeds_list.append(img_emb) - text_embs = self.llm_backbone.embed_tokens(lang_tokens) * sqrt_h + text_embs = self.llm_backbone.embed_text_openpi(lang_tokens) # ─── 4. Concat into prefix (NO state token — state is in lang_tokens) ── prefix_embs = torch.cat([*image_embeds_list, text_embs], dim=1) diff --git a/src/tether/runtime/fast_inference/pi05.py b/src/tether/runtime/fast_inference/pi05.py index 9fc374ab..950bb5c0 100644 --- a/src/tether/runtime/fast_inference/pi05.py +++ b/src/tether/runtime/fast_inference/pi05.py @@ -799,9 +799,8 @@ def predict_action( ) # Embed via the paligemma text model's embed_tokens. # Pi0.5 scales the embedding by sqrt(hidden_dim) — matches the FluxVLA reference. - text_model = self._vla.llm_backbone.model.model.language_model - lang_emb = text_model.embed_tokens(lang_tokens[0, :prompt_len]) - lang_emb = (lang_emb * math.sqrt(lang_emb.shape[-1])).bfloat16() + lang_emb = self._vla.llm_backbone.embed_text_openpi(lang_tokens[0, :prompt_len]) + lang_emb = lang_emb.bfloat16() chunk_size = self.n_action_steps if noise is None: diff --git a/tests/_lerobot_pi_golden.py b/tests/_lerobot_pi_golden.py new file mode 100644 index 00000000..9018e74c --- /dev/null +++ b/tests/_lerobot_pi_golden.py @@ -0,0 +1,270 @@ +"""Golden pi0/pi0.5 reference from the qualified OLD stack (independent of LeRobot 0.6.1). + +Why: the LeRobot 0.6.1 port tests compared Tether's patched forward with +native 0.6.1 in the same interpreter. Under transformers 5.3 native 0.6.1 is +itself mis-scaled (image and language embeddings ~1/sqrt(2048) of openpi), so +"matches native" passed while GPU parity failed at cosine 0.47. This fixture +pins outputs from the stack that was GPU-qualified before the migration: +LeRobot 0.5.1 + transformers 5.3.0 + Tether cdb94dc. Any stack that drifts +from openpi scale fails against it. + + capture (old stack): python tests/_lerobot_pi_golden.py capture {pi0,pi05} OUT.pt + compare (new stack): python tests/_lerobot_pi_golden.py compare {pi0,pi05} GOLDEN.pt + export (new stack): python tests/_lerobot_pi_golden.py export {pi0,pi05} GOLDEN.pt DIR + +The tiny model shrinks constructor sizes only (two layers, width 32, two-layer +SigLIP, vocabulary 64, 16x16 images, three-action chunk). Weights are saved in +the fixture, so initialisation order differences between versions cannot leak in. +""" + +from __future__ import annotations + +import sys + +NUM_STEPS = 10 +# float32 tolerance: the old stack divides image features by sqrt(2048) inside +# transformers 5.3 and LeRobot 0.5.1 multiplies them back; 5.4+ skips both, +# and language scaling moved into GemmaTextScaledWordEmbedding. Equal maths, +# different rounding, so allow a few float32 ulps at these magnitudes. +RTOL = 1e-5 +ATOL = 1e-5 + + +def _modules(family): + if family == "pi0": + from lerobot.policies.pi0 import modeling_pi0 as module + + return module, module.PI0Pytorch, module.PI0Config + from lerobot.policies.pi05 import modeling_pi05 as module + + return module, module.PI05Pytorch, module.PI05Config + + +def build(family): + import torch + from unittest.mock import patch + + torch.set_num_threads(1) + torch.manual_seed(123) + module, cls, config_cls = _modules(family) + original_pali = module.PaliGemmaForConditionalGenerationWithPiGemma + original_expert = module.PiGemmaForCausalLM + + def tiny_pali(config): + config.text_config.vocab_size = 64 + config.vision_config.hidden_size = 32 + config.vision_config.intermediate_size = 64 + config.vision_config.num_hidden_layers = 2 + config.vision_config.num_attention_heads = 4 + config.vision_config.patch_size = 4 + config.vision_config.projection_dim = 32 + return original_pali(config=config) + + def tiny_expert(config): + config.vocab_size = 64 + return original_expert(config=config) + + config = config_cls( + device="cpu", + dtype="float32", + image_resolution=(16, 16), + max_action_dim=4, + max_state_dim=4, + chunk_size=3, + n_action_steps=3, + compile_model=False, + ) + with ( + patch.object( + module, "get_gemma_config", return_value=module.GemmaConfig(32, 2, 64, 8, 1, 4) + ), + patch.object(module, "PaliGemmaForConditionalGenerationWithPiGemma", tiny_pali), + patch.object(module, "PiGemmaForCausalLM", tiny_expert), + ): + model = cls(config).eval() + return model + + +def make_inputs(): + import torch + + g = torch.Generator().manual_seed(7) + # Images in [-1, 1] like LeRobot's preprocessed SigLIP input. + return { + "img0": torch.rand(2, 3, 16, 16, generator=g) * 2 - 1, + "img1": torch.rand(2, 3, 16, 16, generator=g) * 2 - 1, + "img2": torch.rand(2, 3, 16, 16, generator=g) * 2 - 1, + "imask0": torch.tensor([True, True]), + "imask1": torch.tensor([True, False]), + "imask2": torch.tensor([False, True]), + "tokens": torch.randint(0, 64, (2, 4), generator=g), + "lmask": torch.tensor([[True, True, False, False], [True, True, True, False]]), + "state": torch.randn(2, 4, generator=g), + "noise": torch.randn(2, 3, 4, generator=g), + "actions": torch.randn(2, 3, 4, generator=g), + "time": torch.tensor([0.7, 0.3]), + } + + +def _obs(family, x): + obs = [ + [x["img0"], x["img1"], x["img2"]], + [x["imask0"], x["imask1"], x["imask2"]], + x["tokens"], + x["lmask"], + ] + return obs + ([x["state"]] if family == "pi0" else []) + + +def outputs(family, model, x): + """Native outputs, then Tether's patched export forward (installs patches).""" + import copy + + import torch + + obs = _obs(family, x) + native = copy.deepcopy(model) + out = {} + with torch.no_grad(): + prefix, pad, _ = native.embed_prefix(*obs[:4]) + out["prefix_embs"] = prefix + out["prefix_pad"] = pad.to(torch.int64) + out["loss"] = native(*obs, x["actions"], x["noise"], x["time"]) + out["actions_1"] = native.sample_actions(*obs, noise=x["noise"].clone(), num_steps=1) + out["actions_10"] = native.sample_actions( + *obs, noise=x["noise"].clone(), num_steps=NUM_STEPS + ) + from tether.exporters import monolithic + + monolithic.apply_export_patches() + if family == "pi05": + monolithic._apply_pi05_denoise_step_patch() + out["patched_actions_10"] = model.sample_actions( + *obs, noise=x["noise"].clone(), num_steps=NUM_STEPS + ) + return out + + +def capture(family, path): + import importlib.metadata as md + + import torch + + model = build(family) + x = make_inputs() + state = {k: v.clone().contiguous() for k, v in model.state_dict().items()} + out = outputs(family, model, x) + meta = {p: md.version(p) for p in ("lerobot", "transformers", "torch")} + torch.save({"state_dict": state, "inputs": x, "outputs": out, "meta": meta}, path) + print(family, "captured", meta) + + +def _load(family, path): + import torch + + golden = torch.load(path, weights_only=True) + model = build(family) + missing, unexpected = model.load_state_dict(golden["state_dict"], strict=False) + # Tied lm_head weights may be keyed differently across versions; anything + # else missing means the architectures diverged and must fail loudly. + assert not unexpected, unexpected + assert all("lm_head" in k for k in missing), missing + return golden, model + + +def _report(name, actual, expected): + import torch + + diff = (actual.double() - expected.double()).abs() + cos = torch.nn.functional.cosine_similarity( + actual.double().flatten(), expected.double().flatten(), dim=0 + ) + scale = actual.double().norm() / expected.double().norm() + print( + f"{name}: max_abs={diff.max().item():.3e} mean_abs={diff.mean().item():.3e} " + f"cos={cos.item():.9f} norm_ratio={scale.item():.6f}" + ) + + +def compare(family, path): + import torch + + golden, model = _load(family, path) + out = outputs(family, model, golden["inputs"]) + failures = [] + for name, expected in golden["outputs"].items(): + actual = out[name] + if actual.dtype.is_floating_point: + _report(name, actual, expected) + try: + torch.testing.assert_close(actual, expected, rtol=RTOL, atol=ATOL) + except AssertionError as exc: + failures.append(f"{name}: {exc}") + print("golden from", golden["meta"]) + assert not failures, "\n".join(failures) + print(family, "golden parity passed") + + +def export(family, path, out_dir): + """Real monolithic toolchain (onnx-diagnostic patches, torch.export, + torch.onnx.export, Where-dtype fix) on the tiny model, run in ONNX Runtime + and compared with the OLD-stack golden. Mirrors export_{pi0,pi05}_monolithic + minus checkpoint loading and tokenizer probing; B=1 like the exporter.""" + from pathlib import Path + + import numpy as np + import onnxruntime as ort + import torch + from onnx_diagnostic.torch_export_patches import torch_export_patches + + from tether.exporters import monolithic + + golden, model = _load(family, path) + x = golden["inputs"] + monolithic.apply_export_patches() + if family == "pi05": + monolithic._apply_pi05_denoise_step_patch() + monolithic._force_eager_attn(model) + + class Wrapper(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = model + + def forward(self, i0, i1, i2, m0, m1, m2, tokens, lmask, *rest): + kwargs = {"noise": rest[-1], "num_steps": NUM_STEPS} + args = [[i0, i1, i2], [m0, m1, m2], tokens, lmask, *rest[:-1]] + return self.model.sample_actions(*args, **kwargs) + + names = ["img0", "img1", "img2", "imask0", "imask1", "imask2", "tokens", "lmask"] + names += (["state"] if family == "pi0" else []) + ["noise"] + # Row 0 of the golden batch: all three cameras present for row 0 except + # camera 2; B=1 matches the production exporter's trace shape. + feed = {n: x[n][:1].clone() for n in names} + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + onnx_path = out_dir / "model.onnx" + wrapper = Wrapper().eval() + with torch_export_patches(patch_transformers=True): + ep = torch.export.export(wrapper, tuple(feed.values()), dynamic_shapes=None, strict=False) + torch.onnx.export( + ep, + tuple(feed.values()), + str(onnx_path), + input_names=names, + output_names=["actions"], + opset_version=19, + ) + monolithic._fix_onnx_where_dtype_mismatches(onnx_path) + session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"]) + (actions,) = session.run(None, {n: v.numpy() for n, v in feed.items()}) + actual = torch.from_numpy(np.asarray(actions)) + expected = golden["outputs"]["actions_10"][:1] + _report("onnx_actions_10", actual, expected) + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4) + print(family, "onnx export golden parity passed") + + +if __name__ == "__main__": + mode, family = sys.argv[1], sys.argv[2] + {"capture": capture, "compare": compare, "export": export}[mode](family, *sys.argv[3:]) diff --git a/tests/fixtures/lerobot_pi_golden/SHA256SUMS b/tests/fixtures/lerobot_pi_golden/SHA256SUMS new file mode 100644 index 00000000..1ba19560 --- /dev/null +++ b/tests/fixtures/lerobot_pi_golden/SHA256SUMS @@ -0,0 +1,2 @@ +992d9f5cedcb8d13fbe7202f584f57267df6e08ac65bee0ee3dec7cc0986bde1 pi0.pt +3f57a88d0f220703a201581d8f1033f06aed7ba311b9e61f5ec9115a6ad69c3e pi05.pt diff --git a/tests/fixtures/lerobot_pi_golden/pi0.pt b/tests/fixtures/lerobot_pi_golden/pi0.pt new file mode 100644 index 00000000..917821a4 Binary files /dev/null and b/tests/fixtures/lerobot_pi_golden/pi0.pt differ diff --git a/tests/fixtures/lerobot_pi_golden/pi05.pt b/tests/fixtures/lerobot_pi_golden/pi05.pt new file mode 100644 index 00000000..ebaceb3e Binary files /dev/null and b/tests/fixtures/lerobot_pi_golden/pi05.pt differ diff --git a/tests/integration/test_eval_aloha_smoke.py b/tests/integration/test_eval_aloha_smoke.py index 569de695..dcc1c39a 100644 --- a/tests/integration/test_eval_aloha_smoke.py +++ b/tests/integration/test_eval_aloha_smoke.py @@ -44,6 +44,6 @@ def test_aloha_two_episode_customer_flow(tmp_path): assert all(type(row["success"]) is bool for row in report["episodes"]) assert report["confidence_interval"]["method"] == "wilson-score" assert receipt["policy"]["revision"] == REVISION - assert receipt["lerobot_version"] == "0.5.1" and receipt["device"] == "cpu" + assert receipt["lerobot_version"] == "0.6.1" and receipt["device"] == "cpu" assert receipt["migrated"] is True assert (output / "eval_info.json").is_file() diff --git a/tests/test_cli_export_end_to_end.py b/tests/test_cli_export_end_to_end.py index 2a12c198..eb88717f 100644 --- a/tests/test_cli_export_end_to_end.py +++ b/tests/test_cli_export_end_to_end.py @@ -18,11 +18,14 @@ from __future__ import annotations import logging +import re import sys import types import pytest +from tether._lerobot_version import QUALIFIED_TRANSFORMERS_VERSION + def test_monolithic_module_importable(): """The module must be importable even without the [monolithic] extras @@ -34,16 +37,40 @@ def test_monolithic_module_importable(): assert hasattr(monolithic, "apply_export_patches") -def test_dispatch_by_model_id(monkeypatch): - """`export_monolithic` picks the right backend from the model_id.""" +@pytest.mark.parametrize("model_id,family", [ + ("lerobot/smolvla_base", "smolvla"), + ("lerobot/pi0_base", "pi0"), + ("lerobot/pi05_base", "pi05"), +]) +def test_dispatch_by_model_id(monkeypatch, model_id, family): + """`export_monolithic` picks the right backend from the model_id. + + The family exporters are stubbed: calling the real one starts a full + export whenever the [monolithic] extras happen to be installed, which + made this test pass or fail depending on the interpreter. + """ from tether.exporters import monolithic monkeypatch.setattr(monolithic, "_model_type_from_hub_config", lambda _: None) - with pytest.raises(ImportError, match="Missing dependencies|monolithic"): - # Wrong transformers version or missing deps -> ImportError - monolithic.export_monolithic( - "lerobot/smolvla_base", "/tmp/should_not_run", + called = [] + for name in ("smolvla", "pi0", "pi05", "gr00t"): + monkeypatch.setattr( + monolithic, f"export_{name}_monolithic", + lambda *a, _n=name, **k: called.append(_n) or {}, ) + monolithic.export_monolithic(model_id, "/tmp/should_not_run") + assert called == [family] + + +def test_dep_check_refuses_torch_without_export_patch_symbol(monkeypatch): + """onnx-diagnostic 0.9.x patches a torch 2.8+ symbol on entry, so the + gate must refuse torch<2.8 instead of crashing mid-export.""" + from tether.exporters import monolithic + import torch + + monkeypatch.setattr(torch, "__version__", "2.7.1") + with pytest.raises(ImportError, match=r"torch>=2\.8 \(torch 2\.7\.1 detected"): + monolithic._require_monolithic_deps() def test_unsupported_model_type_raises(monkeypatch): @@ -67,28 +94,40 @@ def test_dep_check_catches_wrong_transformers(monkeypatch): import transformers monkeypatch.setattr(transformers, "__version__", "4.99.0") - with pytest.raises(ImportError, match="5.3.0"): + with pytest.raises(ImportError, match=re.escape(QUALIFIED_TRANSFORMERS_VERSION)): monolithic._require_monolithic_deps() -def test_dep_check_catches_wrong_lerobot(monkeypatch): +def test_dep_check_refuses_transformers_53(monkeypatch): + """5.3.0 + LeRobot 0.6.1 mis-scales pi0/pi0.5 embeddings ~45x (GPU parity + cosine 0.47, Studio ADR 0003 correction 2026-10-12); the gate must refuse it.""" + from tether.exporters import monolithic + import transformers + + monkeypatch.setattr(transformers, "__version__", "5.3.0") + with pytest.raises(ImportError, match="mis-scaled"): + monolithic._require_monolithic_deps() + + +@pytest.mark.parametrize("lerobot_version", ["0.4.4", "0.5.1"]) +def test_dep_check_catches_wrong_lerobot(monkeypatch, lerobot_version): """_require_monolithic_deps() must fail early on the known-bad lerobot 0.4.x stack instead of letting torch.export fail downstream.""" from tether.exporters import monolithic import transformers - monkeypatch.setattr(transformers, "__version__", "5.3.0") + monkeypatch.setattr(transformers, "__version__", QUALIFIED_TRANSFORMERS_VERSION) for mod_name in ("lerobot", "onnx_diagnostic", "onnxscript", "optree", "scipy"): monkeypatch.setitem(sys.modules, mod_name, types.ModuleType(mod_name)) def _fake_dist_version(dist_name: str) -> str: if dist_name == "lerobot": - return "0.4.4" + return lerobot_version raise monolithic.PackageNotFoundError(dist_name) monkeypatch.setattr(monolithic, "_dist_version", _fake_dist_version) - with pytest.raises(ImportError, match=r"lerobot 0\.4\.4.*lerobot==0\.5\.1"): + with pytest.raises(ImportError, match=r"lerobot " + lerobot_version.replace(".", r"\.") + r".*lerobot==0\.6\.1"): monolithic._require_monolithic_deps() diff --git a/tests/test_finetune_local_dataset.py b/tests/test_finetune_local_dataset.py index 0758d44e..1a678224 100644 --- a/tests/test_finetune_local_dataset.py +++ b/tests/test_finetune_local_dataset.py @@ -169,7 +169,7 @@ def file_hash(path): ] adapter = { "policy": "verified-moments-v1", - "adapter_source_sha256": "11c0a0cc0d5f4cd81de9904ec0d08f0305ba9ecec6e5952c25fe435c857afdf3", + "adapter_source_sha256": "af6969a6663d0b288db781ad46736eb62d33a361dbe75cd578da6c163b075307", "statistics_sha256": canonical(stats), "source_dataset_sha256": identity["sha256"], "split_sha256": split["sha256"], @@ -220,11 +220,11 @@ def file_hash(path): "fps": 10, }, "writer": { - "version": "0.5.1", + "version": "0.6.1", "source_commit": WRITER_SOURCE, "wheel_sha256": WRITER_WHEEL, "source_files": dict(WRITER_FILES), - "runtime_versions": {"lerobot": "0.5.1"}, + "runtime_versions": {"lerobot": "0.6.1"}, }, "statistics_policy": "verified-moments-v1", "statistics_adapter": adapter, @@ -718,7 +718,7 @@ def test_omitted_base_modes_use_only_pinned_defaults(local_cfg): assert result.severity == "ok" assert ( result.detail["normalization_mapping_source"] - == "pinned LeRobot 0.5.1 SmolVLAConfig defaults" + == "pinned LeRobot 0.6.1 SmolVLAConfig defaults" ) config["normalization_mapping"] = None write_json(path, config) @@ -768,3 +768,23 @@ def test_existing_positional_config_arguments_stay_compatible(tmp_path): assert cfg.num_steps == 123 assert cfg.dataset_root is None assert cfg.dataset_manifest_sha256 is None + + +def test_pinned_lerobot_sources_match_installed_qualified_runtime(): + """The pins are reviewed upstream bytes; check them against a real 0.6.1 install when present.""" + from importlib import metadata + + from tether._lerobot_version import QUALIFIED_LEROBOT_VERSION + from tether.finetune import local_dataset + + try: + dist = metadata.distribution("lerobot") + except metadata.PackageNotFoundError: + pytest.skip("lerobot is not installed") + if dist.version != QUALIFIED_LEROBOT_VERSION: + pytest.skip(f"lerobot {dist.version} is not the qualified {QUALIFIED_LEROBOT_VERSION}") + pins = {**local_dataset.WRITER_FILES, **local_dataset.NORMALIZATION_FILES} + actual = { + name: hashlib.sha256(Path(dist.locate_file(name)).read_bytes()).hexdigest() for name in pins + } + assert actual == pins diff --git a/tests/test_finetune_resume_command.py b/tests/test_finetune_resume_command.py index a7de3d3a..285fe2ac 100644 --- a/tests/test_finetune_resume_command.py +++ b/tests/test_finetune_resume_command.py @@ -47,3 +47,20 @@ def test_resume_refuses_checkpoint_missing_saved_train_config(tmp_path): checkpoint(tmp_path, 100, config=False) with pytest.raises(ResumeConfigurationError, match="train_config"): build_resume_command(tmp_path) + + +def test_resume_refuses_legacy_lerobot_config(tmp_path): + model = tmp_path / "checkpoints" / "000003" / "pretrained_model" + model.mkdir(parents=True) + (model / "train_config.json").write_text('{"eval_freq": 100}') + with pytest.raises(ResumeConfigurationError, match="cross-version-resume-unsupported") as error: + build_resume_command(tmp_path) + assert error.value.code == "cross-version-resume-unsupported" + + +def test_resume_refuses_unreadable_train_config(tmp_path): + model = tmp_path / "checkpoints" / "000003" / "pretrained_model" + model.mkdir(parents=True) + (model / "train_config.json").write_text("{not json") + with pytest.raises(ResumeConfigurationError, match="unreadable"): + build_resume_command(tmp_path) diff --git a/tests/test_lerobot06_patched_forward.py b/tests/test_lerobot06_patched_forward.py new file mode 100644 index 00000000..567621f5 --- /dev/null +++ b/tests/test_lerobot06_patched_forward.py @@ -0,0 +1,343 @@ +"""Real, tiny LeRobot forwards; process isolation keeps export patches out of other tests.""" + +from __future__ import annotations + +import importlib.metadata +import subprocess +import sys +from pathlib import Path + +import pytest + + +@pytest.mark.parametrize( + "path", + [ + "pi0", + "pi05", + "decomposed", + "snapflow_pi0", + "snapflow_pi05", + "snapflow_state_out", + "private_api", + ], +) +def test_patched_forward_matches_lerobot06(path): + try: + version = importlib.metadata.version("lerobot") + except importlib.metadata.PackageNotFoundError: + pytest.skip("requires the LeRobot 0.6.1 native extra") + if version != "0.6.1": + pytest.skip("this regression exercises LeRobot 0.6.1") + result = subprocess.run( + [sys.executable, str(Path(__file__).resolve()), path], + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def _check(path): + import copy + import torch + from unittest.mock import patch + from lerobot.policies.pi0 import modeling_pi0 as pi0 + from lerobot.policies.pi05 import modeling_pi05 as pi05 + from lerobot.policies.common.vla_utils import make_att_2d_masks, prepare_attention_masks_4d + from tether.exporters import monolithic, decomposed + from tether.distill import snapflow_pi0_model as snapflow + + if path == "private_api": + import ast + import inspect + + for owner in (pi0.PI0Pytorch, pi05.PI05Pytorch): + assert not hasattr(owner, "_prepare_attention_masks_4d") + assert ( + str(inspect.signature(owner._apply_checkpoint)) == "(self, func, *args, **kwargs)" + ) + for owner in (pi0.PI0Policy, pi05.PI05Policy): + assert list(inspect.signature(owner._preprocess_images).parameters) == ["self", "batch"] + root = Path(__file__).resolve().parents[1] + for directory in (root / "src", root / "scripts"): + for source in directory.rglob("*.py"): + tree = ast.parse(source.read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + assert node.func.attr != "_prepare_attention_masks_4d", str(source) + if isinstance(node, ast.ImportFrom) and (node.module or "").startswith( + "lerobot" + ): + assert all(not name.name.startswith("_") for name in node.names), str( + source + ) + return + + torch.set_num_threads(1) + torch.manual_seed(123) + module = pi0 if path.endswith("pi0") or path == "pi0" else pi05 + cls = module.PI0Pytorch if module is pi0 else module.PI05Pytorch + config_cls = module.PI0Config if module is pi0 else module.PI05Config + # Shrink constructor configurations only. Every embedding, attention, norm, + # velocity, cache and Euler computation remains the real LeRobot code. + original_pali = module.PaliGemmaForConditionalGenerationWithPiGemma + original_expert = module.PiGemmaForCausalLM + + def tiny_pali(config): + config.text_config.vocab_size = 64 + config.vision_config.hidden_size = 32 + config.vision_config.intermediate_size = 64 + config.vision_config.num_hidden_layers = 2 + config.vision_config.num_attention_heads = 4 + config.vision_config.patch_size = 4 + config.vision_config.projection_dim = 32 + return original_pali(config=config) + + def tiny_expert(config): + config.vocab_size = 64 + return original_expert(config=config) + + config = config_cls( + device="cpu", + dtype="float32", + image_resolution=(16, 16), + max_action_dim=4, + max_state_dim=4, + chunk_size=3, + n_action_steps=3, + compile_model=False, + ) + with ( + patch.object( + module, "get_gemma_config", return_value=module.GemmaConfig(32, 2, 64, 8, 1, 4) + ), + patch.object(module, "PaliGemmaForConditionalGenerationWithPiGemma", tiny_pali), + patch.object(module, "PiGemmaForCausalLM", tiny_expert), + ): + model = cls(config).eval() + images = [torch.randn(2, 3, 16, 16) for _ in range(3)] + image_masks = [ + torch.tensor([True, True]), + torch.tensor([True, False]), + torch.tensor([False, True]), + ] + tokens = torch.randint(0, 64, (2, 4)) + masks = torch.tensor([[True, True, False, False], [True, True, True, False]]) + state = torch.randn(2, 4) + noise = torch.randn(2, 3, 4) + actions = torch.randn_like(noise) + time = torch.tensor([0.7, 0.3]) + obs = [images, image_masks, tokens, masks] + ([state] if module is pi0 else []) + native = copy.deepcopy(model) + native_step = cls.denoise_step + with torch.no_grad(): + expected_loss = native(*obs, actions, noise, time) + expected_actions = native.sample_actions(*obs, noise=noise.clone(), num_steps=3) + expected_one_step = native.sample_actions(*obs, noise=noise.clone(), num_steps=1) + prefix, pad, att = native.embed_prefix(images, image_masks, tokens, masks) + _, cache = native.paligemma_with_expert( + attention_mask=prepare_attention_masks_4d(make_att_2d_masks(pad, att)), + position_ids=pad.cumsum(1) - 1, + inputs_embeds=[prefix, None], + use_cache=True, + ) + x_t = time[:, None, None] * noise + (1 - time[:, None, None]) * actions + step_kw = dict(prefix_pad_masks=pad, past_key_values=cache, x_t=x_t, timestep=time) + if module is pi0: + step_kw["state"] = state + expected_velocity = native_step(native, **step_kw) + + if path.startswith("snapflow"): + snapflow.enable_snapflow(model) + if path == "snapflow_state_out": + snapflow.enable_snapflow_state_out(model) + # A zero target-time contribution must reproduce the teacher. + step_kw["target_time"] = torch.ones_like(time) + if path == "snapflow_state_out": + step_kw["state"] = state + # The state-out student's embedding is its intentional extension. + # Execute the UNPATCHED LeRobot denoise forward with that embedding. + ref = copy.deepcopy(model) + embed = ref.embed_suffix + ref.embed_suffix = lambda x, t: embed( + x, t, target_time=torch.ones_like(t), state=state + ) + expected_velocity = native_step( + ref, **{k: v for k, v in step_kw.items() if k not in ("state", "target_time")} + ) + torch.testing.assert_close( + model.denoise_step(**step_kw), expected_velocity, rtol=1e-6, atol=1e-6 + ) + import types + + if path == "snapflow_state_out": + ref.denoise_step = types.MethodType(native_step, ref) + expected_one_step = cls.sample_actions(ref, *obs, noise=noise.clone(), num_steps=1) + sample_obs = obs + ([state] if path == "snapflow_state_out" else []) + torch.testing.assert_close( + model.sample_actions_1step(*sample_obs, noise=noise.clone()), + expected_one_step, + rtol=1e-6, + atol=1e-6, + ) + # Exercise both training-prefix adapters with real model forwards. + from types import SimpleNamespace + from tether.finetune.backends.snapflow_backend import _build_velocity_adapters + from lerobot.utils.constants import ( + OBS_LANGUAGE_TOKENS, + OBS_LANGUAGE_ATTENTION_MASK, + OBS_STATE, + ) + + policy = SimpleNamespace( + model=model, + config=config, + _preprocess_images=lambda batch: (images, image_masks), + prepare_state=lambda batch: state, + ) + _, student_velocity = _build_velocity_adapters( + policy, policy, "pi0" if module is pi0 else "pi05" + ) + torch.testing.assert_close( + student_velocity( + x_t, + time, + target_time=torch.ones_like(time), + **{ + OBS_LANGUAGE_TOKENS: tokens, + OBS_LANGUAGE_ATTENTION_MASK: masks, + OBS_STATE: state, + }, + ), + expected_velocity, + rtol=1e-6, + atol=1e-6, + ) + # Also test a learned target-time correction. A zero-init-only + # test would miss an exporter silently dropping target_time. + model.target_time_embed_mlp[-1].weight.normal_(std=0.01) + model.target_time_embed_mlp[-1].bias.normal_(std=0.01) + ref = copy.deepcopy(model) + embed = ref.embed_suffix + state_kwargs = {"state": state} if path == "snapflow_state_out" else {} + ref.embed_suffix = lambda *args: embed( + *args, target_time=torch.ones_like(args[-1]), **state_kwargs + ) + ref.denoise_step = types.MethodType(native_step, ref) + native_kwargs = {k: v for k, v in step_kw.items() if k != "target_time"} + if path == "snapflow_state_out": + native_kwargs.pop("state") + expected_velocity = native_step(ref, **native_kwargs) + expected_loss = cls.forward(ref, *obs, actions, noise, time) + expected_one_step = cls.sample_actions(ref, *obs, noise=noise.clone(), num_steps=1) + torch.testing.assert_close( + model.denoise_step(**step_kw), expected_velocity, rtol=1e-6, atol=1e-6 + ) + monolithic.apply_export_patches() + if module is pi05: + monolithic._apply_pi05_denoise_step_patch() + if path.startswith("snapflow"): + monolithic._install_snapflow_export_denoise_step(model, is_pi05=module is pi05) + step_kw["past_key_values"] = copy.deepcopy(cache) + actual_velocity = model.denoise_step(**step_kw) + torch.testing.assert_close(actual_velocity, expected_velocity, rtol=1e-6, atol=1e-6) + torch.testing.assert_close( + (noise - actions - actual_velocity).square(), expected_loss, rtol=1e-6, atol=1e-6 + ) + if not path.startswith("snapflow"): + torch.testing.assert_close( + model.sample_actions(*obs, noise=noise.clone(), num_steps=3), + expected_actions, + rtol=1e-6, + atol=1e-6, + ) + if module is pi05 and path != "pi05": + decomposed.PI05_PALIGEMMA_LAYERS = 2 + prefix_wrapper = decomposed._build_prefix_class()(model) + flat = prefix_wrapper(*images, *image_masks, tokens, masks) + for actual, layer_value in zip( + flat[:-1], + [v for layer in cache.layers for v in (layer.keys, layer.values)], + strict=True, + ): + torch.testing.assert_close(actual, layer_value, rtol=1e-6, atol=1e-6) + torch.testing.assert_close(flat[-1], pad, rtol=0, atol=0) + per_step = decomposed._build_expert_per_step_class()(model) + torch.testing.assert_close( + per_step(*flat, x_t, time, *([state] if path == "snapflow_state_out" else [])), + expected_velocity, + rtol=1e-6, + atol=1e-6, + ) + n_steps = 1 if path.startswith("snapflow") else 3 + baked = decomposed._build_expert_class()(model, n_steps) + state_args = (state,) if path == "snapflow_state_out" else () + step_args = (*flat, x_t, time, *state_args) + baked_args = (*flat, noise, *state_args) + expected_baked = expected_one_step if path.startswith("snapflow") else expected_actions + torch.testing.assert_close(baked(*baked_args), expected_baked, rtol=1e-6, atol=1e-6) + prefix_graph = torch.export.export( + prefix_wrapper, (*images, *image_masks, tokens, masks), strict=False + ).module() + for actual, expected in zip( + prefix_graph(*images, *image_masks, tokens, masks), flat, strict=True + ): + torch.testing.assert_close(actual, expected, rtol=1e-6, atol=1e-6) + step_graph = torch.export.export(per_step, step_args, strict=False).module() + torch.testing.assert_close( + step_graph(*step_args), expected_velocity, rtol=1e-6, atol=1e-6 + ) + baked_graph = torch.export.export(baked, baked_args, strict=False).module() + torch.testing.assert_close( + baked_graph(*baked_args), expected_baked, rtol=1e-6, atol=1e-6 + ) + + class ExportStep(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = model + + def forward(self, *args): + from transformers.cache_utils import DynamicCache + + kv = DynamicCache() + for i in range(2): + kv.update(args[2 * i], args[2 * i + 1], i) + kwargs = dict( + prefix_pad_masks=args[4], past_key_values=kv, x_t=args[5], timestep=args[6] + ) + if module is pi0 or path == "snapflow_state_out": + kwargs["state"] = args[7] + if path.startswith("snapflow"): + kwargs["target_time"] = torch.ones_like(args[6]) + return self.model.denoise_step(**kwargs) + + class ExportActions(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = model + + def forward(self, *args): + if path.startswith("snapflow"): + return self.model.sample_actions_1step(*args[:-1], noise=args[-1]) + return self.model.sample_actions(*args[:-1], noise=args[-1], num_steps=3) + + action_args = (*(sample_obs if path.startswith("snapflow") else obs), noise) + exported_actions = torch.export.export(ExportActions(), action_args, strict=False).module() + torch.testing.assert_close( + exported_actions(*action_args), + expected_one_step if path.startswith("snapflow") else expected_actions, + rtol=1e-6, + atol=1e-6, + ) + + flat_cache = tuple(v for layer in cache.layers for v in (layer.keys, layer.values)) + export_args = (*flat_cache, pad, x_t, time, state) + exported = torch.export.export(ExportStep(), export_args, strict=False).module() + torch.testing.assert_close(exported(*export_args), expected_velocity, rtol=1e-6, atol=1e-6) + print(path, "native forward parity passed") + + +if __name__ == "__main__": + _check(sys.argv[1]) diff --git a/tests/test_lerobot_pi_golden.py b/tests/test_lerobot_pi_golden.py new file mode 100644 index 00000000..cd4c99c4 --- /dev/null +++ b/tests/test_lerobot_pi_golden.py @@ -0,0 +1,73 @@ +"""pi0/pi0.5 against a golden captured on the pre-migration stack (see _lerobot_pi_golden.py). + +The fixture was captured with LeRobot 0.5.1 + transformers 5.3.0 + torch 2.7.1 +and Tether cdb94dc (the GPU-qualified stack before ADR 0003): + + PYTHONPATH=/src /bin/python \\ + tests/_lerobot_pi_golden.py capture {pi0,pi05} tests/fixtures/lerobot_pi_golden/{pi0,pi05}.pt +""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import importlib.util +import subprocess +import sys +from pathlib import Path + +import pytest + +HERE = Path(__file__).resolve().parent +FIXTURES = HERE / "fixtures" / "lerobot_pi_golden" +HELPER = HERE / "_lerobot_pi_golden.py" +FAMILIES = ["pi0", "pi05"] + + +def _require_lerobot_061(): + try: + version = importlib.metadata.version("lerobot") + except importlib.metadata.PackageNotFoundError: + pytest.skip("requires the LeRobot 0.6.1 native extra") + if version != "0.6.1": + pytest.skip("this golden check exercises LeRobot 0.6.1") + + +def _run(*args, timeout=600): + result = subprocess.run( + [sys.executable, str(HELPER), *args], capture_output=True, text=True, timeout=timeout + ) + assert result.returncode == 0, result.stdout + result.stderr + return result.stdout + + +def test_golden_fixture_provenance(): + torch = pytest.importorskip("torch") + sums = dict( + reversed(line.split()) for line in (FIXTURES / "SHA256SUMS").read_text().splitlines() + ) + for family in FAMILIES: + path = FIXTURES / f"{family}.pt" + assert hashlib.sha256(path.read_bytes()).hexdigest() == sums[path.name] + meta = torch.load(path, weights_only=True)["meta"] + assert meta == {"lerobot": "0.5.1", "transformers": "5.3.0", "torch": "2.7.1"} + + +@pytest.mark.parametrize("family", FAMILIES) +def test_native_and_patched_match_pre_migration_golden(family): + _require_lerobot_061() + out = _run("compare", family, str(FIXTURES / f"{family}.pt")) + assert f"{family} golden parity passed" in out + + +@pytest.mark.parametrize("family", FAMILIES) +def test_monolithic_onnx_export_matches_pre_migration_golden(family, tmp_path): + _require_lerobot_061() + torch = pytest.importorskip("torch") + if tuple(int(p) for p in torch.__version__.split(".")[:2]) < (2, 8): + pytest.skip("onnx-diagnostic export patches need torch>=2.8") + for module in ("onnx_diagnostic", "onnxruntime", "onnxscript"): + if importlib.util.find_spec(module) is None: + pytest.skip(f"requires the monolithic extra ({module})") + out = _run("export", family, str(FIXTURES / f"{family}.pt"), str(tmp_path), timeout=900) + assert f"{family} onnx export golden parity passed" in out diff --git a/tests/test_lora_argv_parser.py b/tests/test_lora_argv_parser.py index 44847f2e..53993adb 100644 --- a/tests/test_lora_argv_parser.py +++ b/tests/test_lora_argv_parser.py @@ -1,7 +1,7 @@ -"""Parse the real LoRA argv with LeRobot 0.5.1, offline and without weights. +"""Parse the real LoRA argv with LeRobot 0.6.1, offline and without weights. Mirrors Studio's backend/tests/test_lora_argv_parser.py at the Tether boundary. -Requires LeRobot 0.5.1 and PEFT. Inline resume only parses its flag; saved-config +Requires LeRobot 0.6.1 and PEFT. Inline resume only parses its flag; saved-config resume restores the checkpoint path. Neither case runs training or qualifies dataset admission or interrupted-training recovery. """ @@ -29,7 +29,7 @@ def no_network(*args, **kwargs): monkeypatch.setattr(socket.socket, "connect_ex", no_network) pytest.importorskip("lerobot") pytest.importorskip("peft") - assert version("lerobot") == "0.5.1" + assert version("lerobot") == "0.6.1" import draccus from lerobot.configs import parser, policies @@ -41,7 +41,7 @@ def no_network(*args, **kwargs): assert Path(run.__file__).resolve().is_relative_to(source) # Keep the real config loader and policy decoding; stub only hub lookup. policy_file = tmp_path / "config.json" - policy_file.write_text(json.dumps(draccus.encode(SmolVLAConfig(device="cpu")))) + policy_file.write_text(json.dumps({**draccus.encode(SmolVLAConfig(device="cpu")), "type": "smolvla"})) def hub_config(*args, **kwargs): assert kwargs["repo_id"] == "lerobot/smolvla_base" @@ -80,7 +80,8 @@ def _config(monkeypatch, tmp_path, rank, dataset="remote", resume=False): def verified_export(actual_root, actual_digest): assert actual_root == root assert actual_digest == cfg.dataset_manifest_sha256 - return snapshot + # Real verify_local_export shape (lerobot_v3.verify_local_dataset reads the receipt). + return {**snapshot, "receipt": {"artifacts": []}, "info": {}} monkeypatch.setattr(admission, "verify_local_export", verified_export) monkeypatch.setattr(schema, "check_schema", lambda cfg: SimpleNamespace(severity="ok")) diff --git a/tests/test_one_command_deploy.py b/tests/test_one_command_deploy.py index 8c25a7bd..11bc16b6 100644 --- a/tests/test_one_command_deploy.py +++ b/tests/test_one_command_deploy.py @@ -338,7 +338,7 @@ def test_export_dep_missing_errors_with_monolithic_install_hint(self, runner, cl with patch( "tether.exporters.monolithic.export_monolithic", - side_effect=ImportError("Missing dependencies: lerobot==0.5.1, onnx-diagnostic"), + side_effect=ImportError("Missing dependencies: lerobot==0.6.1, onnx-diagnostic"), ): result = runner.invoke(cli_app, [ "go", diff --git a/tests/test_paligemma_backbone.py b/tests/test_paligemma_backbone.py index 31344857..7bf35768 100644 --- a/tests/test_paligemma_backbone.py +++ b/tests/test_paligemma_backbone.py @@ -215,3 +215,25 @@ def __init__(self): ) return _Stub() + + +@pytest.mark.parametrize("module_scale", [1.0, 8.0**0.5]) +def test_embed_text_openpi_is_independent_of_embedding_module_scaling(module_scale): + """transformers 5.3 Gemma embeds unscaled; 5.4+ GemmaTextScaledWordEmbedding + already multiplies by sqrt(hidden). openpi text = table lookup * sqrt(hidden) + either way, never sqrt(hidden) twice.""" + stub = _make_stub_paligemma(text_hidden=8, vocab_size=16) + table = stub.model.language_model.embed_tokens + + class Scaled(nn.Embedding): + def forward(self, ids): + return super().forward(ids) * module_scale + + scaled = Scaled(16, 8) + scaled.weight = table.weight + stub.model.language_model.embed_tokens = scaled + ids = torch.tensor([[1, 5, 15]]) + expected = table.weight[ids] * 8.0**0.5 + torch.testing.assert_close( + PaliGemmaBackbone(model=stub).embed_text_openpi(ids), expected, rtol=0, atol=0 + ) diff --git a/tests/test_pi05_vla.py b/tests/test_pi05_vla.py index 352ffffb..914e05a9 100644 --- a/tests/test_pi05_vla.py +++ b/tests/test_pi05_vla.py @@ -22,6 +22,7 @@ from tether.models.base_vla import BaseVLA from tether.models.heads import VLAHead from tether.models.llm import LLMBackbone +from tether.models.llm.paligemma_backbone import PaliGemmaBackbone from tether.models.vision import VisionBackbone from tether.models.vlas.pi05 import Pi05VLA from tether.registry.components import VLAS @@ -231,6 +232,8 @@ def forward(self, images): class _StubPaliGemmaForPi05(LLMBackbone, nn.Module): """Stub PaliGemma — exposes the attrs Pi05VLA.predict_action uses.""" + embed_text_openpi = PaliGemmaBackbone.embed_text_openpi + def __init__( self, text_hidden: int = 2048, diff --git a/tests/test_pi0_vla.py b/tests/test_pi0_vla.py index 71271c60..4cd83103 100644 --- a/tests/test_pi0_vla.py +++ b/tests/test_pi0_vla.py @@ -23,6 +23,7 @@ from tether.models.base_vla import BaseVLA from tether.models.heads import VLAHead from tether.models.llm import LLMBackbone +from tether.models.llm.paligemma_backbone import PaliGemmaBackbone from tether.models.projectors import Projector from tether.models.vision import VisionBackbone from tether.models.vlas.pi0 import Pi0VLA @@ -271,6 +272,8 @@ class _StubPaliGemmaForPi0(LLMBackbone, nn.Module): """Stub PaliGemma exposing the property accessors Pi0VLA.predict_action relies on. Returns past_key_values that look like a transformers Cache.""" + embed_text_openpi = PaliGemmaBackbone.embed_text_openpi + def __init__( self, text_hidden: int = 2048, diff --git a/tests/test_snapflow_backend.py b/tests/test_snapflow_backend.py index 855baa10..6683e8dc 100644 --- a/tests/test_snapflow_backend.py +++ b/tests/test_snapflow_backend.py @@ -284,7 +284,6 @@ def fake_build_prefix(policy, obs): torch.zeros(1, 8, 256), canned_prefix_masks, torch.ones(1, 8, dtype=torch.long), ) - p_mock.model._prepare_attention_masks_4d.return_value = torch.ones(1, 1, 8, 8) p_mock.model.paligemma_with_expert.forward.return_value = (None, canned_past_kv) # state_proj.weight.dtype is read for the state-dtype cast; set it # to match canned_state so the cast is a no-op under mocks. @@ -347,7 +346,6 @@ def test_student_passes_target_time_through_denoise_step(self): torch.ones(1, 8, dtype=torch.bool), torch.ones(1, 8, dtype=torch.long), ) - student_mock.model._prepare_attention_masks_4d.return_value = torch.ones(1, 1, 8, 8) student_mock.model.paligemma_with_expert.forward.return_value = (None, [("fake_cache",)]) student_mock.model.paligemma_with_expert.paligemma.model.language_model.config = MagicMock() student_mock.model.paligemma_with_expert.gemma_expert.model.config = MagicMock()