[RyanPCo]

HPT.init defaulted self.device to torch.device('cuda') and moved the
freshly built policy there. Under submitit every rank sees all 8 GPUs, so all
8 ranks resolved that to cuda:0 and built their policy on the same device:
8 x 3.34B, which OOMs during construction. The 1B arm fit and survived, which
made the failure look config-specific.

PI already builds on CPU and leaves self.device unset for Lightning to fill in
through ModelWrapper.on_fit_start / on_validation_start. HPT now does the same:
no device is chosen in init and nothing leaves CPU there.

Lightning's write then has to reach the HPTModel the algo holds. That model
reads self.device to build the loss accumulators in compute_loss, while
HPT.compute_losses accumulates on the algo's, so a write that stopped at the
algo would leave the two on different devices -- 'Expected all tensors to be
on the same device, but found at least two devices, cuda:0 and cpu!'. That was
7 of the unit tests (4 test_hpt_train_step, 2 test_norm_cache, 1
test_resolve_once) whenever the suite ran on a GPU node. HPT.device is
therefore a property whose setter also writes nets['policy'].device, so the two
cannot diverge through the path the trainer uses.

Note this sets a device, it does not move anything: the parameters follow
because Lightning moves the LightningModule itself. rollout.py::_load_policy is
the only construction path outside Lightning and already does both its own
.to(policy_device) and the device write, which now reaches the HPTModel too.

Verified on 8xH100 with the 3.34B hpt_3b arch-comparison arm: every rank
reports CUDA_VISIBLE_DEVICES [0..7], all 8 reach fit start and train an epoch
without an OOM.

---

Folded in (to keep the stack under Graphite's 50-branch cap):

• algo: DeterministicEvalMixin, the seeded eval helpers out of hpt.py #665 — algo: DeterministicEvalMixin, the seeded eval helpers out of hpt.py

🤖 Generated with Claude Code

https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

[RyanPCo]

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

• hpt: build the policy on CPU and let Lightning place it #664 : 2 dependent PRs (#665 , #666 ) 👈 (View in Graphite)
• hpt: never re-initialise pretrained stems; hash-check them; mask text padding #645 : 1 other dependent PR (#646 )
• obs: proprio history (K proprio steps, one token each, learned time embedding) #650
• loss: mask the repeat-last action padding out of the HPT loss #649
• resnet: keep a frozen backbone in eval mode; give backbones their own LR group #652
• augs: per-sample image augmentation for HPT train-mode augs #651
• eval: deterministic validation (seeded prompts + flow sampling); log-only grad-norm MAD #644
• data: mecka_fold_flagship_top3_hpt_6d, the 3-operator flagship train split #643
• val: subsampled metric loaders + pinned contiguous video loaders #642
• infra: mecka 6D / keypoint recipes, 300M models, defaults, and the rollout revert #640
• eval: layout-aware frame-invariant metrics, and a train-split viz head #659
• data: stop letting one bad episode or a rotation channel take a run down #658
• norm: guard degenerate channels, horizon pooling, and no rot6d normalization #657
• keypoints: the 144-D wrist-first hand-keypoint action becomes the human default #656
• mecka: fix the double-mirrored LEFT wrist frame, and correct existing zarrs #655
• 6d: PI goes 6D-only; Eva and Human gain the continuous-rotation modes #654
• 6d: continuous rotation representation, bimanual layouts, shared norm helpers #653
• train: give each Slurm slurm job gets its own cache dir to avoid conflict #638 : 1 other dependent PR (#639 )
• write git version into the run dir for reproducibility #637
• train: don't load PI base weights on resume and verify that checkpoint path exists #636
• configs: declare proprio and actions dims once per HPT model config #634 : 1 other dependent PR (#635 )
• cache norm stats for data configs to minimize recomputation #633
• import cloud packages lazily to support local training #631 : 1 other dependent PR (#632 )
• cache dataset resolver for faster training initialization #630
• validate episode hashes in data config before training #628
• modify pi camera map to match data #629
• add tests to ensure training works between every model and data vendor #626
• run CI tests on every PR #625
• add env variables for user-specific paths #624
• build: set pyproject and uv as the dependency managers #623
• fix(configs): fix paths for data hydra yaml configs #622
• refactor(algo): delete ACT #627 : 1 other dependent PR (#621 )
• tests: create dedicated test folder with unit and integration tests #620
• add pyright type checker #619
• main

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): 2 nit.

[RyanPCo]

nit (dead-code): Once self.device = None is set, the following model.device = self.device assigns None to an HPTModel whose init already sets device None. The new setter also handles propagation once nets['policy'] exists, so that line can go. Any device kwarg is now silently ignored; no shipped config passes one, but a pop-with-warning would make the change explicit.

[RyanPCo]

nit (comments): The module docstring includes roadmap and status prose ('nothing else uses it yet', 'Seven of these eight methods are algo-agnostic ... small follow-up') that goes stale on the next edit or when PI adopts the mixin. Keep the seeding formula and the list of attributes the host algo must set, and track the PI follow-up in an issue or TODO.

[RyanPCo]

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

Fixed

• model.device = self.device is a no-op after the change: Removed the line. A device kwarg is still ignored, but HPT now prints a yellow warning saying so.
• Module docstring carries roadmap/status prose: The docstring keeps the seeding formula and the list of attributes the host must set. The PI follow-up is now a TODO above the class, and the method count is gone.