Validate HPU lazy mode on hardware (mark_step, probe sync, and torch.save workarounds are inert in eager mode)
- 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
Bewertung
Dieses Issue wurde noch nicht bewertet.