aidotse / aidotse/LeakPro

Validate HPU lazy mode on hardware (mark_step, probe sync, and torch.save workarounds are inert in eager mode)

Aperta
#449 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
23
Fork
28
Merge medio
4g 12h
PR unite (30g)
5

Descrizione

Follow-up to #425.

#425 wires HPU support through every attack family and was smoke-tested on physical Gaudi2 (HL-225) hardware. That run was in **eager mode**, which is now correctly the default (`_HPU_LAZY_MODE = False`, matching Habana's own `is_lazy()`).

The gap: **lazy mode is the only mode in which this PR's HPU-specific code actually does anything**, and it has not been run on hardware.

Specifically, these are all inert in eager mode:

- **`mark_step()`** — `leakpro/utils/device.py` returns early when not lazy, so none of the ~25 call sites ever reach `_htcore.mark_step()`. That includes the four in the GIA reconstruction loop (`abstract_gia.py:120,181,195`, `step_strategies.py:118`) and the ones in every MIA training loop.
- **The probe's `.cpu()` sync** in `_probe_hpu_acquisition` — only load-bearing in lazy mode, where `torch.zeros(1, device="hpu")` alone would be queued rather than executed.
- **The five "move tensors to CPU before `torch.save`" workarounds** (`diff_mi/train_util.py`, `gan_handler.py`, `generator_handler.py`, `cifar_model_handler.py`, `celebA_plgmi_handler.py`) — added for a Habana storage-copy bug that surfaces via the lazy-mode graph.

So the eager-mode smoke test demonstrates that `mark_step()` doesn't *crash*; it can't demonstrate that it works. The README's current phrasing ("gradients, the meta-optimizer step, and `mark_step()` all executed correctly") is stronger than the evidence supports and should be narrowed to eager mode.

### What's needed

1. Run at least one MIA example and one GIA example on Gaudi with `PT_HPU_LAZY_MODE=1` and Habana's torch fork.
2. Confirm `mark_step()` placement is correct in lazy mode (after `optimizer.step()`, once per iteration) and that the reconstruction loop in `abstract_gia.py` doesn't accumulate an unbounded graph across its `deepcopy` of optimizer state.
3. Confirm the `torch.save` workarounds are still necessary, and whether they can be consolidated into one helper in `leakpro/utils/device.py`.

### Not blocking #425

Eager mode is the shipped default and is validated. This tracks the untested path.

Guida per i contributori

Apri la guida per i contributori

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.