Skip to content

hpt: build the policy on CPU and let Lightning place it - #681

Closed
RyanPCo wants to merge 1 commit into
ryanco/rdtfrom
ryanco/hpt-cpu-build
Closed

RyanPCo wants to merge 1 commit into
ryanco/rdtfrom
ryanco/hpt-cpu-build

Conversation

@RyanPCo

@RyanPCo RyanPCo commented Sep 20, 2026

Copy link
Copy Markdown
Contributor

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, so the existing
device property -- which already propagates to the HPTModel it holds -- is the
only place a device is ever set.

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.
🤖 Generated with Claude Code

https://claude.ai/code/session_01MrpKGcZMHFDQ4ded57hqga

RyanPCo commented Sep 20, 2026 •

Copy link
Copy Markdown
Contributor Author

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.

A cancelled DDP job left GPUs allocated with nvidia-smi blaming PIDs absent
from /proc. The driver is fine: the memory is held by live processes
reparented to init. gpu19 had 126 pt_data_worker survivors across three
cancellations, each holding /dev/nvidia0-7, pinning 7 GPUs and ~1.7 TB of RAM.

Workers fork after CUDA init, so they inherit the rank's device fds, and ranks
open every visible GPU -- one survivor pins the node. A worker also inherits
both ends of its result pipe, so a dead rank never sends EPIPE; it blocks in
pipe_write and never reaches PyTorch's ppid check. Whole ranks leak the same
way: one wedged in torch.cuda.synchronize() inside inductor cudagraph_trees
ignores SIGTERM until Slurm gives up (UnkillableStepTimeout) and drops the job
with its ranks still spinning. ProctrackType=proctrack/linuxproc walks the PPID
tree, so it never sweeps anything reparented to init.

egomimic/utils/gpu_orphans.py:

* orphan_guarded/die_with_parent put PR_SET_PDEATHSIG on every worker, so the
  kernel kills it with its rank whatever it is blocked on. Wired into the train
  and val loaders and the norm-stats pass.
* arm_rank_pdeathsig does the same for the rank process.
* reap_orphans runs from trainHydra.main before CUDA init and clears what an
  earlier job left, so a leak cannot outlive the job that made it. It only
  signals processes that are ours, reparented to init and holding an nvidia fd.
  Slurm-gated; EGOMIMIC_REAP_GPU_ORPHANS=0 disables it.

reap_gpu_orphans.sh is the manual version for a node we are not submitting to.

The cluster-level fix is ProctrackType=proctrack/cgroup, which needs admin.

Co-Authored-By: Claude Opus 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01DF1j2AzwzFztRsqvbTC35b
@RyanPCo

RyanPCo commented Sep 20, 2026

Copy link
Copy Markdown
Contributor Author

Superseded. The hpt: build the policy on CPU and let Lightning place it commit has been folded into #664, which now carries both the CPU build and the device propagation as one change — #664 alone was only half the fix, since it kept the torch.device('cuda') default that causes the 8-ranks-on-cuda:0 OOM.

What was left on this branch is the slurm GPU-orphan-reap commit, which #682 already carries. Closing in favour of #682.

@RyanPCo RyanPCo closed this Sep 20, 2026
This was referenced Sep 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant