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

egomimic/algo/algo.py

@@ -72,3 +72,13 @@ def log_info(self, info):
             loss_log (dict): name -> summary statistic
         """
         raise NotImplementedError("Must implement log_info in subclass")
+
+    def compile_for_training(self, mode=None, dynamic=None):
+        """Hand the hot submodules to ``torch.compile``.
+
+        Called once by ``ModelWrapper`` when ``model.compile.enabled`` is set,
+        before DDP wraps anything -- inductor then sees one module at a time and
+        DDP's bucketing is unaffected. Returns the names of the modules
+        compiled; the default compiles nothing and returns [].
+        """
+        return []

egomimic/algo/hpt.py

@@ -26,6 +26,7 @@
 )
 from egomimic.models.image_augs import per_sample
 from egomimic.rldb.embodiment.embodiment import get_embodiment, get_embodiment_id
+from egomimic.utils.compile_utils import compile_modules
 from egomimic.utils.hf_utils import download_from_huggingface
 from egomimic.utils.tensor_utils import EinOpsRearrange, get_sinusoid_encoding_table
 
@@ -1031,6 +1032,23 @@ def __init__(
 
         self.training_step = 0
 
+    @override
+    def compile_for_training(self, mode=None, dynamic=None):
+        """Compile the trunk, the flow head's denoiser and the image encoders.
+
+        Only these three are invoked through ``__call__``; the stems run via
+        ``compute_latent`` and the heads via ``compute_loss``, which
+        ``Module.compile`` does not route (and which measured no gain anyway).
+        """
+        policy = self.nets["policy"]
+        targets = [("trunk", policy.trunk["trunk"] if not policy.no_trunk else None)]
+        targets += [
+            (f"head[{name}].model", getattr(head, "model", None))
+            for name, head in policy.heads.items()
+        ]
+        targets += [(f"encoder[{name}]", enc) for name, enc in policy.encoders.items()]
+        return compile_modules(targets, mode=mode, dynamic=dynamic)
+
     @override
     def process_batch_for_training(self, batch):
         """
@@ -1098,7 +1116,11 @@ def process_batch_for_training(self, batch):
             # TODO make this work with any fp type
             for key, value in processed_batch[embodiment_id].items():
                 if isinstance(value, torch.Tensor):
-                    value = value.to(self.device)
+                    # non_blocking: a no-op from pageable memory, and from a
+                    # pin_memory loader it lets the copy overlap the previous
+                    # step. Everything downstream is queued on the same stream,
+                    # so ordering holds without an explicit wait.
+                    value = value.to(self.device, non_blocking=True)
                     if value.is_floating_point():
                         value = value.float()
                     processed_batch[embodiment_id][key] = value
@@ -1142,7 +1164,10 @@ def forward_training(self, batch):
                 "domain": embodiment_name,  # readability on config side
                 "data": data,
             }
-            hpt_batches[embodiment_id] = self._clone_batch(hpt_batch)
+            # Only the OT loss reads this second copy; cloning every image
+            # tensor each step for a loss that is off is pure allocation.
+            if self.ot:
+                hpt_batches[embodiment_id] = self._clone_batch(hpt_batch)
 
             if self.freeze_repr:
                 loss = self.nets["policy"].compute_loss_depth(
@@ -1323,12 +1348,23 @@ def _forward_ot(self, batch, embodiment1_id, embodiment2_id):
         )
 
     def _apply_image_augs(self, images, short):
-        """Augment one encoded camera's (B, 3, H, W) batch with the mode's augs."""
+        """Augment one encoded camera's (B, 3, H, W) batch with the mode's augs.
+
+        An absent camera arrives as an all-zero image and must stay zero
+        (normalizing it would make it -mean/std). The test is per sample, so a
+        batch mixing present and absent samples is covered, and it is a
+        torch.where rather than an `if`, which would read the predicate on the
+        host and serialise the step against the next batch's copy.
+        """
         if short not in self.encoders:
             return images
         # No eval-aug fallback in train mode: train_image_augs=None means none.
         augs = self.train_image_augs if self.nets.training else self.eval_image_augs
