perf: vectorized image augs and a torch.compile knob for HPT and Pi; Pi clips gradients like openpi by RyanPCo · Pull Request #676 · GaTech-RL2/EgoVerse

RyanPCo · 2026-09-20T16:14:55Z

Makes the HPT and Pi training step faster: batched image augs, an opt-in torch.compile flag, and a few smaller fixes. It also changes how Pi trains (see below).

Behaviour changes

Pi now clips gradients at 1.0 and keeps logging the grad norm, as openpi does. Pi runs before this PR trained unclipped, so don't compare Pi runs across it as if only the named variable changed.
- enable_grad_norm was never passed to ModelWrapper, so the false in pi0.5_base.yaml did nothing. It is wired through now.
- Pi drops the false, so it keeps logging Train/policy_grad_norms_*.
- Pi sets gradient_clip_val: 1.0 in its model config. trainHydra hands that to the Trainer (_model_trainer_defaults) unless trainer.gradient_clip_val is set explicitly.
Batched ColorJitter draws the op order once per batch, not once per image. The factors are still per-sample.
Absent cameras are now detected per sample. The old all-zero guard ran one all() over the whole batch, so a batch mixing present and absent samples normalized the absent ones to -mean/std. The check moved into _apply_image_augs and runs per sample, on the device. The augs now run on every sample, so RNG use no longer depends on the data.

Speed (one H100, 360x640 mecka frames)

	before	after
HPT 300M, batch 64, batched augs	245 ms/step	173 ms
+ `model.compile.enabled=true`		109 ms (2.24x)
Pi, batch 16, compiled	827 ms/step	588 ms (1.41x)

Batched augs. BatchedColorJitter draws the factors as one on-device tensor instead of making 64 per-image Python round trips. A degenerate factor range matches torchvision exactly (tests/unit/test_image_augs.py). Anything it can't batch falls back to the per-image loop.
Compile flag. model.compile.enabled=true calls <algo>.compile_for_training(), which uses nn.Module.compile so state_dict keys keep their names; torch.compile(module) would prefix them _orig_mod. and break existing checkpoints. It is off by default, and the compiled fp32 loss was bit-identical to eager.
Smaller fixes:
- Attention uses need_weights=False, so it dispatches to SDPA.
- Frames decode to float32 instead of float64, halving collate and host-to-device traffic (338 → 169 MB per camera per batch, bit-identical).
- One device→host sync per camera per step is removed.
- The batch is no longer deep-copied every step for an OT loss that is off.
- Jitter factors are drawn in float32, since uniform_ into bf16 gives only ~40 distinct values.

Accuracy check

Measured with one-off harnesses that are kept out of this PR:

The old attention path and SDPA give the same fp32 step to 1.3e-6 relative, with gradients at cosine ≥ 0.999999.
Each sample's aug distribution is unchanged (two-sample KS over 8 statistics, worst p = 0.61).
A 2000-step A/B at two seeds shows no regression (held-out loss −9.0% and −23.2%). This rules out a regression; it is not evidence the rewrite helps. Two seeds can't carry that, and the arms consume different amounts of RNG.

Notes

What didn't pay: persistent_workers, pin_memory, more workers, max-autotune.
CUDA-graph compile modes (reduce-overhead, max-autotune) are where a rank can wedge in cudagraph_trees and keep its GPUs after a cancel. The compile config comments say so.

Tests

New and updated: test_image_augs.py (per-sample absent mask, float32 factors), test_image_decode_dtype.py, and test_train_perf_wiring.py (enable_grad_norm wiring, Pi clip default).
The full unit suite passes on transformers 4.53.2.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

RyanPCo · 2026-09-20T16:15:29Z

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

This stack of pull requests is managed by Graphite. Learn more about stacking.

RyanPCo

Automated adversarial review (read-only; findings were checked by a second skeptical pass before posting): 1 medium, 4 low, 2 nit.

RyanPCo · 2026-10-30T20:55:34Z


        if config_tree is not None:
            self.model = self._instantiate_model(config_tree, norm_stats_state)
+            self._maybe_compile(config_tree)


