From a519280db15680af6a0172912025191e11682cb6 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sat, 19 Sep 2026 19:02:55 +0000 Subject: [PATCH] perf: vectorized image augs and a torch.compile knob for HPT and Pi 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 .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) Claude-Session: https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw --- egomimic/algo/algo.py | 10 ++ egomimic/algo/hpt.py | 52 +++++-- egomimic/algo/pi.py | 25 +++- .../model/hpt_bc_mecka_6d_300M.yaml | 22 +++ egomimic/hydra_configs/model/pi0.5_base.yaml | 31 ++++- egomimic/models/denoising_nets.py | 10 +- egomimic/models/image_augs.py | 130 +++++++++++++++++- egomimic/pl_utils/pl_model.py | 29 ++++ egomimic/rldb/zarr/utils.py | 7 + .../rldb/zarr/zarr_dataset_action_expert.py | 4 +- egomimic/rldb/zarr/zarr_dataset_multi.py | 13 +- egomimic/trainHydra.py | 26 +++- egomimic/utils/compile_utils.py | 48 +++++++ tests/unit/test_image_augs.py | 109 +++++++++++++++ tests/unit/test_image_decode_dtype.py | 31 +++++ tests/unit/test_train_perf_wiring.py | 104 ++++++++++++++ 16 files changed, 625 insertions(+), 26 deletions(-) create mode 100644 egomimic/utils/compile_utils.py create mode 100644 tests/unit/test_image_augs.py create mode 100644 tests/unit/test_image_decode_dtype.py create mode 100644 tests/unit/test_train_perf_wiring.py diff --git a/egomimic/algo/algo.py b/egomimic/algo/algo.py index a01e4cc86..36b5e1c08 100644 --- a/egomimic/algo/algo.py +++ b/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 [] diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 7789f0323..f1ef9eca5 100644 --- a/egomimic/algo/hpt.py +++ b/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: diff --git a/egomimic/algo/pi.py b/egomimic/algo/pi.py index 243dd31e2..6f384bb53 100644 --- a/egomimic/algo/pi.py +++ b/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 ) diff --git a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml b/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml index 48dea6f4f..8555bb2de 100644 --- a/egomimic/hydra_configs/model/hpt_bc_mecka_6d_300M.yaml +++ b/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 diff --git a/egomimic/hydra_configs/model/pi0.5_base.yaml b/egomimic/hydra_configs/model/pi0.5_base.yaml index 8201d07d9..4a05b646e 100644 --- a/egomimic/hydra_configs/model/pi0.5_base.yaml +++ b/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 \ No newline at end of file +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 diff --git a/egomimic/models/denoising_nets.py b/egomimic/models/denoising_nets.py index 18471c59f..640734042 100644 --- a/egomimic/models/denoising_nets.py +++ b/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) diff --git a/egomimic/models/image_augs.py b/egomimic/models/image_augs.py index 0765dc352..497aaed94 100644 --- a/egomimic/models/image_augs.py +++ b/egomimic/models/image_augs.py @@ -2,13 +2,129 @@ torchvision's random transforms draw one parameter set per call, so applying a ``ColorJitter`` to a stacked (B, 3, H, W) batch jitters every image in the batch -by the same factors (and in the same op order). ``PerSampleAugs`` wraps the -configured transform and calls it once per image, so every sample draws its own -parameters. +by the same factors (and in the same op order). ``PerSampleAugs`` gives every +sample its own parameters. + +The obvious way to do that -- call the transform once per image -- costs one +Python round trip and a few dozen tiny CUDA launches *per image*: measured at +89 ms of a 196 ms HPT train step at batch 64 (H100, 2026-09-19), i.e. 46% of +the step, because every factor is a CPU scalar that has to be copied to the +device. ``BatchedColorJitter`` draws the factors as one on-device (B,) tensor +and runs each adjustment once over the whole batch instead. """ +from typing import List, Optional, Tuple + import torch from torch import nn +from torchvision import transforms as T +from torchvision.transforms import _functional_tensor as _ft + +_GRAY_W = (0.2989, 0.587, 0.114) + + +def _rgb_to_gray(img: torch.Tensor) -> torch.Tensor: + w = torch.tensor(_GRAY_W, dtype=img.dtype, device=img.device).view(3, 1, 1) + return (img * w).sum(dim=-3, keepdim=True) + + +def _blend(img1: torch.Tensor, img2: torch.Tensor, ratio: torch.Tensor) -> torch.Tensor: + return (ratio * img1 + (1.0 - ratio) * img2).clamp(0.0, 1.0) + + +class BatchedColorJitter(nn.Module): + """``torchvision.transforms.ColorJitter`` with per-sample factors, batched. + + Same adjustments and same factor distributions as torchvision, drawn once + per image instead of once per call. The op ORDER is drawn once per call + (torchvision's own behaviour), not once per image as the per-image loop + did; each adjustment clamps to [0, 1], so the order is not quite inert, but + it is a permutation of the same four factors either way. + + Input is (B, 3, H, W) float in [0, 1]. + """ + + def __init__( + self, + brightness: Optional[Tuple[float, float]] = None, + contrast: Optional[Tuple[float, float]] = None, + saturation: Optional[Tuple[float, float]] = None, + hue: Optional[Tuple[float, float]] = None, + ): + super().__init__() + self.brightness = brightness + self.contrast = contrast + self.saturation = saturation + self.hue = hue + + @classmethod + def from_color_jitter(cls, jitter: T.ColorJitter) -> "BatchedColorJitter": + return cls(jitter.brightness, jitter.contrast, jitter.saturation, jitter.hue) + + @staticmethod + def _factors(bounds, n, img) -> torch.Tensor: + # Drawn in float32: uniform_ into bfloat16 yields ~40 distinct values, + # and into uint8 it raises. + lo, hi = bounds + f = torch.empty(n, 1, 1, 1, dtype=torch.float32, device=img.device) + return f.uniform_(lo, hi).to(img.dtype) + + def _apply_brightness(self, img, f): + return (img * f).clamp(0.0, 1.0) + + def _apply_contrast(self, img, f): + mean = _rgb_to_gray(img).mean(dim=(-3, -2, -1), keepdim=True) + return _blend(img, mean, f) + + def _apply_saturation(self, img, f): + return _blend(img, _rgb_to_gray(img), f) + + def _apply_hue(self, img, f): + hsv = _ft._rgb2hsv(img) + h, s, v = hsv.unbind(dim=-3) + h = (h + f.squeeze(-3)) % 1.0 + return _ft._hsv2rgb(torch.stack((h, s, v), dim=-3)).clamp(0.0, 1.0) + + def forward(self, img: torch.Tensor) -> torch.Tensor: + n = img.shape[0] + ops = [ + (self.brightness, self._apply_brightness), + (self.contrast, self._apply_contrast), + (self.saturation, self._apply_saturation), + (self.hue, self._apply_hue), + ] + for idx in torch.randperm(len(ops)).tolist(): + bounds, fn = ops[idx] + if bounds is None: + continue + img = fn(img, self._factors(bounds, n, img)) + return img + + def extra_repr(self) -> str: + return ( + f"brightness={self.brightness}, contrast={self.contrast}, " + f"saturation={self.saturation}, hue={self.hue}" + ) + + +# Transforms that are already batch-safe AND sample no parameters, so running +# them on the stacked batch is identical to running them per image. +_DETERMINISTIC = (T.Normalize, T.Resize, T.CenterCrop, T.ConvertImageDtype) + + +def _vectorize(augs) -> Optional[List[nn.Module]]: + """The batched equivalent of ``augs``, or None if any member has no batched + form (then the caller keeps the per-image loop).""" + members = augs.transforms if isinstance(augs, T.Compose) else [augs] + out: List[nn.Module] = [] + for t in members: + if isinstance(t, T.ColorJitter): + out.append(BatchedColorJitter.from_color_jitter(t)) + elif isinstance(t, _DETERMINISTIC): + out.append(t) + else: + return None + return out class PerSampleAugs(nn.Module): @@ -22,9 +138,15 @@ class PerSampleAugs(nn.Module): def __init__(self, augs): super().__init__() self.augs = augs + vectorized = _vectorize(augs) + self.vectorized = nn.ModuleList(vectorized) if vectorized is not None else None def forward(self, images): - return torch.stack([self.augs(image) for image in images]) + if self.vectorized is None: + return torch.stack([self.augs(image) for image in images]) + for t in self.vectorized: + images = t(images) + return images def extra_repr(self): return repr(self.augs) diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index 662db1ddb..bfc201968 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/egomimic/pl_utils/pl_model.py @@ -103,6 +103,31 @@ def _instantiate_model(self, config_tree, norm_stats_state): norm_stats=norm_stats, ) + def _maybe_compile(self, config_tree) -> None: + """Hand the algo's hot submodules to ``torch.compile`` when + ``model.compile.enabled`` is set. + + Here, not in the algo's ``__init__``: this runs on every DDP rank in + ``setup("fit")``, before Lightning wraps the module, so each rank + compiles for the GPU it is about to land on. Not in ``__init__`` either: + ``load_from_checkpoint`` (robot rollout, eval) re-runs it from the saved + hparams and would compile for inference. Compilation itself is lazy -- + nothing is traced until the first batch. + """ + cfg = self._as_config(config_tree) + opts = cfg.model.get("compile") + if not opts or not opts.get("enabled", False): + return + compiled = self.model.compile_for_training( + mode=opts.get("mode"), dynamic=opts.get("dynamic") + ) + if not compiled: + print( + f"[compile] model.compile.enabled is set but " + f"{type(self.model).__name__} compiled nothing", + flush=True, + ) + # batch is now a dict, handle on model side def training_step(self, batch, batch_idx): self.train() @@ -374,6 +399,10 @@ def configure_optimizers(self) -> Dict[str, Any]: } return {"optimizer": optimizer} + def setup(self, stage: str) -> None: + if stage == "fit" and self.hparams.get("config_tree") is not None: + self._maybe_compile(self.hparams.config_tree) + def on_fit_start(self): self.model.device = self.device print( diff --git a/egomimic/rldb/zarr/utils.py b/egomimic/rldb/zarr/utils.py index 3c55af5ed..6f6bb1ca5 100644 --- a/egomimic/rldb/zarr/utils.py +++ b/egomimic/rldb/zarr/utils.py @@ -16,3 +16,10 @@ def set_global_seed(seed: int = 42): torch.cuda.manual_seed_all(seed) os.environ["PYTHONHASHSEED"] = str(seed) + + +def jpeg_pixels_to_float(decoded: np.ndarray) -> np.ndarray: + """``(..., H, W, 3)`` uint8 RGB -> ``(..., 3, H, W)`` float32 in [0, 1].""" + # float32: float64 doubles collate/pin/H2D bytes and is bit-identical here + # (see tests/unit/test_image_decode_dtype.py). + return np.moveaxis(decoded, -1, -3).astype(np.float32) / np.float32(255.0) diff --git a/egomimic/rldb/zarr/zarr_dataset_action_expert.py b/egomimic/rldb/zarr/zarr_dataset_action_expert.py index abb581a9e..05402df78 100644 --- a/egomimic/rldb/zarr/zarr_dataset_action_expert.py +++ b/egomimic/rldb/zarr/zarr_dataset_action_expert.py @@ -18,6 +18,7 @@ import torch from egomimic.rldb.embodiment.embodiment import get_embodiment_id +from egomimic.rldb.zarr.utils import jpeg_pixels_to_float from egomimic.rldb.zarr.zarr_dataset_multi import ( LocalEpisodeResolver, MultiDataset, @@ -90,8 +91,7 @@ def _load_obs_at(self, zarr_key: str, frame_idx: int): """Read a single obs key at one frame, decoding JPEG/JSON as needed.""" val = self.episode_reader.read({zarr_key: (frame_idx, None)})[zarr_key] if zarr_key in self._image_keys: - decoded = simplejpeg.decode_jpeg(val, colorspace="RGB") - return np.transpose(decoded, (2, 0, 1)) / 255.0 + return jpeg_pixels_to_float(simplejpeg.decode_jpeg(val, colorspace="RGB")) if zarr_key in self._json_keys: return self._decode_json_entry(val) return val diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index f23d0aad9..de201b1cd 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -2281,7 +2281,18 @@ def _next(reason: str, key: str = "") -> int: idx = _next("JPEG decode failed", key=k) retry = True break - data[k] = np.transpose(decoded, (2, 0, 1)) / 255.0 + # float32, not numpy's float64 default for `uint8 / float`: + # process_batch_for_training casts to float32 on the GPU + # anyway, so the extra 4 bytes a pixel buy nothing and + # double the collate, the pin and the H2D copy. At 360x640 + # and batch 64 that is 338 MB a camera a batch instead of + # 169. The result is bit-identical: there are only 256 + # 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( + np.float32 + ) / np.float32(255.0) elif zarr_key in self._json_keys: if isinstance(data[k], np.ndarray): data[k] = [self._decode_json_entry(v) for v in data[k]] diff --git a/egomimic/trainHydra.py b/egomimic/trainHydra.py index fab2b61ad..6f9ed07e8 100644 --- a/egomimic/trainHydra.py +++ b/egomimic/trainHydra.py @@ -88,6 +88,24 @@ def _build_model_config_tree(cfg: DictConfig) -> DictConfig: return tree +def _model_trainer_defaults(cfg: DictConfig) -> dict: + """Trainer kwargs a model config supplies: the gradient clip it was + published with (pi0.5: openpi's 1.0), unless ``trainer.gradient_clip_val`` + is set explicitly.""" + clip = cfg.model.get("gradient_clip_val") + if clip is None or cfg.trainer.get("gradient_clip_val") is not None: + return {} + return {"gradient_clip_val": clip} + + +def _model_wrapper_kwargs(cfg: DictConfig) -> dict: + """``ModelWrapper`` kwargs the model config supplies.""" + return { + "scheduler_interval": cfg.model.get("scheduler_interval", "step"), + "enable_grad_norm": cfg.model.get("enable_grad_norm", True), + } + + def _requeue_resume_path(cfg: DictConfig) -> Optional[str]: """``/last.ckpt`` when this process is a Slurm requeue (``SLURM_RESTART_COUNT`` > 0), else None. The dir is the ModelCheckpoint @@ -733,7 +751,7 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: model: LightningModule = ModelWrapper( config_tree=_build_model_config_tree(cfg), norm_stats_state=norm_stats.to_state(), - scheduler_interval=cfg.model.get("scheduler_interval", "step"), + **_model_wrapper_kwargs(cfg), ) _log_dataset_frame_counts( @@ -758,7 +776,11 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: plugins.append(SLURMEnvironment(requeue_signal=signal.SIGUSR1)) print("SLURM REQUEUE ENABLED") trainer: Trainer = hydra.utils.instantiate( - cfg.trainer, callbacks=callbacks, logger=logger, plugins=plugins or None + cfg.trainer, + callbacks=callbacks, + logger=logger, + plugins=plugins or None, + **_model_trainer_defaults(cfg), ) object_dict = { diff --git a/egomimic/utils/compile_utils.py b/egomimic/utils/compile_utils.py new file mode 100644 index 000000000..69a463170 --- /dev/null +++ b/egomimic/utils/compile_utils.py @@ -0,0 +1,48 @@ +"""torch.compile wiring shared by the algos. + +``nn.Module.compile`` (not ``torch.compile(module)``) is what the algos call: +it swaps the module's call implementation IN PLACE, so the module keeps its +identity and -- the part that matters here -- its ``state_dict`` keys. Wrapping +a module in ``torch.compile`` instead returns an ``OptimizedModule`` whose +parameters are all prefixed ``_orig_mod.``, which silently breaks every +checkpoint written before the flag was turned on. + +The other half of the contract: ``Module.compile`` only routes ``__call__``. +A submodule the algo invokes through a named method (``stem.compute_latent``, +``head.compute_loss``) runs uncompiled no matter what, so only modules that are +actually *called* are worth handing over. +""" + +from __future__ import annotations + +import logging +from typing import Iterable, List, Optional + +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +def compile_modules( + modules: Iterable[tuple[str, Optional[nn.Module]]], + mode: Optional[str] = None, + dynamic: Optional[bool] = None, +) -> List[str]: + """``module.compile(...)`` for each named module; returns the names done.""" + kwargs = {"dynamic": dynamic} + if mode: + kwargs["mode"] = mode + done = [] + for name, module in modules: + if not isinstance(module, nn.Module): + continue # a head with no sub-network, or a None slot + module.compile(**kwargs) + done.append(name) + if done: + logger.info( + "torch.compile (mode=%s, dynamic=%s): %s", + mode or "default", + dynamic, + ", ".join(done), + ) + return done diff --git a/tests/unit/test_image_augs.py b/tests/unit/test_image_augs.py new file mode 100644 index 000000000..4854370d5 --- /dev/null +++ b/tests/unit/test_image_augs.py @@ -0,0 +1,109 @@ +"""BatchedColorJitter must match torchvision's per-image adjustments, keep +per-sample randomness, and hand anything it cannot batch back to the loop.""" + +import pytest +import torch +import torchvision.transforms.functional as F +from torchvision import transforms as T + +from egomimic.models.image_augs import BatchedColorJitter, PerSampleAugs + +IMAGENET = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + + +@pytest.fixture +def images(): + torch.manual_seed(0) + return torch.rand(8, 3, 16, 16) + + +@pytest.mark.parametrize( + "kwargs, reference", + [ + ({"brightness": (0.7, 0.7)}, lambda x: F.adjust_brightness(x, 0.7)), + ({"contrast": (0.6, 0.6)}, lambda x: F.adjust_contrast(x, 0.6)), + ({"saturation": (1.3, 1.3)}, lambda x: F.adjust_saturation(x, 1.3)), + ({"hue": (0.04, 0.04)}, lambda x: F.adjust_hue(x, 0.04)), + ], +) +def test_matches_torchvision(images, kwargs, reference): + """A degenerate factor range pins every sample to one factor, so the batched + adjustment must equal torchvision's scalar one elementwise.""" + out = BatchedColorJitter(**kwargs)(images) + assert torch.allclose(out, reference(images), atol=1e-5) + + +def test_factors_are_per_sample(images): + jitter = BatchedColorJitter(brightness=(0.5, 1.5)) + ratios = (jitter(images) / images.clamp(min=1e-6)).flatten(1).median(dim=1).values + assert ratios.std() > 1e-3, "every sample got the same brightness factor" + + +def test_vectorizes_the_shipped_aug_list(): + augs = PerSampleAugs( + T.Compose( + [ + T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05), + IMAGENET, + ] + ) + ) + assert augs.vectorized is not None + assert isinstance(augs.vectorized[0], BatchedColorJitter) + + +def test_falls_back_for_unbatchable_transforms(images): + augs = PerSampleAugs(T.Compose([T.RandomHorizontalFlip(), IMAGENET])) + assert augs.vectorized is None + assert augs(images).shape == images.shape + + +def test_deterministic_aug_list_is_unchanged(images): + """Eval augs have no random member, so the batched path must be exact.""" + augs = PerSampleAugs(T.Compose([IMAGENET])) + assert augs.vectorized is not None + assert torch.allclose(augs(images), IMAGENET(images), atol=1e-6) + + +def test_output_stays_in_range(images): + augs = PerSampleAugs( + T.Compose( + [T.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.2)] + ) + ) + out = augs(images) + assert out.min() >= 0.0 and out.max() <= 1.0 + + +def test_factors_are_drawn_in_float32(): + """uniform_ into bfloat16 collapses 4096 draws onto a few dozen values, and + into uint8 it raises.""" + img = torch.zeros(4096, 3, 1, 1, dtype=torch.bfloat16) + f = BatchedColorJitter._factors((0.9, 1.1), 4096, img) + assert f.dtype == torch.bfloat16 + assert f.float().unique().numel() > 30 + BatchedColorJitter._factors( + (0.9, 1.1), 4, torch.zeros(4, 3, 1, 1, dtype=torch.uint8) + ) + + +def test_absent_camera_stays_zero_per_sample(): + """An all-zero sample is an absent camera: it must not be normalized to + -mean/std, even when the rest of its batch is present.""" + from types import SimpleNamespace + + from egomimic.algo.hpt import HPT + + stub = SimpleNamespace( + nets=SimpleNamespace(training=True), + train_image_augs=PerSampleAugs(T.Compose([T.ColorJitter(0.4), IMAGENET])), + eval_image_augs=None, + encoders={"front_img_1": object()}, + ) + torch.manual_seed(0) + images = torch.rand(3, 3, 8, 8) + images[1] = 0.0 + out = HPT._apply_image_augs(stub, images, "front_img_1") + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert not torch.equal(out[0], images[0]), "present samples are augmented" + assert torch.equal(HPT._apply_image_augs(stub, images, "not_an_encoder"), images) diff --git a/tests/unit/test_image_decode_dtype.py b/tests/unit/test_image_decode_dtype.py new file mode 100644 index 000000000..f35c37168 --- /dev/null +++ b/tests/unit/test_image_decode_dtype.py @@ -0,0 +1,31 @@ +"""Decoded frames leave the dataset as float32. + +`uint8 / 255.0` is float64 in numpy, which doubles what the collate, the +pinned-memory staging and the H2D copy move for every image in every batch -- +and the model casts to float32 on arrival regardless. +""" + +import numpy as np + +from egomimic.rldb.zarr.utils import jpeg_pixels_to_float + + +def test_float32_division_is_bit_identical_to_float64(): + """The cheap dtype is also the same number: there are only 256 inputs, so + this is exhaustive, not a sample.""" + u = np.arange(256, dtype=np.uint8) + wide = (u / 255.0).astype(np.float32) + narrow = u.astype(np.float32) / np.float32(255.0) + assert np.array_equal(wide, narrow) + + +def test_jpeg_pixels_to_float(): + decoded = np.random.randint(0, 256, size=(8, 6, 3), dtype=np.uint8) + out = jpeg_pixels_to_float(decoded) + assert out.dtype == np.float32 + assert out.shape == (3, 8, 6) + assert np.array_equal( + out, (np.transpose(decoded, (2, 0, 1)) / 255.0).astype(np.float32) + ) + window = np.stack([decoded, decoded]) + assert jpeg_pixels_to_float(window).shape == (2, 3, 8, 6) diff --git a/tests/unit/test_train_perf_wiring.py b/tests/unit/test_train_perf_wiring.py new file mode 100644 index 000000000..49cea555c --- /dev/null +++ b/tests/unit/test_train_perf_wiring.py @@ -0,0 +1,104 @@ +"""The torch.compile knob: off unless asked, in-place so checkpoints survive.""" + +import torch.nn as nn +from omegaconf import OmegaConf + +from egomimic.pl_utils.pl_model import ModelWrapper +from egomimic.utils.compile_utils import compile_modules + + +class _Tiny(nn.Module): + def __init__(self): + super().__init__() + self.lin = nn.Linear(4, 4) + + def forward(self, x): + return self.lin(x) + + +def test_compile_keeps_state_dict_keys(): + """nn.Module.compile is in-place; torch.compile(module) would prefix every + key with `_orig_mod.` and break checkpoints written before the flag.""" + module = _Tiny() + before = list(module.state_dict()) + compile_modules([("tiny", module)]) + assert list(module.state_dict()) == before + + +def test_compile_modules_skips_non_modules_and_reports_names(): + """A head without a sub-network hands back None (or something that is not a + Module); compiling it must be a skip, not a crash.""" + assert compile_modules([("a", _Tiny()), ("b", None), ("c", object())]) == ["a"] + + +class _SpyAlgo: + def __init__(self): + self.calls = [] + self.nets = nn.ModuleDict({"policy": _Tiny()}) + + def compile_for_training(self, mode=None, dynamic=None): + self.calls.append((mode, dynamic)) + return ["policy"] + + +def _wrapper_with(compile_cfg, monkeypatch): + algo = _SpyAlgo() + monkeypatch.setattr(ModelWrapper, "_instantiate_model", lambda *a, **k: algo) + model = {"robomimic_model": {}} + if compile_cfg is not None: + model["compile"] = compile_cfg + wrapper = ModelWrapper( + config_tree=OmegaConf.create({"model": model}), norm_stats_state=None + ) + wrapper.setup("fit") + return algo + + +def test_compile_is_off_unless_asked(monkeypatch): + assert _wrapper_with(None, monkeypatch).calls == [] + assert _wrapper_with({"enabled": False, "mode": None}, monkeypatch).calls == [] + + +def test_compile_only_when_fitting(monkeypatch): + """load_from_checkpoint (robot rollout, eval) re-runs __init__ from the saved + hparams; it must not compile the policy for inference.""" + algo = _SpyAlgo() + monkeypatch.setattr(ModelWrapper, "_instantiate_model", lambda *a, **k: algo) + tree = {"model": {"robomimic_model": {}, "compile": {"enabled": True}}} + wrapper = ModelWrapper(config_tree=OmegaConf.create(tree), norm_stats_state=None) + assert algo.calls == [] + wrapper.setup("validate") + assert algo.calls == [] + wrapper.setup("fit") + assert algo.calls == [(None, None)] + + +def test_compile_forwards_mode_and_dynamic(monkeypatch): + cfg = {"enabled": True, "mode": "reduce-overhead", "dynamic": True} + assert _wrapper_with(cfg, monkeypatch).calls == [("reduce-overhead", True)] + + +def test_enable_grad_norm_reaches_the_wrapper(compose_resolve): + from egomimic.trainHydra import _model_wrapper_kwargs + + cfg = compose_resolve( + "train_zarr_cartesian", ["model=pi0.5_base", "+model.enable_grad_norm=false"] + ) + assert _model_wrapper_kwargs(cfg)["enable_grad_norm"] is False + cfg = compose_resolve("train_zarr_cartesian", ["model=pi0.5_base"]) + assert _model_wrapper_kwargs(cfg)["enable_grad_norm"] is True + + +def test_pi_clips_like_openpi_unless_the_trainer_says_otherwise(compose_resolve): + from egomimic.trainHydra import _model_trainer_defaults + + cfg = compose_resolve("train_zarr_cartesian", ["model=pi0.5_base"]) + assert cfg.model.gradient_clip_val == 1.0 + assert cfg.model.get("enable_grad_norm", True) is True + assert _model_trainer_defaults(cfg) == {"gradient_clip_val": 1.0} + + cfg.trainer.gradient_clip_val = 0.5 + assert _model_trainer_defaults(cfg) == {} + + hpt = compose_resolve("train_zarr_cartesian", []) + assert _model_trainer_defaults(hpt) == {}