-        return augs(images) if augs is not None else images
+        if augs is None:
+            return images
+        absent = (images == 0).flatten(1).all(1)
+        out = augs(images)
+        return torch.where(absent.view(-1, *[1] * (out.dim() - 1)), 0.0, out)
 
     def _stem_history_len(self, domain: str, modality: str) -> int:
         """``history_len`` of ``domain``'s ``modality`` stem (1 when the stem
@@ -1381,11 +1417,9 @@ def _robomimic_to_hpt_data(
         for key in cam_keys:
             if key in batch:
                 short = key.rsplit(".", 1)[-1]
-                _data = batch[key]
-                if not torch.all(_data == 0):
-                    _data = self._apply_image_augs(_data, short)
-
-                data[short] = _data.unsqueeze(1).unsqueeze(1)
+                data[short] = (
+                    self._apply_image_augs(batch[key], short).unsqueeze(1).unsqueeze(1)
+                )
 
         for key in lang_keys:
             if key in batch:

egomimic/algo/pi.py

@@ -30,6 +30,7 @@
     ConverterRegistry,
     pad_to_width,
 )
+from egomimic.utils.compile_utils import compile_modules
 
 logger = logging.getLogger(__name__)
 # Ensure logger propagates to root logger and has appropriate level
@@ -424,7 +425,9 @@ def process_batch_for_training(self, batch):
 
             for key, value in processed_batch[embodiment_id].items():
                 if isinstance(value, torch.Tensor):
-                    value = value.to(self.device)
+                    # See HPT.process_batch_for_training: no-op from pageable
+                    # memory, overlapped copy from a pin_memory loader.
+                    value = value.to(self.device, non_blocking=True)
                     if value.is_floating_point():
                         value = value.float()
                     processed_batch[embodiment_id][key] = value
@@ -437,6 +440,20 @@ def process_batch_for_training(self, batch):
 
         return processed_batch
 
+    @override
+    def compile_for_training(self, mode=None, dynamic=None):
+        """Compile openpi's ``PI0Pytorch.forward`` (the flow-matching training
+        pass: SigLIP tower, PaliGemma prefix and the action expert).
+
+        ``sample_actions`` is already compiled by openpi itself, so only the
+        training pass is left. The observation preprocessing at the top of that
+        forward is Python-level dict work, so expect dynamo to graph-break there
+        and compile the transformer stack behind it.
+        """
+        return compile_modules(
+            [("pi0", self.nets["policy"])], mode=mode, dynamic=dynamic
+        )
+
     @override
     def forward_training(self, batch):
         """
@@ -462,7 +479,9 @@ def forward_training(self, batch):
                 embodiment_name,
             )
 
-            losses = self.nets["policy"].forward(processed_obs, action)
+            # __call__, not .forward: nn.Module.compile only routes __call__, so
+            # a direct .forward() would silently run uncompiled.
+            losses = self.nets["policy"](processed_obs, action)
             loss = self._reduce_loss(losses, _batch[ac_key], embodiment_id, ac_key)
 
             predictions[f"{embodiment_name}_{ac_key}"] = _batch[ac_key]
@@ -500,7 +519,7 @@ def forward_eval(self, batch):
                 )
 
                 # Flow-matching val loss — same call as forward_training.
-                losses = self.nets["policy"].forward(processed_obs, action)
+                losses = self.nets["policy"](processed_obs, action)
                 unnorm_preds[f"{embodiment_name}_loss"] = self._reduce_loss(
                     losses, _batch[ac_key], embodiment_id, ac_key
                 )

egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml

@@ -183,3 +183,25 @@ optimizer:
   weight_decay: 0.0001
 
 scheduler: null
+
+# torch.compile for the training step: ModelWrapper._maybe_compile hands the
+# trunk, the flow head's denoiser and the image encoders to torch.compile
+# (HPT.compile_for_training) before DDP wraps anything. Measured on one H100 at
+# batch 64 and the 360x640 frames mecka decodes to (2026-09-19): 173 ms/step ->
+# 109 ms (1.58x), for ~40 s of compile at the first batch. mode="reduce-overhead"
+# is worth a further ~5% for roughly double that compile, but CUDA-graph modes
+# (reduce-overhead, max-autotune) are where a rank can wedge in inductor's
+# cudagraph_trees, ignore SIGTERM and keep its GPUs after a cancel.
+#
+# 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.
+# dynamic: null is torch.compile's automatic dynamic shapes -- a new batch shape
+# (the last partial batch, validation, another embodiment's batch) recompiles
+# once into a shape-generic graph; false would recompile per shape until
+# dynamo's recompile limit drops the module back to eager.
+# Model configs without this block take +model.compile.enabled=true.
+compile:
+  enabled: false
+  mode: null
+  dynamic: null

egomimic/hydra_configs/model/pi0.5_base.yaml

@@ -58,7 +58,11 @@ robomimic_model:
         size: 224
         interpolation: 3
 
-enable_grad_norm: false
+# openpi's trainers clip the global grad norm at 1.0 (AdamW
+# clip_gradient_norm) and log the pre-clip norm every step. trainHydra hands this
+# to the Trainer unless trainer.gradient_clip_val is set; ModelWrapper logs the
+# norm as Train/policy_grad_norms_raw (enable_grad_norm: false turns that off).
+gradient_clip_val: 1.0
 
 optimizer:
   _target_: torch.optim.AdamW
@@ -74,4 +78,27 @@ scheduler:
   _partial_: true
   num_warmup_steps: 2000
 
-scheduler_interval: step
+scheduler_interval: step
+
+# torch.compile for the training step: ModelWrapper._maybe_compile compiles
+# openpi's PI0Pytorch.forward (PI.compile_for_training). Measured on one H100 at
+# batch 16 (2026-09-19): 827 ms/step -> 588 ms (1.41x), for ~4 min of compile at
+# the first batch. `sample_actions` is compiled by openpi itself and is not
+# affected by this flag.
+#
+# CUDA-graph modes (reduce-overhead, max-autotune) are where a rank can wedge in
+# inductor's cudagraph_trees, ignore SIGTERM and keep its GPUs after a cancel;
+# openpi already compiles sample_actions with max-autotune whatever this says.
+#
+# 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.
+# dynamic: null is torch.compile's automatic dynamic shapes -- a new batch shape
+# (the last partial batch, validation, another embodiment's batch) recompiles
+# once into a shape-generic graph; false would recompile per shape until
+# dynamo's recompile limit drops the module back to eager.
+# Model configs without this block take +model.compile.enabled=true.
+compile:
+  enabled: false
+  mode: null
+  dynamic: null

egomimic/models/denoising_nets.py

@@ -575,7 +575,7 @@ def forward(self, x, t, cond, timestep_table):
         res = x
         x = self.ln1(x)
         x = self.aln1(x, t, cond, timestep_table)
-        (x, _) = self.mha(x, x, x)
+        (x, _) = self.mha(x, x, x, need_weights=False)
         x = self.alns1(x, t, cond, timestep_table)
         x = x + res
         res = x
@@ -617,13 +617,17 @@ def __init__(
         self.cond_proj = nn.Linear(cond_dim, hidden_dim)
 
     def forward_cross(self, x, cond):
+        # need_weights=False: the default True makes nn.MultiheadAttention take
+        # the slow path -- it materialises the B x H x L x L scores and averages
+        # them over heads to return weights nobody here reads -- instead of
+        # dispatching to scaled_dot_product_attention. The outputs are the same.
         res = x
         x = self.ln1(x)
-        x, _ = self.mha(x, x, x)
+        x, _ = self.mha(x, x, x, need_weights=False)
         x = x + res
         res = x
         x = self.ln2(x)
-        x, _ = self.cmha(x, cond, cond)
+        x, _ = self.cmha(x, cond, cond, need_weights=False)
         x = x + res
         res = x
         x = self.ln3(x)