medium (correctness): _maybe_compile runs in __init__ from config_tree, which save_hyperparameters persists. ModelWrapper.load_from_checkpoint in robot/rollout.py therefore re-applies model.compile.enabled=true for any HPT checkpoint trained with the flag. The trunk, encoders and head denoiser are then compiled with inductor on the first real-time inference call: about 40 s or more, plus recompiles for batch-1 and eval mode, and the robot box needs a working triton toolchain. PI is unaffected, because rollout goes through sample_actions, which Module.compile does not route. Consider applying compilation only for training, for example by calling compile_for_training from trainHydra after constructing the wrapper or in setup(stage="fit"), or by forcing compile.enabled=false in rollout's checkpoint patching.

RyanPCo · 2026-10-30T20:55:34Z

+    dynamic: bool = False,
+) -> List[str]:
+    """``module.compile(...)`` for each named module; returns the names done."""
+    kwargs = {"dynamic": dynamic}


low (performance): dynamic is always passed as a bool. The yaml defaults to false, and _maybe_compile coerces null to False, so torch.compile's default automatic-dynamic mode (dynamic=None) can't be selected. No train loader sets drop_last, so the final partial batch, validation batches and each embodiment's batch/token shape in cotraining each trigger a full static recompile (about 40 s for HPT). With enough shapes, dynamo's recompile limit sends those modules back to eager without warning. Consider passing dynamic through as-is with a null default, or document the per-shape recompiles next to the ~40 s / ~4 min figures in the yaml.

RyanPCo · 2026-10-30T20:55:35Z

