feat(train): log per-rank HBM around LIBERO eval boundary

[shuheng-liu]

What this does

Adds a per-rank HBM probe around the LIBERO eval block in src/opentau/scripts/train.py:

• Just before the eval rollouts: snapshot torch.cuda.memory_allocated() / memory_reserved() and call reset_peak_memory_stats() so the peak counters cover only the eval window.
• Just after the post-eval accelerator.wait_for_everyone() barrier: gather {pre, peak, post} allocated/reserved per rank via gather_object, and emit one structured logging.info line per rank on rank 0 with the post − pre retained delta.

No-op when eval_envs is None (the existing eval gate) or when CUDA is unavailable.

Motivation

We have empirical evidence from an 8-GPU ZeRO-2 training run that crashes correlate with the eval_freq step boundary:

• Two CUDA OOM failures in accelerator.backward ~3 training steps after LIBERO eval finishes (78.8 GiB used / 310 MiB free on rank 0, vs. the steady-state training peak of ~50 GiB for the same cam=4 ZeRO-2 sdpa+ckpt cell in feat(pi07,fsdp): enable FSDP-FULL_SHARD for full unfreeze + profile_step audit + ZeRO-2 vs FSDP matrix #273's matrix).
• One NCCL ALLGATHER timeout where only one rank had enqueued the next collective post-eval — the other 7 were stuck in a non-collective code path. Consistent with asymmetric env teardown.

Both patterns are consistent with LIBERO sim envs / vectorized rollouts retaining CUDA buffers across the eval → training boundary, but we can't confirm without a measurement. This probe makes the retained-vs-peak distinction visible directly in the run log every eval_freq steps (~once every 24 min of wall-clock on the current config), so the next failure has actionable data attached.

How it was tested

• pre-commit run --files src/opentau/scripts/train.py — all hooks pass (ruff, ruff-format, pyupgrade, typos, bandit, license-header, secret-scan).
• pytest -m "not gpu" -n auto tests/scripts/ — 155 passed locally.

End-to-end verification needs a GPU node and is not covered by CI; on the next training launch with eval_freq > 0 and eval_envs configured, the log will get one new line per rank per eval boundary in the form:

Eval HBM probe step=13000 rank=0 pre=50.21/50.78 peak=72.40/78.83 post=58.10/60.22 retained_alloc=+7.89 GiB

Overhead

Negligible. torch.cuda.memory_* reads poll allocator bookkeeping (no kernel launch, no CUDA sync). gather_object of a ~50-byte dict across 8 ranks is sub-ms and is placed immediately after the existing accelerator.wait_for_everyone() so it doesn't add a sync point. Fires once per eval_freq training steps.

How to checkout & try? (for the reviewer)

git fetch origin claude/awesome-booth-b2f762
git checkout claude/awesome-booth-b2f762
uv sync --extra dev --extra libero

CPU sanity:

pytest -m "not gpu" -n auto tests/scripts/

Inspect the new lines:

sed -n '700,795p' src/opentau/scripts/train.py

To exercise the probe end-to-end (needs a GPU + LIBERO env):

opentau-train \
    --accelerate-config configs/examples/accelerate_deepspeed_config.yaml \
    --config_path=configs/examples/pi05_training_config.json \
    --eval_freq=10 --steps=20

Checklist

• I have added Google-style docstrings to important functions and ensured function parameters are typed.
• My PR includes policy-related changes.
  • If the above is checked: I have run the GPU pytests (pytest -m "gpu") and regression tests.