aidotse / aidotse/LeakPro

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

Offen
#449 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
Vorherrschende Sprache
Python
Sterne
23
Forks
28
Ø Merge
4 T. 12 Std.
Gemergte PRs (30 T.)
5

Beschreibung

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.

Beitragsleitfaden

Beitragsleitfaden öffnen

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.