+                    # possible values and float32(u)/float32(255) equals
+                    # float32(u / 255.0) for every one of them (see
+                    # tests/unit/test_image_decode_dtype.py).
+                    data[k] = np.transpose(decoded, (2, 0, 1)).astype(


low (duplication): ZarrActionExpertDataset._load_obs_at (zarr_dataset_action_expert.py:94) still decodes with / 255.0 and returns float64, so the two zarr datasets now emit different dtypes for the same frame. Consider a small shared decode_jpeg_chw_float32(buf) helper and calling it from both sites. The 9-line comment could then shrink to one line pointing at the bit-identity test.

RyanPCo · 2026-10-30T20:55:35Z

+    assert np.array_equal(wide, narrow)
+
+
+def test_decode_path_emits_float32():


low (test-quality): test_decode_path_emits_float32 copies the expression instead of running the dataset code, so it still passes if ZarrDataset.getitem reverts to / 255.0. Its docstring also names MultiDataset._read_keys, which does not exist; the code is in ZarrDataset.getitem. Test a shared decode helper directly, or drop this test and keep the exhaustive bit-identity one.

RyanPCo · 2026-10-30T20:55:35Z

+# Off by default -- a run that never reaches steady state pays the compile for
+# nothing. Per run: model.compile.enabled=true. `mode` goes straight to
+# torch.compile; null is the default inductor mode the numbers were taken with.
+compile:


low (duplication): Only this file and hpt_bc_mecka_6d_300M.yaml define compile:. The default model (hpt_bc_flow_eva) and the other HPT configs, which HPT.compile_for_training would also speed up, do not, so the documented model.compile.enabled=true fails under Hydra struct mode for them and needs +model.compile.enabled=true. Consider putting the block in a shared place (e.g. hpt_bc_flow_eva / a defaults include) or documenting the + spelling once instead of repeating the long comment in each YAML.

RyanPCo · 2026-10-30T20:55:35Z

+        before DDP wraps anything -- inductor then sees one module at a time and
+        DDP's bucketing is unaffected. Algos that have nothing worth compiling
+        leave this alone; the default no-ops so a run that turns the flag on
+        against such an algo is slower to start, not broken.


nit (docs): The default compiles nothing, so a run with the flag on against such an algo starts exactly like an uncompiled one, not slower (ModelWrapper prints a warning). Suggest ending the docstring with: 'Returns the names of the modules compiled; the default compiles nothing and returns [].'

RyanPCo · 2026-10-30T20:55:35Z

+    cfg = {"enabled": True, "mode": "reduce-overhead", "dynamic": True}
+    assert _wrapper_with(cfg, monkeypatch).calls == [("reduce-overhead", True)]
+
+


nit (test-quality): This greps trainHydra's source for an exact kwarg string. A reformat or rename breaks it, and it checks nothing about behavior. Consider moving the ModelWrapper kwargs into a helper, as _model_trainer_defaults already does for the clip, and asserting on its output. Otherwise drop the test. The docstring should say what is guaranteed, not what used to happen.

The 300M HPT train step was launch-bound, not compute-bound: ~14,600 CUDA kernel launches per step, 230 ms of CPU against 111 ms of GPU. Two things dominated, and both are fixed here. PerSampleAugs called torchvision's ColorJitter once per image so each sample drew its own factors -- 64 Python round trips a step, each launching a few dozen kernels and copying its factors from the host. BatchedColorJitter draws the factors as one on-device tensor and runs each adjustment once over the batch; a degenerate factor range reproduces torchvision's scalar adjustment exactly (tests/unit/test_image_augs.py). Anything it cannot batch still goes through the per-image loop. model.compile.enabled=true routes to <algo>.compile_for_training(), which uses nn.Module.compile rather than torch.compile(module) so state_dict keys survive -- the wrapper form renames every parameter _orig_mod.* and would break checkpoints written before the flag. That only routes __call__, so HPT compiles the trunk, the flow head's denoiser and the image encoders (the stems and heads are reached through named methods), and PI's call site moves to __call__. On one H100 at batch 64 and the 360x640 frames mecka decodes to: HPT 245 -> 173 ms from the augs, -> 109 ms compiled (2.24x). Pi at batch 16: 827 -> 588 ms compiled (1.41x). In fp32 the compiled loss is bit-identical to eager. Also here: the flow head's attentions pass need_weights=False so they dispatch to scaled_dot_product_attention; frames decode to float32 instead of numpy's float64 default, halving the collate and H2D traffic (338 -> 169 MB a camera a batch, bit-identical); a per-camera per-step device->host sync is gone; and the batch is no longer deep-copied every step for an OT loss that is off. The all-zero (absent camera) guard moves into _apply_image_augs and becomes per-sample: the old test was one all() over the batch, so a batch mixing present and absent samples normalized the absent ones to -mean/std. The augs now run on every sample and the selector keeps the absent ones at zero, so RNG use no longer depends on the data. Jitter factors are drawn in float32 (uniform_ into bfloat16 gives ~40 distinct values; into uint8 it raises). enable_grad_norm was never passed to ModelWrapper, so pi0.5_base's false did nothing. It is wired now, and Pi follows openpi instead of opting out: it keeps logging the pre-clip grad norm every step and clips at openpi's 1.0, set as gradient_clip_val in the model config and handed to the Trainer unless trainer.gradient_clip_val is set. Pi recipes trained unclipped until now, so this is a behaviour change for every Pi run. One grad-norm pass remains (the second was removed with the log-only MAD detector). The compile comments note that CUDA-graph modes are where a rank can wedge in cudagraph_trees and keep its GPUs after a cancel. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

RyanPCo · 2026-10-30T23:09:26Z

Review findings: status after the fixup pass. The fixes are squashed into this PR's commits, and the whole stack was force-pushed.

Fixed

Compile flag in hparams recompiles on rollout load_from_checkpoint: compile moved from init to setup('fit'); test checks that init and validate do not compile
dynamic=False forces a recompile per batch shape: dynamic now defaults to None (automatic dynamic shapes) in configs, compile_modules and the algo signatures; yaml comment explains why
compile_for_training docstring wrong about the no-op case: docstring now gives the return contract and says the default compiles nothing and returns []
float32 decode fix applied to one of two decode sites: shared jpeg_pixels_to_float in rldb/zarr/utils.py, used by the action-expert dataset (676 fixup) and ZarrDataset (in the c952 fixup, which last edited that line); comment cut to one line. rldb: move MultiDataset and format-neutral pieces out of zarr/ #717 moves set_global_seed out to utils/seed.py and leaves jpeg_pixels_to_float in rldb/zarr/utils.py
mirror test re-types the production expression: mirror test replaced by a test of jpeg_pixels_to_float (dtype, values, window shape)
enable_grad_norm test greps trainHydra source: new _model_wrapper_kwargs helper; test asserts its output for pi0.5_base with and without enable_grad_norm=false

Partial

compile block duplicated, missing from other model configs: compile block comments now document +model.compile.enabled=true for configs without the block; no shared defaults group added

RyanPCo force-pushed the ryanco/zarr-read-coalescing branch from 732de9b to ec36005 Compare September 20, 2026 19:09

RyanPCo force-pushed the ryanco/torch-compile-train branch 2 times, most recently from 43cdc6c to e10ccbe Compare September 21, 2026 18:45

RyanPCo force-pushed the ryanco/zarr-read-coalescing branch from 183054e to 796c29c Compare September 22, 2026 13:18

RyanPCo force-pushed the ryanco/torch-compile-train branch from e10ccbe to 4e35316 Compare September 22, 2026 13:18

RyanPCo changed the title ~~perf: vectorized image augs and a torch.compile knob for HPT and Pi~~ perf: vectorized image augs and a torch.compile knob for HPT and Pi; Pi clips gradients like openpi Sep 22, 2026

RyanPCo force-pushed the ryanco/zarr-read-coalescing branch from 7f5dcda to 2a00698 Compare October 12, 2026 21:50

RyanPCo force-pushed the ryanco/torch-compile-train branch from 446a560 to fcdbf6e Compare October 12, 2026 23:19

RyanPCo force-pushed the ryanco/zarr-read-coalescing branch from 2a00698 to ee8dfd2 Compare October 12, 2026 23:19

RyanPCo mentioned this pull request Oct 19, 2026

lance migration #716

Open

RyanPCo force-pushed the ryanco/zarr-read-coalescing branch from ee8dfd2 to 76b71f7 Compare October 19, 2026 19:49

RyanPCo force-pushed the ryanco/torch-compile-train branch 2 times, most recently from f1dee44 to d3e5b0d Compare October 19, 2026 20:22

RyanPCo force-pushed the ryanco/torch-compile-train branch from d3e5b0d to f20e8aa Compare October 30, 2026 18:27

RyanPCo force-pushed the ryanco/zarr-read-coalescing branch from 119657d to bbbca3e Compare October 30, 2026 18:27

RyanPCo commented Oct 30, 2026

View reviewed changes

RyanPCo force-pushed the ryanco/torch-compile-train branch from f20e8aa to a519280 Compare October 30, 2026 23:05

RyanPCo force-pushed the ryanco/zarr-read-coalescing branch from bbbca3e to 9ce191d Compare October 30, 2026 23:05

RyanPCo mentioned this pull request Oct 30, 2026

augs: per-sample image augmentation for HPT train-mode augs #651

Draft

Navigation Menu

Repository navigation

perf: vectorized image augs and a torch.compile knob for HPT and Pi; Pi clips gradients like openpi - #676

perf: vectorized image augs and a torch.compile knob for HPT and Pi; Pi clips gradients like openpi#676
RyanPCo wants to merge 1 commit into
ryanco/zarr-read-coalescingfrom
ryanco/torch-compile-train

RyanPCo commented Sep 20, 2026 •

edited

Loading

Uh oh!

RyanPCo commented Sep 20, 2026 •

edited

Loading

Uh oh!

RyanPCo left a comment

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo Oct 30, 2026

Uh oh!

RyanPCo commented Oct 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		assert np.array_equal(wide, narrow)


		def test_decode_path_emits_float32():

		cfg = {"enabled": True, "mode": "reduce-overhead", "dynamic": True}
		assert _wrapper_with(cfg, monkeypatch).calls == [("reduce-overhead", True)]

Conversation

RyanPCo commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Behaviour changes

Speed (one H100, 360x640 mecka frames)

Accuracy check

Notes

Tests

Uh oh!

RyanPCo commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

RyanPCo left a comment

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

RyanPCo commented Oct 30, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

RyanPCo commented Sep 20, 2026 •

edited

Loading

RyanPCo commented Sep 20, 2026 •

edited

Loading